From 685cb01d6d4aa354ba818b013f9ab002bbc6e8c1 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 21 Aug 2026 17:01:04 +0800 Subject: [PATCH] feat: add grouped function column bindings (#3994) Function applications from the canonical remote contract cannot currently declare scalar or grouped computed-column outputs atomically. This adds the remote-only declaration contract for scalar, struct-as-one-column, and expanded named-struct outputs. It validates result mappings, fixes exact input/output Arrow schemas in the request, persists grouped sibling metadata, and keeps local Function execution unsupported. Unknown newer application or binding metadata remains readable, while schema-changing mutations fail closed instead of rewriting it. Stable Lance field IDs are deliberately not a declaration prerequisite in this slice. Inputs bind by parameter name and field path; Sophon remains responsible for exact-version validation, atomic all-NULL sibling creation, binding identity and revision allocation, and persisted output identities. --- python/python/lancedb/_lancedb.pyi | 3 + python/python/lancedb/functions.py | 56 +- python/python/lancedb/remote/table.py | 5 +- python/python/lancedb/table.py | 63 +- .../tests/test_first_class_function_slice1.py | 86 ++ python/src/table.rs | 18 + rust/lancedb/src/function.rs | 69 +- rust/lancedb/src/remote/table.rs | 244 +++- rust/lancedb/src/table.rs | 11 + rust/lancedb/src/table/add_columns.rs | 82 +- rust/lancedb/src/table/computed_columns.rs | 1084 ++++++++++++++++- rust/lancedb/src/table/merge.rs | 4 + rust/lancedb/src/table/refresh.rs | 3 + rust/lancedb/src/table/schema_evolution.rs | 17 + rust/lancedb/src/table/update.rs | 4 + .../tests/first_class_function_slice1.rs | 2 + .../v1/remote_function_binding.canonical.json | 2 +- .../v1/remote_function_binding.json | 12 + .../remote_grouped_declaration_request.json | 45 + .../v1/remote_scalar_declaration_request.json | 41 + 20 files changed, 1799 insertions(+), 52 deletions(-) create mode 100644 rust/lancedb/tests/fixtures/first_class_functions/v1/remote_grouped_declaration_request.json create mode 100644 rust/lancedb/tests/fixtures/first_class_functions/v1/remote_scalar_declaration_request.json diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 22878fd85..ea5d3e972 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -341,6 +341,9 @@ class Table: async def add_computed_columns( self, columns: list[tuple[str, str]] ) -> AddColumnsResult: ... + async def add_function_columns( + self, application_json: str, output_name: Optional[str] + ) -> AddColumnsResult: ... async def refresh_column(self, column: str) -> RefreshColumnResult: ... async def refresh_column_async(self, column: str) -> Job: ... async def add_columns_with_schema(self, schema: pa.Schema) -> AddColumnsResult: ... diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 9c4b063cd..df781f665 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -123,6 +123,23 @@ class _RemoteValue(BaseModel): ) +class _OpenRemoteValue(_RemoteValue): + """Forward-readable value whose extras stay out of canonical encoding.""" + + if _PYDANTIC_V2: + model_config = {"extra": "allow", "frozen": True} + else: + + class Config: + allow_mutation = False + extra = "allow" + + def _unknown_field_names(self) -> set[str]: + if _PYDANTIC_V2: + return set((self.__pydantic_extra__ or {}).keys()) + return set(self.__dict__) - set(self.__fields__) + + class FunctionArtifact(_RemoteValue): """Content-addressed Python artifact identity.""" @@ -137,13 +154,13 @@ class FunctionParameter(_RemoteValue): nullable: bool -class FunctionResultField(_RemoteValue): +class FunctionResultField(_OpenRemoteValue): name: str arrow_type: str nullable: bool -class FunctionOutput(_RemoteValue): +class FunctionOutput(_OpenRemoteValue): """Scalar or ordered named-struct output; unknown kinds remain decodable.""" kind: str @@ -211,12 +228,12 @@ class FunctionVersion(_RemoteValue): created_at: str -class FunctionVersionRef(_RemoteValue): +class FunctionVersionRef(_OpenRemoteValue): name: str version: str -class ApplicationInput(_RemoteValue): +class ApplicationInput(_OpenRemoteValue): """One parameter value. Slice 1 freezes integers, strings, booleans, nulls, arrays, and objects. @@ -234,7 +251,7 @@ class ApplicationInput(_RemoteValue): return _validate_literal(value) -class FunctionApplication(_RemoteValue): +class FunctionApplication(_OpenRemoteValue): """Immutable pre-declaration application of an exact Function version.""" function: FunctionVersionRef @@ -243,6 +260,33 @@ class FunctionApplication(_RemoteValue): group_id: str columns: Mapping[str, str] = Field(default_factory=dict) + def _known_dict(self) -> dict[str, Any]: + value = super()._known_dict() + for name in self._unknown_field_names(): + value.pop(name, None) + return value + + def _ensure_declarable(self) -> None: + unknown = {f"application.{name}" for name in self._unknown_field_names()} + unknown.update( + f"function.{name}" for name in self.function._unknown_field_names() + ) + for index, input_value in enumerate(self.inputs): + unknown.update( + f"inputs[{index}].{name}" for name in input_value._unknown_field_names() + ) + unknown.update(f"output.{name}" for name in self.output._unknown_field_names()) + for index, field in enumerate(self.output.fields): + unknown.update( + f"output.fields[{index}].{name}" + for name in field._unknown_field_names() + ) + if unknown: + raise ValueError( + "Function application contains fields from a newer contract: " + f"{sorted(unknown)!r}" + ) + def rename(self, *, columns: Mapping[str, str]) -> FunctionApplication: """Return a copy with result-field to table-column aliases.""" if self.output.kind != "named_struct": @@ -293,6 +337,8 @@ class FunctionBinding(_RemoteValue): group_id: str inputs: tuple[InputBinding, ...] outputs: tuple[OutputMapping, ...] + input_schema: Optional[Mapping[str, Any]] = None + output_schema: Optional[Mapping[str, Any]] = None class RefreshColumnResult(_RemoteValue): diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 25363cf8f..b97f8f194 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -49,6 +49,7 @@ from lancedb.index import ( LabelList, ) from lancedb.job import Job +from lancedb.functions import FunctionApplication from lancedb.remote.db import LOOP from lancedb.table import IndexConfigType, KNOWN_METRICS import pyarrow as pa @@ -960,7 +961,9 @@ class RemoteTable(Table): def add_columns( self, - transforms: Dict[str, str] | None = None, + transforms: Dict[str, str | FunctionApplication] + | FunctionApplication + | None = None, *, computed: Dict[str, str] | None = None, ) -> AddColumnsResult: diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 79e67fdba..913ab5289 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -72,6 +72,7 @@ from .index import ( FTS, ) from .expr import Expr +from .functions import FunctionApplication from .merge import LanceMergeInsertBuilder from .pydantic import LanceModel, model_to_dict from .query import ( @@ -1942,7 +1943,8 @@ class Table(ABC): @abstractmethod def add_columns( self, - transforms: Dict[str, str] + transforms: Dict[str, str | FunctionApplication] + | FunctionApplication | pa.Field | List[pa.Field] | pa.Schema @@ -1955,13 +1957,21 @@ class Table(ABC): Parameters ---------- - transforms: Dict[str, str], pa.Field, List[pa.Field], pa.Schema + transforms: Dict[str, str | FunctionApplication], FunctionApplication, + pa.Field, List[pa.Field], pa.Schema A map of column name to a SQL expression to use to calculate the value of the new column. These expressions will be evaluated for each row in the table, and can reference existing columns. Alternatively, a pyarrow Field or Schema can be provided to add new columns with the specified data types. The new columns will be initialized with null values. + + A mapping with one ``FunctionApplication`` value keeps its scalar + or named-struct result in the named table column. A bare + named-struct application expands its ordered result fields as one + atomic sibling group; aliases come from ``rename(columns=...)``. + Function columns are supported only on LanceDB Cloud and + Enterprise. computed: Dict[str, str], optional A map of column name to a SQL expression defining the column. The column's type and inputs are derived from the expression, so no @@ -4056,9 +4066,10 @@ class LanceTable(Table): def add_columns( self, - transforms: Dict[str, str] - | pa.field - | List[pa.field] + transforms: Dict[str, str | FunctionApplication] + | FunctionApplication + | pa.Field + | List[pa.Field] | pa.Schema | None = None, *, @@ -5992,9 +6003,10 @@ class AsyncTable: async def add_columns( self, - transforms: dict[str, str] - | pa.field - | List[pa.field] + transforms: dict[str, str | FunctionApplication] + | FunctionApplication + | pa.Field + | List[pa.Field] | pa.Schema | None = None, *, @@ -6005,12 +6017,19 @@ class AsyncTable: Parameters ---------- - transforms: Dict[str, str] + transforms: Dict[str, str | FunctionApplication] or FunctionApplication A map of column name to a SQL expression to use to calculate the value of the new column. These expressions will be evaluated for each row in the table, and can reference existing columns. Alternatively, you can pass a pyarrow field or schema to add new columns with NULLs. + + A mapping with one ``FunctionApplication`` value keeps its scalar + or named-struct result in the named table column. A bare + named-struct application expands its ordered result fields as one + atomic sibling group; aliases come from ``rename(columns=...)``. + Function columns are supported only on LanceDB Cloud and + Enterprise. computed: Dict[str, str], optional A map of column name to a SQL expression defining the column. The column's type and inputs are derived from the expression. @@ -6034,6 +6053,32 @@ class AsyncTable: version: the new version number of the table after adding columns. """ + function_application = None + function_output_name = None + if isinstance(transforms, FunctionApplication): + function_application = transforms + elif isinstance(transforms, dict) and any( + isinstance(value, FunctionApplication) for value in transforms.values() + ): + if len(transforms) != 1 or not all( + isinstance(value, FunctionApplication) for value in transforms.values() + ): + raise ValueError( + "one add_columns call declares exactly one Function sibling group" + ) + function_output_name, function_application = next(iter(transforms.items())) + + if function_application is not None: + if computed: + raise ValueError( + "add_columns cannot mix a Function application with SQL " + "computed columns" + ) + function_application._ensure_declarable() + return await self._inner.add_function_columns( + function_application.to_canonical_json(), function_output_name + ) + if isinstance(transforms, pa.Field): transforms = [transforms] if isinstance(transforms, list) and all( diff --git a/python/python/tests/test_first_class_function_slice1.py b/python/python/tests/test_first_class_function_slice1.py index 9f934507f..fead28bc8 100644 --- a/python/python/tests/test_first_class_function_slice1.py +++ b/python/python/tests/test_first_class_function_slice1.py @@ -14,6 +14,7 @@ from lancedb.functions import ( PythonRuntimeSpec, RefreshColumnResult, ) +from lancedb.table import AsyncTable FIXTURES = ( @@ -166,6 +167,8 @@ def test_binding_and_refresh_result_keep_stable_remote_fields(): assert binding.revision == 3 assert binding.function.version == "fv_01K3TEXT" assert [output.output_ordinal for output in binding.outputs] == [0, 1] + assert binding.input_schema is not None + assert binding.output_schema is not None result = RefreshColumnResult.from_json( json.dumps(job_result("remote_refresh_job.json")) @@ -223,3 +226,86 @@ def test_canonical_client_values_contain_secret_names_only(): canonical = json.loads(version.to_canonical_json()) assert canonical["required_secrets"] == ["HF_TOKEN"] assert_no_secret_values(canonical) + + +class _FunctionDeclarationInner: + def __init__(self): + self.calls = [] + + async def add_function_columns(self, application_json, output_name): + self.calls.append((json.loads(application_json), output_name)) + return "declared" + + +def known_application() -> FunctionApplication: + value = json.loads(fixture("remote_function_application.json")) + value.pop("future_application") + return FunctionApplication(**value) + + +@pytest.mark.asyncio +async def test_add_columns_routes_struct_as_one_and_grouped_expansion_atomically(): + inner = _FunctionDeclarationInner() + table = AsyncTable(inner) + application = known_application() + + result = await table.add_columns( + {"features": application._copy(update={"columns": {}})} + ) + assert result == "declared" + assert inner.calls[-1][1] == "features" + + bare = application._copy(update={"columns": {}}).rename( + columns={"normalized_text": "search_text"} + ) + result = await table.add_columns(bare) + assert result == "declared" + assert inner.calls[-1][1] is None + assert inner.calls[-1][0]["columns"] == {"normalized_text": "search_text"} + + +@pytest.mark.asyncio +async def test_add_columns_rejects_mixed_groups_and_unknown_newer_application(): + inner = _FunctionDeclarationInner() + table = AsyncTable(inner) + application = known_application() + + with pytest.raises(ValueError, match="exactly one Function sibling group"): + await table.add_columns({"a": application, "b": application}) + + future = json.loads(fixture("remote_function_application.json")) + application = FunctionApplication(**future) + with pytest.raises(ValueError, match="newer contract"): + await table.add_columns(application) + + future.pop("future_application") + future["output"]["assignment"] = "cell_flag" + application = FunctionApplication(**future) + assert "assignment" not in json.loads(application.to_canonical_json())["output"] + with pytest.raises(ValueError, match="output.assignment"): + await table.add_columns(application) + assert inner.calls == [] + + +def test_rename_requires_named_struct_and_keeps_partial_mapping_immutable(): + scalar = FunctionApplication.from_json( + json.dumps( + { + "function": {"name": "embed", "version": "fv_exact"}, + "inputs": [], + "output": { + "kind": "scalar", + "arrow_type": "list", + "nullable": False, + }, + "group_id": "fg_scalar", + } + ) + ) + with pytest.raises(ValueError, match="named-struct"): + scalar.rename(columns={"value": "embedding"}) + + application = known_application()._copy(update={"columns": {}}) + renamed = application.rename(columns={"normalized_text": "search_text"}) + assert dict(application.columns) == {} + assert dict(renamed.columns) == {"normalized_text": "search_text"} diff --git a/python/src/table.rs b/python/src/table.rs index 0e3eb4cf8..cb4752cce 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1551,6 +1551,24 @@ impl Table { }) } + pub fn add_function_columns( + self_: PyRef<'_, Self>, + application_json: String, + output_name: Option, + ) -> PyResult> { + let application = + lancedb::function::FunctionApplication::from_json(&application_json).infer_error()?; + let inner = self_.inner_ref()?.clone(); + future_into_py(self_.py(), async move { + let builder = match output_name { + Some(name) => inner.add_columns().function_as(name, application), + None => inner.add_columns().function(application), + }; + let result = builder.execute().await.infer_error()?; + Ok(AddColumnsResult::from(result)) + }) + } + pub fn refresh_column(self_: PyRef<'_, Self>, column: String) -> PyResult> { let inner = self_.inner_ref()?.clone(); future_into_py(self_.py(), async move { diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index 087a00b90..fe91f1680 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -74,6 +74,46 @@ fn validate_literal(value: &Value) -> Result<()> { } } +fn has_unknown_keys(value: &Value, allowed: &[&str]) -> bool { + value + .as_object() + .is_some_and(|object| object.keys().any(|key| !allowed.contains(&key.as_str()))) +} + +fn application_has_unknown_nested_fields(value: &Value) -> bool { + let Some(application) = value.as_object() else { + return false; + }; + if application + .get("function") + .is_some_and(|value| has_unknown_keys(value, &["name", "version"])) + { + return true; + } + if application + .get("inputs") + .and_then(Value::as_array) + .is_some_and(|inputs| { + inputs + .iter() + .any(|input| has_unknown_keys(input, &["parameter", "kind", "value"])) + }) + { + return true; + } + application.get("output").is_some_and(|output| { + has_unknown_keys(output, &["kind", "arrow_type", "nullable", "fields"]) + || output + .get("fields") + .and_then(Value::as_array) + .is_some_and(|fields| { + fields + .iter() + .any(|field| has_unknown_keys(field, &["name", "arrow_type", "nullable"])) + }) + }) +} + macro_rules! impl_json { ($type:ty) => { impl $type { @@ -358,6 +398,10 @@ pub struct FunctionApplication { group_id: String, #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] columns: BTreeMap, + #[serde(default, flatten, skip_serializing)] + unknown_fields: BTreeMap, + #[serde(default, skip)] + unknown_nested_fields: bool, } impl FunctionApplication { @@ -380,10 +424,18 @@ impl FunctionApplication { pub fn columns(&self) -> &BTreeMap { &self.columns } + /// Whether a newer writer attached application fields this client cannot + /// validate. Such applications remain readable but must not be declared. + pub fn has_unknown_fields(&self) -> bool { + !self.unknown_fields.is_empty() || self.unknown_nested_fields + } /// Decode a remote application after validating the Slice 1 literal domain. pub fn from_json(json: &str) -> Result { - let application: Self = from_json(json)?; + let value: Value = from_json(json)?; + let has_unknown_nested_fields = application_has_unknown_nested_fields(&value); + let mut application: Self = serde_json::from_value(value).map_err(invalid_json)?; + application.unknown_nested_fields = has_unknown_nested_fields; application .inputs .iter() @@ -433,6 +485,13 @@ pub struct FunctionBinding { group_id: String, inputs: Vec, outputs: Vec, + /// Exact Arrow schema presented to the Function, encoded with the Lance + /// Namespace Arrow JSON representation. + #[serde(default, skip_serializing_if = "Option::is_none")] + input_schema: Option, + /// Exact physical Arrow schema of the grouped table outputs. + #[serde(default, skip_serializing_if = "Option::is_none")] + output_schema: Option, } impl FunctionBinding { @@ -459,6 +518,14 @@ impl FunctionBinding { pub fn outputs(&self) -> &[OutputMapping] { &self.outputs } + + pub fn input_schema(&self) -> Option<&Value> { + self.input_schema.as_ref() + } + + pub fn output_schema(&self) -> Option<&Value> { + self.output_schema.as_ref() + } } impl_json!(FunctionBinding); diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 328f2a708..1e5691554 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -2147,6 +2147,7 @@ impl BaseTable for RemoteTable { self.check_mutable().await?; let table_schema = self.schema().await?; + crate::table::computed_columns::ensure_supported_function_metadata(table_schema.as_ref())?; let table_def = TableDefinition::try_from_rich_schema(table_schema.clone())?; let num_partitions = if self.server_version.support_multipart_write() { @@ -2698,6 +2699,10 @@ impl BaseTable for RemoteTable { _read_columns: Option>, ) -> Result { self.check_mutable().await?; + crate::table::computed_columns::ensure_no_function_bindings_for_mutation( + self.schema().await?.as_ref(), + "schema evolution", + )?; match transforms { NewColumnTransform::SqlExpressions(expressions) => { let body = expressions @@ -2746,6 +2751,10 @@ impl BaseTable for RemoteTable { async fn add_computed_columns(&self, columns: &[(String, String)]) -> Result { self.check_mutable().await?; + crate::table::computed_columns::ensure_no_function_bindings_for_mutation( + self.schema().await?.as_ref(), + "schema evolution", + )?; // The server plans the declaration: expression validation, type // inference and the persisted binding all happen there. let entries = columns @@ -2785,6 +2794,63 @@ impl BaseTable for RemoteTable { Ok(result) } + async fn add_function_columns( + &self, + application: &crate::function::FunctionApplication, + output_name: Option<&str>, + ) -> Result { + self.check_mutable().await?; + let schema = self.schema().await?; + let plan = crate::table::computed_columns::plan_function_application( + schema.as_ref(), + application, + output_name, + )?; + let new_columns = plan + .outputs + .iter() + .map(|output| { + serde_json::json!({ + "name": output.output_name, + "all_null": true, + }) + }) + .collect::>(); + let mut body = serde_json::json!({ + "new_columns": new_columns, + "function": { + "application": plan.application, + "binding_metadata_version": plan.binding_metadata_version, + "input_bindings": plan.input_bindings, + "input_schema": plan.input_schema, + "output_schema": plan.output_schema, + "outputs": plan.outputs, + }, + }); + self.apply_branch_body(&mut body); + let request = self + .client + .post(&format!("/v1/table/{}/add_columns/", self.identifier)) + .json(&body); + let (request_id, response) = self.send(request, true).await?; + let response = self.check_table_response(&request_id, response).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + + if body.trim().is_empty() { + return Ok(AddColumnsResult { version: 0 }); + } + + let result: AddColumnsResult = serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse add Function columns response: {e}").into(), + request_id, + status_code: None, + })?; + + self.invalidate_schema_cache(); + self.track_write_version(result.version); + Ok(result) + } + async fn refresh_column(&self, _column: &str) -> Result { // The server runs a refresh as a job and does not report a fill // count, so the blocking form has no honest result to return. @@ -3810,11 +3876,14 @@ mod tests { assert_eq!(rename, "y"); if old_server { - http::Response::builder().status(200).body("{}").unwrap() + http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap() } else { http::Response::builder() .status(200) - .body(r#"{"version": 43}"#) + .body(r#"{"version": 43}"#.to_string()) .unwrap() } } else { @@ -3945,11 +4014,14 @@ mod tests { assert_eq!(predicate, "id in (1, 2, 3)"); if old_server { - http::Response::builder().status(200).body("{}").unwrap() + http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap() } else { http::Response::builder() .status(200) - .body(r#"{"version": 43}"#) + .body(r#"{"version": 43}"#.to_string()) .unwrap() } } else { @@ -6516,7 +6588,9 @@ mod tests { #[tokio::test] async fn test_add_columns(#[case] old_server: bool) { let table = Table::new_with_handler("my_table", move |request| { - if request.url().path() == "/v1/table/my_table/add_columns/" { + if request.url().path() == "/v1/table/my_table/describe/" { + simple_describe_response() + } else if request.url().path() == "/v1/table/my_table/add_columns/" { assert_eq!(request.method(), "POST"); assert_eq!( request.headers().get("Content-Type").unwrap(), @@ -6540,11 +6614,14 @@ mod tests { assert_eq!(expression, "cast(NULL as int32)"); if old_server { - http::Response::builder().status(200).body("{}").unwrap() + http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap() } else { http::Response::builder() .status(200) - .body(r#"{"version": 43}"#) + .body(r#"{"version": 43}"#.to_string()) .unwrap() } } else { @@ -6569,19 +6646,22 @@ mod tests { /// plan; the client never types the expression itself. #[tokio::test] async fn test_add_computed_columns_sends_the_expression() { - let table = Table::new_with_handler("my_table", |request| { - assert_eq!(request.method(), "POST"); - assert_eq!(request.url().path(), "/v1/table/my_table/add_columns/"); - let body = request.body().unwrap().as_bytes().unwrap(); - let value: serde_json::Value = serde_json::from_slice(body).unwrap(); - assert_eq!( - value["new_columns"], - serde_json::json!([{"name": "doubled", "computed": "x * 2"}]) - ); - http::Response::builder() - .status(200) - .body(r#"{"version": 7}"#) - .unwrap() + let table = Table::new_with_handler("my_table", |request| match request.url().path() { + "/v1/table/my_table/describe/" => simple_describe_response(), + "/v1/table/my_table/add_columns/" => { + assert_eq!(request.method(), "POST"); + let body = request.body().unwrap().as_bytes().unwrap(); + let value: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!( + value["new_columns"], + serde_json::json!([{"name": "doubled", "computed": "x * 2"}]) + ); + http::Response::builder() + .status(200) + .body(r#"{"version": 7}"#.to_string()) + .unwrap() + } + path => panic!("Unexpected path: {path}"), }); let result = table @@ -6593,6 +6673,129 @@ mod tests { assert_eq!(result.version, 7); } + #[tokio::test] + async fn test_add_scalar_function_column_sends_atomic_null_declaration() { + let table = Table::new_with_handler("my_table", |request| { + match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body( + r#"{"version":1,"schema":{"fields":[{"name":"description","nullable":true,"type":{"type":"string"}}]}}"#, + ) + .unwrap(), + "/v1/table/my_table/add_columns/" => { + let actual: serde_json::Value = serde_json::from_slice( + request.body().unwrap().as_bytes().unwrap(), + ) + .unwrap(); + let expected: serde_json::Value = serde_json::from_str(include_str!( + "../../tests/fixtures/first_class_functions/v1/remote_scalar_declaration_request.json" + )) + .unwrap(); + assert_eq!(actual, expected); + http::Response::builder() + .status(200) + .body(r#"{"version":8}"#) + .unwrap() + } + path => panic!("Unexpected path: {path}"), + } + }); + let application = crate::function::FunctionApplication::from_json( + r#"{ + "function":{"name":"embed","version":"fv_01K3EXACT"}, + "inputs":[{"parameter":"text","kind":"column","value":{"path":"description"}}], + "output":{"kind":"scalar","arrow_type":"list","nullable":false}, + "group_id":"fg_scalar" + }"#, + ) + .unwrap(); + + let result = table + .add_columns() + .function_as("embedding", application) + .execute() + .await + .unwrap(); + assert_eq!(result.version, 8); + } + + #[tokio::test] + async fn test_add_named_struct_function_expands_one_atomic_sibling_group() { + let table = Table::new_with_handler("my_table", |request| match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body( + r#"{"version":1,"schema":{"fields":[ + {"name":"title","nullable":true,"type":{"type":"string"}}, + {"name":"body","nullable":true,"type":{"type":"string"}} + ]}}"#, + ) + .unwrap(), + "/v1/table/my_table/add_columns/" => { + let actual: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + let expected: serde_json::Value = serde_json::from_str(include_str!( + "../../tests/fixtures/first_class_functions/v1/remote_grouped_declaration_request.json" + )) + .unwrap(); + assert_eq!(actual, expected); + http::Response::builder() + .status(200) + .body(r#"{"version":9}"#) + .unwrap() + } + path => panic!("Unexpected path: {path}"), + }); + let application = crate::function::FunctionApplication::from_json( + r#"{ + "function":{"name":"text_features","version":"fv_01K3TEXT"}, + "inputs":[ + {"parameter":"title","kind":"column","value":{"path":"title"}}, + {"parameter":"body","kind":"column","value":{"path":"body"}} + ], + "output":{"kind":"named_struct","fields":[ + {"name":"normalized_text","arrow_type":"utf8","nullable":false}, + {"name":"token_count","arrow_type":"int64","nullable":false} + ]}, + "group_id":"fg_01K3TEXT", + "columns":{"normalized_text":"search_text"} + }"#, + ) + .unwrap(); + + let result = table + .add_columns() + .function(application) + .execute() + .await + .unwrap(); + assert_eq!(result.version, 9); + } + + #[tokio::test] + async fn test_add_columns_fails_closed_on_newer_function_binding_metadata() { + let table = Table::new_with_handler("my_table", |request| { + match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body( + r#"{"version":1,"schema":{"fields":[{"name":"x","nullable":true,"type":{"type":"int32"}}],"metadata":{"lancedb::function_bindings":"{\"version\":2,\"bindings\":[]}"}}}"#, + ) + .unwrap(), + path => panic!("mutation request must not be sent: {path}"), + } + }); + + let err = table + .add_columns() + .computed("doubled", "x * 2") + .execute() + .await + .unwrap_err(); + assert!(matches!(err, Error::NotSupported { .. })); + } + /// A remote refresh is a server job: the async form returns its handle, /// and the blocking form refuses rather than invent a fill count. #[tokio::test] @@ -10911,6 +11114,7 @@ mod tests { .status(200) .body("{}".to_string()) .unwrap(), + "/v1/table/my_table/describe/" => simple_describe_response(), "/v1/table/my_table/add_columns/" | "/v1/table/my_table/alter_columns/" | "/v1/table/my_table/drop_columns/" => { diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 2e16b0940..9228b4baf 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -760,6 +760,16 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { message: "computed columns are not supported on this table type".into(), }) } + /// Declare one immutable registered-Function output group. + async fn add_function_columns( + &self, + _application: &crate::function::FunctionApplication, + _output_name: Option<&str>, + ) -> Result { + Err(Error::NotSupported { + message: "Function columns are supported only on LanceDB Cloud and Enterprise".into(), + }) + } /// Fill a computed column's unfilled rows. /// /// The default returns `NotSupported`; Lance-backed tables override it. @@ -3158,6 +3168,7 @@ impl BaseTable for NativeTable { let ds = self.dataset.get().await?; let table_schema = Schema::from(&ds.schema().clone()); + computed_columns::ensure_supported_function_metadata(&table_schema)?; computed_columns::ensure_not_written( &table_schema, add.data.schema().fields().iter().map(|f| f.name().as_str()), diff --git a/rust/lancedb/src/table/add_columns.rs b/rust/lancedb/src/table/add_columns.rs index 67764c346..3b91c30e4 100644 --- a/rust/lancedb/src/table/add_columns.rs +++ b/rust/lancedb/src/table/add_columns.rs @@ -9,6 +9,7 @@ use lance::dataset::NewColumnTransform; use super::BaseTable; use super::schema_evolution::AddColumnsResult; +use crate::function::FunctionApplication; use crate::{Error, Result}; /// Adds columns to a table. See [`Table::add_columns`](super::Table::add_columns). @@ -16,6 +17,7 @@ pub struct AddColumnsBuilder { parent: Arc, transform: Option, computed: Vec<(String, String)>, + function: Option<(FunctionApplication, Option)>, read_columns: Option>, } @@ -25,6 +27,7 @@ impl std::fmt::Debug for AddColumnsBuilder { .field("parent", &self.parent) .field("has_transform", &self.transform.is_some()) .field("computed", &self.computed) + .field("has_function", &self.function.is_some()) .field("read_columns", &self.read_columns) .finish() } @@ -36,6 +39,7 @@ impl AddColumnsBuilder { parent, transform: None, computed: Vec::new(), + function: None, read_columns: None, } } @@ -83,6 +87,48 @@ impl AddColumnsBuilder { self } + /// Declare every field of a named-struct Function result as one atomic + /// sibling group. Result-field aliases come from + /// [`FunctionApplication::columns`](crate::function::FunctionApplication::columns). + /// + /// ``` + /// # use lancedb::Table; + /// # use lancedb::function::FunctionApplication; + /// # async fn declare(table: &Table, application: FunctionApplication) -> lancedb::Result<()> { + /// table.add_columns().function(application).execute().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn function(mut self, application: FunctionApplication) -> Self { + self.function = Some((application, None)); + self + } + + /// Declare a scalar or entire named-struct Function result as one table + /// column. The physical column starts all-null and is materialized by the + /// remote Function refresh path. + /// + /// ``` + /// # use lancedb::Table; + /// # use lancedb::function::FunctionApplication; + /// # async fn declare(table: &Table, application: FunctionApplication) -> lancedb::Result<()> { + /// table + /// .add_columns() + /// .function_as("embedding", application) + /// .execute() + /// .await?; + /// # Ok(()) + /// # } + /// ``` + pub fn function_as( + mut self, + name: impl Into, + application: FunctionApplication, + ) -> Self { + self.function = Some((application, Some(name.into()))); + self + } + /// Limit which existing columns a [`NewColumnTransform::BatchUDF`] mapper /// receives. Every other transform, and a computed column, determines what /// it reads, so setting this alongside one is an error rather than a silent @@ -98,21 +144,23 @@ impl AddColumnsBuilder { parent, transform, computed, + function, read_columns, } = self; - match (transform, computed.is_empty()) { - (None, true) => Err(Error::InvalidInput { + let declaration_count = usize::from(!computed.is_empty()) + usize::from(function.is_some()); + if transform.is_some() && declaration_count != 0 || declaration_count > 1 { + return Err(Error::InvalidInput { + message: "add_columns cannot mix transforms, SQL computed columns, and a Function application; they cannot be added atomically in one call" + .into(), + }); + } + + match (transform, computed.is_empty(), function) { + (None, true, None) => Err(Error::InvalidInput { message: "add_columns requires a transform or a computed column".into(), }), - // The two commit through different transforms, so one call covering - // both would be two commits and could half-apply. - (Some(_), false) => Err(Error::InvalidInput { - message: "add_columns cannot mix a transform with computed columns; \ - they cannot be added atomically in one call" - .into(), - }), - (Some(transform), true) => { + (Some(transform), true, None) => { if read_columns.is_some() && !matches!(transform, NewColumnTransform::BatchUDF(_)) { return Err(Error::InvalidInput { message: "read_columns applies only to a BatchUDF transform; \ @@ -122,7 +170,7 @@ impl AddColumnsBuilder { } parent.add_columns(transform, read_columns).await } - (None, false) => { + (None, false, None) => { if read_columns.is_some() { return Err(Error::InvalidInput { message: "read_columns applies only to a BatchUDF transform; \ @@ -132,6 +180,18 @@ impl AddColumnsBuilder { } parent.add_computed_columns(&computed).await } + (None, true, Some((application, output_name))) => { + if read_columns.is_some() { + return Err(Error::InvalidInput { + message: "read_columns does not apply to a Function application; its inputs are already bound" + .into(), + }); + } + parent + .add_function_columns(&application, output_name.as_deref()) + .await + } + _ => unreachable!("mixed add_columns modes were rejected above"), } } } diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index 9a6a2585d..841b13856 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -12,13 +12,14 @@ //! where the column's type and inputs come from. A SQL expression is //! self-describing -- both are derived from the expression, so a caller writes //! neither -- while a kind resolved through a registry cannot be typed without -//! consulting it. Only SQL exists today; the tag is what lets another kind be -//! added without a second reading of the same key. +//! consulting it. Registered Functions use an exact remote version plus a +//! schema-level grouped binding; unknown newer kinds remain readable and fail +//! closed before mutation. //! //! [`computed_columns`] and [`computed_column_from_field`] read declarations //! back off a schema. -use std::collections::HashMap; +use std::collections::{BTreeSet, HashMap}; use std::sync::Arc; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef}; @@ -26,7 +27,11 @@ use datafusion_common::tree_node::TreeNode; use datafusion_physical_plan::PhysicalExpr; use lance::dataset::NewColumnTransform; use lance_datafusion::planner::Planner; +use lance_namespace::models::{JsonArrowDataType, JsonArrowField, JsonArrowSchema}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use crate::function::{FunctionApplication, FunctionBinding}; use crate::{Error, Result}; /// Field metadata key marking a column as computed. The value is `"true"`. @@ -41,9 +46,28 @@ pub const EXPRESSION_META_KEY: &str = "computed_column.expression"; /// Field metadata key holding the column's inputs, as a JSON array of names. pub const INPUTS_META_KEY: &str = "computed_column.inputs"; +/// Field metadata key holding the grouped Function binding identity. +pub const FUNCTION_BINDING_ID_META_KEY: &str = "computed_column.function.binding_id"; + +/// Field metadata key holding this sibling's ordered Function output ordinal. +pub const FUNCTION_OUTPUT_ORDINAL_META_KEY: &str = "computed_column.function.output_ordinal"; + +/// Schema metadata key holding all immutable grouped Function bindings. +pub const FUNCTION_BINDINGS_META_KEY: &str = "lancedb::function_bindings"; + +/// Version of the schema-level grouped Function binding envelope. +pub const FUNCTION_BINDINGS_VERSION: u32 = 1; + /// Value of [`KIND_META_KEY`] for a column defined by a SQL expression. pub const SQL_KIND: &str = "sql"; +/// Value of [`KIND_META_KEY`] for a registered Function binding. +pub const FUNCTION_KIND: &str = "function"; + +/// Synthetic result identity used when the entire Function result maps to one +/// table column (scalar or struct-as-one-column). +pub const WHOLE_RESULT_FIELD: &str = "$value"; + /// The rule that defines a computed column's values. /// /// Non-exhaustive: a kind added later is an additive change, and a caller that @@ -57,6 +81,14 @@ pub enum ComputedColumnKind { /// The expression. expression: String, }, + /// One physical output in an immutable grouped registered-Function + /// binding. The full binding lives in schema metadata. + Function { + /// Shared immutable binding identity. + binding_id: String, + /// Position of this field in the binding's ordered sibling outputs. + output_ordinal: u32, + }, /// A kind this version does not understand, written by a newer one. /// /// Reported rather than hidden so a caller can tell a column it cannot @@ -97,6 +129,251 @@ fn computed_column_metadata(expression: &str, inputs: &[String]) -> HashMap HashMap { + HashMap::from([ + (COMPUTED_COLUMN_META_KEY.to_string(), "true".to_string()), + (KIND_META_KEY.to_string(), FUNCTION_KIND.to_string()), + ( + FUNCTION_BINDING_ID_META_KEY.to_string(), + binding_id.to_string(), + ), + ( + FUNCTION_OUTPUT_ORDINAL_META_KEY.to_string(), + output_ordinal.to_string(), + ), + ( + INPUTS_META_KEY.to_string(), + serde_json::to_string(inputs).unwrap_or_else(|_| "[]".to_string()), + ), + ]) +} + +#[derive(Debug, Serialize, Deserialize)] +struct FunctionBindingEnvelope { + version: u32, + bindings: Vec, +} + +/// Encode immutable grouped bindings for schema-level persistence. +pub fn function_bindings_metadata(bindings: &[FunctionBinding]) -> Result { + let bindings = bindings + .iter() + .map(serde_json::to_value) + .collect::, _>>() + .map_err(|e| Error::InvalidInput { + message: format!("invalid Function binding metadata: {e}"), + })?; + serde_json::to_string(&FunctionBindingEnvelope { + version: FUNCTION_BINDINGS_VERSION, + bindings, + }) + .map_err(|e| Error::InvalidInput { + message: format!("invalid Function binding metadata: {e}"), + }) +} + +/// Decode known grouped Function bindings without rewriting their raw schema +/// metadata. Unknown envelope versions fail closed. +pub fn function_bindings(schema: &ArrowSchema) -> Result> { + let Some(envelope) = function_binding_envelope(schema)? else { + return Ok(Vec::new()); + }; + envelope + .bindings + .into_iter() + .map(|binding| { + serde_json::from_value(binding).map_err(|e| Error::InvalidInput { + message: format!("invalid Function binding metadata: {e}"), + }) + }) + .collect() +} + +fn function_binding_envelope(schema: &ArrowSchema) -> Result> { + let Some(raw) = schema.metadata().get(FUNCTION_BINDINGS_META_KEY) else { + return Ok(None); + }; + let envelope: FunctionBindingEnvelope = + serde_json::from_str(raw).map_err(|e| Error::InvalidInput { + message: format!("invalid Function binding metadata: {e}"), + })?; + if envelope.version != FUNCTION_BINDINGS_VERSION { + return Err(Error::NotSupported { + message: format!( + "Function binding metadata version {} is not supported by this client", + envelope.version + ), + }); + } + Ok(Some(envelope)) +} + +/// Validate metadata before a schema mutation. Read-only access remains +/// possible for older datasets, while incomplete or newer contracts cannot be +/// silently rewritten by this client. +pub(crate) fn ensure_supported_function_metadata(schema: &ArrowSchema) -> Result<()> { + let raw_bindings = function_binding_envelope(schema)? + .map(|envelope| envelope.bindings) + .unwrap_or_default(); + for value in &raw_bindings { + ensure_known_binding_shape(value)?; + } + let bindings = raw_bindings + .into_iter() + .map(|binding| { + serde_json::from_value(binding).map_err(|e| Error::InvalidInput { + message: format!("invalid Function binding metadata: {e}"), + }) + }) + .collect::>>()?; + let mut binding_ids = BTreeSet::new(); + for binding in &bindings { + if !binding_ids.insert(binding.binding_id().to_string()) { + return Err(Error::InvalidInput { + message: format!("duplicate Function binding '{}'", binding.binding_id()), + }); + } + if binding.revision() == 0 || binding.outputs().is_empty() { + return Err(Error::InvalidInput { + message: format!( + "Function binding '{}' has no immutable revision or outputs", + binding.binding_id() + ), + }); + } + if binding.function().name.is_empty() + || binding.function().version.is_empty() + || binding.group_id().is_empty() + { + return Err(Error::InvalidInput { + message: format!( + "Function binding '{}' has no exact version or group identity", + binding.binding_id() + ), + }); + } + if binding.input_schema().is_none() || binding.output_schema().is_none() { + return Err(Error::NotSupported { + message: format!( + "Function binding '{}' does not contain exact Arrow schemas", + binding.binding_id() + ), + }); + } + for (ordinal, output) in binding.outputs().iter().enumerate() { + if output.output_ordinal != ordinal as u32 { + return Err(Error::InvalidInput { + message: format!( + "Function binding '{}' has non-canonical output ordinals", + binding.binding_id() + ), + }); + } + } + ensure_binding_matches_schema(schema, binding)?; + } + + let bindings_by_id = bindings + .iter() + .map(|binding| (binding.binding_id(), binding)) + .collect::>(); + for field in schema.fields() { + if field + .metadata() + .get(COMPUTED_COLUMN_META_KEY) + .map(String::as_str) + != Some("true") + { + continue; + } + match computed_column_from_field(field) { + Some(ComputedColumn { + kind: + ComputedColumnKind::Function { + binding_id, + output_ordinal, + }, + .. + }) => { + let binding = + bindings_by_id + .get(binding_id.as_str()) + .ok_or_else(|| Error::InvalidInput { + message: format!( + "Function output '{}' references missing binding '{}'", + field.name(), + binding_id + ), + })?; + let output = binding + .outputs() + .get(output_ordinal as usize) + .ok_or_else(|| Error::InvalidInput { + message: format!( + "Function output '{}' has invalid ordinal {}", + field.name(), + output_ordinal + ), + })?; + if output.output_name != field.name().as_str() { + return Err(Error::InvalidInput { + message: format!( + "Function output '{}' does not match binding destination '{}'", + field.name(), + output.output_name + ), + }); + } + } + Some(ComputedColumn { + kind: ComputedColumnKind::Sql { .. }, + .. + }) => {} + Some(ComputedColumn { + kind: ComputedColumnKind::Unrecognized { kind }, + .. + }) => { + return Err(Error::NotSupported { + message: format!( + "computed column '{}' uses unsupported kind '{}'", + field.name(), + kind + ), + }); + } + None => { + return Err(Error::InvalidInput { + message: format!( + "computed column '{}' has incomplete declaration metadata", + field.name() + ), + }); + } + } + } + Ok(()) +} + +pub(crate) fn ensure_no_function_bindings_for_mutation( + schema: &ArrowSchema, + operation: &str, +) -> Result<()> { + ensure_supported_function_metadata(schema)?; + if !function_bindings(schema)?.is_empty() { + return Err(Error::NotSupported { + message: format!( + "{operation} is not supported on a table with registered Function bindings" + ), + }); + } + Ok(()) +} + /// Read a field's computed-column declaration, if it carries one. /// /// A field flagged computed but carrying no kind, or a SQL one missing its @@ -114,6 +391,22 @@ pub fn computed_column_from_field(field: &ArrowField) -> Option SQL_KIND => ComputedColumnKind::Sql { expression: metadata.get(EXPRESSION_META_KEY)?.clone(), }, + FUNCTION_KIND => match ( + metadata.get(FUNCTION_BINDING_ID_META_KEY), + metadata + .get(FUNCTION_OUTPUT_ORDINAL_META_KEY) + .and_then(|value| value.parse::().ok()), + ) { + (Some(binding_id), Some(output_ordinal)) if !binding_id.is_empty() => { + ComputedColumnKind::Function { + binding_id: binding_id.clone(), + output_ordinal, + } + } + _ => ComputedColumnKind::Unrecognized { + kind: FUNCTION_KIND.to_string(), + }, + }, other => ComputedColumnKind::Unrecognized { kind: other.to_string(), }, @@ -142,6 +435,552 @@ pub fn computed_columns(schema: &ArrowSchema) -> Vec { .collect() } +#[derive(Debug, Clone, Serialize)] +pub(crate) struct FunctionOutputTarget { + pub result_field: String, + pub output_name: String, + pub output_ordinal: u32, +} + +#[derive(Debug, Clone, Serialize)] +pub(crate) struct FunctionInputTarget { + pub parameter: String, + pub field_path: String, + pub arrow_type: String, + pub nullable: bool, +} + +#[derive(Debug, Clone, Serialize)] +pub(crate) struct FunctionDeclarationPlan { + pub application: FunctionApplication, + pub binding_metadata_version: u32, + pub input_bindings: Vec, + pub input_schema: JsonArrowSchema, + pub output_schema: JsonArrowSchema, + pub outputs: Vec, +} + +fn invalid_function(message: impl Into) -> Error { + Error::InvalidInput { + message: message.into(), + } +} + +fn reject_unknown_object_fields(value: &Value, allowed: &[&str], context: &str) -> Result<()> { + let object = value.as_object().ok_or_else(|| { + invalid_function(format!( + "invalid Function binding metadata: {context} must be an object" + )) + })?; + let unknown = object + .keys() + .filter(|key| !allowed.contains(&key.as_str())) + .cloned() + .collect::>(); + if unknown.is_empty() { + Ok(()) + } else { + Err(Error::NotSupported { + message: format!( + "Function binding metadata contains newer {context} fields: {unknown:?}" + ), + }) + } +} + +fn ensure_known_binding_shape(value: &Value) -> Result<()> { + reject_unknown_object_fields( + value, + &[ + "binding_id", + "revision", + "function", + "group_id", + "inputs", + "outputs", + "input_schema", + "output_schema", + ], + "binding", + )?; + let object = value.as_object().unwrap(); + reject_unknown_object_fields( + object + .get("function") + .ok_or_else(|| invalid_function("Function binding is missing its exact version"))?, + &["name", "version"], + "version reference", + )?; + for input in object + .get("inputs") + .and_then(Value::as_array) + .ok_or_else(|| invalid_function("Function binding inputs must be an array"))? + { + reject_unknown_object_fields( + input, + &[ + "parameter", + "field_id", + "field_path", + "arrow_type", + "nullable", + ], + "input binding", + )?; + } + for output in object + .get("outputs") + .and_then(Value::as_array) + .ok_or_else(|| invalid_function("Function binding outputs must be an array"))? + { + reject_unknown_object_fields( + output, + &[ + "result_field", + "output_name", + "output_field_id", + "output_ordinal", + "arrow_type", + "nullable", + ], + "output mapping", + )?; + } + Ok(()) +} + +fn resolve_field_path<'a>(schema: &'a ArrowSchema, path: &str) -> Result<&'a ArrowField> { + let parts = lance_core::datatypes::parse_field_path(path).map_err(|e| { + invalid_function(format!("invalid Function input field path '{path}': {e}")) + })?; + let Some((root, children)) = parts.split_first() else { + return Err(invalid_function( + "Function input field path cannot be empty", + )); + }; + let mut field = schema + .field_with_name(root) + .map_err(|_| invalid_function(format!("unknown Function input column '{path}'")))?; + for child in children { + let DataType::Struct(fields) = field.data_type() else { + return Err(invalid_function(format!( + "Function input field path '{path}' traverses a non-struct field" + ))); + }; + field = fields + .iter() + .find(|field| field.name() == child) + .map(AsRef::as_ref) + .ok_or_else(|| invalid_function(format!("unknown Function input column '{path}'")))?; + } + Ok(field) +} + +fn canonical_input_arrow_type(field: &JsonArrowField) -> Result { + if field.r#type.fields.is_none() && field.r#type.length.is_none() { + Ok(field.r#type.r#type.clone()) + } else { + serde_json::to_string(field.r#type.as_ref()).map_err(|e| { + invalid_function(format!("could not encode exact Function input type: {e}")) + }) + } +} + +fn parse_output_arrow_type(raw: &str) -> Result { + fn parse(raw: &str) -> Result { + let raw = raw.trim(); + if raw.starts_with('{') { + return serde_json::from_str(raw).map_err(|e| { + invalid_function(format!("invalid Function Arrow type '{raw}': {e}")) + }); + } + if let Some(inner) = raw + .strip_prefix("list<") + .and_then(|value| value.strip_suffix('>')) + { + let mut data_type = JsonArrowDataType::new("list".to_string()); + data_type.fields = Some(vec![JsonArrowField::new( + "item".to_string(), + false, + parse(inner)?, + )]); + return Ok(data_type); + } + if let Some(inner) = raw + .strip_prefix("large_list<") + .and_then(|value| value.strip_suffix('>')) + { + let mut data_type = JsonArrowDataType::new("large_list".to_string()); + data_type.fields = Some(vec![JsonArrowField::new( + "item".to_string(), + false, + parse(inner)?, + )]); + return Ok(data_type); + } + let normalized = match raw { + "boolean" => "bool", + "string" => "utf8", + "large_string" => "large_utf8", + "halffloat" => "float16", + "float" => "float32", + "double" => "float64", + other => other, + }; + Ok(JsonArrowDataType::new(normalized.to_string())) + } + + let data_type = parse(raw)?; + lance_namespace::schema::convert_json_arrow_type(&data_type) + .map_err(|e| invalid_function(format!("unsupported Function Arrow type '{raw}': {e}")))?; + Ok(data_type) +} + +fn ensure_binding_matches_schema(schema: &ArrowSchema, binding: &FunctionBinding) -> Result<()> { + let mut input_fields = Vec::with_capacity(binding.inputs().len()); + for input in binding.inputs() { + let field = resolve_field_path(schema, &input.field_path)?; + if field + .metadata() + .get(COMPUTED_COLUMN_META_KEY) + .map(String::as_str) + == Some("true") + { + return Err(invalid_function(format!( + "Function input '{}' is computed", + input.field_path + ))); + } + if field.is_nullable() != input.nullable { + return Err(invalid_function(format!( + "Function input '{}' no longer matches binding '{}'", + input.field_path, + binding.binding_id() + ))); + } + let parameter_field = ArrowField::new( + input.parameter.clone(), + field.data_type().clone(), + field.is_nullable(), + ) + .with_metadata(field.metadata().clone()); + let json = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![ + parameter_field.clone(), + ])) + .map_err(|e| invalid_function(format!("invalid Function input schema: {e}")))?; + let json_field = json.fields.into_iter().next().unwrap(); + if canonical_input_arrow_type(&json_field)? != input.arrow_type { + return Err(invalid_function(format!( + "Function input '{}' type no longer matches binding '{}'", + input.field_path, + binding.binding_id() + ))); + } + input_fields.push(parameter_field); + } + let input_schema = + lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(input_fields)) + .map_err(|e| invalid_function(format!("invalid Function input schema: {e}")))?; + let input_schema = serde_json::to_value(input_schema).map_err(|e| { + invalid_function(format!("could not encode exact Function input schema: {e}")) + })?; + if binding.input_schema() != Some(&input_schema) { + return Err(invalid_function(format!( + "Function binding '{}' input schema does not match its inputs", + binding.binding_id() + ))); + } + + let mut output_fields = Vec::with_capacity(binding.outputs().len()); + for output in binding.outputs() { + let field = schema.field_with_name(&output.output_name).map_err(|_| { + invalid_function(format!( + "Function binding '{}' output '{}' is missing", + binding.binding_id(), + output.output_name + )) + })?; + if field.name() != &output.output_name || !field.is_nullable() || output.nullable { + return Err(invalid_function(format!( + "Function output '{}' no longer matches binding '{}'", + output.output_name, + binding.binding_id() + ))); + } + let expected_type = parse_output_arrow_type(&output.arrow_type)?; + let expected_type = lance_namespace::schema::convert_json_arrow_type(&expected_type) + .map_err(|e| invalid_function(format!("invalid Function output type: {e}")))?; + if field.data_type() != &expected_type { + return Err(invalid_function(format!( + "Function output '{}' type no longer matches binding '{}'", + output.output_name, + binding.binding_id() + ))); + } + output_fields.push(ArrowField::new( + field.name().clone(), + field.data_type().clone(), + true, + )); + } + let output_schema = + lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(output_fields)) + .map_err(|e| invalid_function(format!("invalid Function output schema: {e}")))?; + let output_schema = serde_json::to_value(output_schema).map_err(|e| { + invalid_function(format!( + "could not encode exact Function output schema: {e}" + )) + })?; + if binding.output_schema() != Some(&output_schema) { + return Err(invalid_function(format!( + "Function binding '{}' output schema does not match physical siblings", + binding.binding_id() + ))); + } + Ok(()) +} + +/// Resolve a Function application against a table schema before any request is +/// serialized. Input paths and the complete sibling output schema are fixed in +/// one plan. +pub(crate) fn plan_function_application( + schema: &ArrowSchema, + application: &FunctionApplication, + output_name: Option<&str>, +) -> Result { + ensure_no_function_bindings_for_mutation(schema, "Function binding declaration")?; + if application.has_unknown_fields() { + return Err(Error::NotSupported { + message: "Function application contains fields from a newer contract".into(), + }); + } + if application.function().name.is_empty() + || application.function().version.is_empty() + || application.group_id().is_empty() + { + return Err(invalid_function( + "Function application requires an exact version and group identity", + )); + } + + let mut parameters = BTreeSet::new(); + let mut input_bindings = Vec::with_capacity(application.inputs().len()); + let mut input_fields = Vec::with_capacity(application.inputs().len()); + for input in application.inputs() { + if !parameters.insert(input.parameter.as_str()) { + return Err(invalid_function(format!( + "duplicate Function parameter '{}'", + input.parameter + ))); + } + if input.kind != "column" { + return Err(Error::NotSupported { + message: format!( + "Function input kind '{}' is not supported for column declaration", + input.kind + ), + }); + } + let source = input.value.as_object().ok_or_else(|| { + invalid_function(format!( + "Function parameter '{}' has an invalid column source", + input.parameter + )) + })?; + if source.len() != 1 { + return Err(Error::NotSupported { + message: format!( + "Function parameter '{}' uses a newer column source contract", + input.parameter + ), + }); + } + let path = source.get("path").and_then(Value::as_str).ok_or_else(|| { + invalid_function(format!( + "Function parameter '{}' requires a column path", + input.parameter + )) + })?; + let field = resolve_field_path(schema, path)?; + if field + .metadata() + .get(COMPUTED_COLUMN_META_KEY) + .map(String::as_str) + == Some("true") + { + return Err(invalid_function(format!( + "Function input '{path}' is computed; computed-on-computed bindings are not supported" + ))); + } + let parameter_field = ArrowField::new( + input.parameter.clone(), + field.data_type().clone(), + field.is_nullable(), + ) + .with_metadata(field.metadata().clone()); + let input_schema = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![ + parameter_field.clone(), + ])) + .map_err(|e| invalid_function(format!("invalid Function input schema: {e}")))?; + let json_field = input_schema.fields.into_iter().next().unwrap(); + input_bindings.push(FunctionInputTarget { + parameter: input.parameter.clone(), + field_path: path.to_string(), + arrow_type: canonical_input_arrow_type(&json_field)?, + nullable: field.is_nullable(), + }); + input_fields.push(parameter_field); + } + let input_schema = + lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(input_fields)) + .map_err(|e| invalid_function(format!("invalid Function input schema: {e}")))?; + + let output = application.output(); + let mut outputs = Vec::new(); + let mut output_fields = Vec::new(); + match output.kind.as_str() { + "scalar" => { + if !application.columns().is_empty() { + return Err(invalid_function( + "scalar Function applications cannot rename result fields", + )); + } + let name = output_name.ok_or_else(|| { + invalid_function( + "a scalar Function application must be mapped to one output column", + ) + })?; + if output.nullable != Some(false) { + return Err(invalid_function( + "Function logical outputs must be non-nullable during NULL assignment", + )); + } + let data_type = + parse_output_arrow_type(output.arrow_type.as_deref().ok_or_else(|| { + invalid_function("scalar Function output is missing its Arrow type") + })?)?; + outputs.push(FunctionOutputTarget { + result_field: WHOLE_RESULT_FIELD.to_string(), + output_name: name.to_string(), + output_ordinal: 0, + }); + output_fields.push(JsonArrowField::new(name.to_string(), true, data_type)); + } + "named_struct" => { + if output.fields.is_empty() { + return Err(invalid_function( + "named-struct Function output requires at least one field", + )); + } + let result_names = output + .fields + .iter() + .map(|field| field.name.as_str()) + .collect::>(); + if result_names.len() != output.fields.len() { + return Err(invalid_function( + "named-struct Function result field names must be unique", + )); + } + if output.fields.iter().any(|field| field.nullable) { + return Err(invalid_function( + "Function logical outputs must be non-nullable during NULL assignment", + )); + } + let unknown = application + .columns() + .keys() + .filter(|name| !result_names.contains(name.as_str())) + .cloned() + .collect::>(); + if !unknown.is_empty() { + return Err(invalid_function(format!( + "unknown Function result fields: {unknown:?}" + ))); + } + + if let Some(name) = output_name { + if !application.columns().is_empty() { + return Err(invalid_function( + "a named-struct mapped to one column cannot also rename expanded fields", + )); + } + let fields = output + .fields + .iter() + .map(|field| { + Ok(JsonArrowField::new( + field.name.clone(), + false, + parse_output_arrow_type(&field.arrow_type)?, + )) + }) + .collect::>>()?; + let mut data_type = JsonArrowDataType::new("struct".to_string()); + data_type.fields = Some(fields); + outputs.push(FunctionOutputTarget { + result_field: WHOLE_RESULT_FIELD.to_string(), + output_name: name.to_string(), + output_ordinal: 0, + }); + output_fields.push(JsonArrowField::new(name.to_string(), true, data_type)); + } else { + let mut destinations = BTreeSet::new(); + for (ordinal, field) in output.fields.iter().enumerate() { + let name = application + .columns() + .get(&field.name) + .unwrap_or(&field.name); + if !destinations.insert(name.as_str()) { + return Err(invalid_function( + "Function output destinations must be unique", + )); + } + outputs.push(FunctionOutputTarget { + result_field: field.name.clone(), + output_name: name.clone(), + output_ordinal: ordinal as u32, + }); + output_fields.push(JsonArrowField::new( + name.clone(), + true, + parse_output_arrow_type(&field.arrow_type)?, + )); + } + } + } + kind => { + return Err(Error::NotSupported { + message: format!( + "Function output kind '{kind}' is not supported for column declaration" + ), + }); + } + } + + for output in &outputs { + if output.output_name.is_empty() { + return Err(invalid_function( + "Function output column name cannot be empty", + )); + } + if schema.field_with_name(&output.output_name).is_ok() { + return Err(Error::ColumnAlreadyExists { + name: output.output_name.clone(), + }); + } + } + + Ok(FunctionDeclarationPlan { + application: application.clone(), + binding_metadata_version: FUNCTION_BINDINGS_VERSION, + input_bindings, + input_schema, + output_schema: JsonArrowSchema::new(output_fields), + outputs, + }) +} + /// Reject a schema change to a column some declaration reads. /// /// A binding is SQL text naming its inputs, so renaming, retyping or dropping @@ -783,7 +1622,7 @@ mod tests { let err = add_computed(&table, &[("embedding".into(), "x * 2".into())]) .await .unwrap_err(); - assert!(matches!(err, Error::ColumnAlreadyExists { name } if name == "embedding")); + assert!(matches!(err, Error::NotSupported { .. })); } /// A kind is what makes a declaration readable at all, so the flag alone @@ -1345,4 +2184,241 @@ mod tests { table.drop_columns(&["doubled"]).await.unwrap(); assert!(declared(&table).await.is_empty()); } + + fn function_input_schema() -> ArrowSchema { + ArrowSchema::new(vec![ + ArrowField::new("title", DataType::Utf8, true), + ArrowField::new("body", DataType::Utf8, true), + ]) + } + + fn named_struct_application(columns: &str) -> FunctionApplication { + FunctionApplication::from_json(&format!( + r#"{{ + "function":{{"name":"text_features","version":"fv_exact"}}, + "inputs":[ + {{"parameter":"title","kind":"column","value":{{"path":"title"}}}}, + {{"parameter":"body","kind":"column","value":{{"path":"body"}}}} + ], + "output":{{"kind":"named_struct","fields":[ + {{"name":"normalized_text","arrow_type":"utf8","nullable":false}}, + {{"name":"token_count","arrow_type":"int64","nullable":false}} + ]}}, + "group_id":"fg_exact", + "columns":{columns} + }}"# + )) + .unwrap() + } + + #[test] + fn test_function_binding_metadata_survives_schema_round_trip() { + let binding = FunctionBinding::from_json(include_str!( + "../../tests/fixtures/first_class_functions/v1/remote_function_binding.json" + )) + .unwrap(); + let raw = function_bindings_metadata(std::slice::from_ref(&binding)).unwrap(); + let mut fields = vec![ + ArrowField::new("title", DataType::Utf8, true), + ArrowField::new("body", DataType::Utf8, true), + ]; + fields.extend( + binding + .outputs() + .iter() + .map(|output| { + let data_type = match output.arrow_type.as_str() { + "utf8" => DataType::Utf8, + "int64" => DataType::Int64, + other => panic!("unexpected fixture output type {other}"), + }; + let metadata = function_computed_column_metadata( + binding.binding_id(), + output.output_ordinal, + &["title".into(), "body".into()], + ); + ArrowField::new(&output.output_name, data_type, true).with_metadata(metadata) + }) + .collect::>(), + ); + let schema = ArrowSchema::new_with_metadata( + fields, + HashMap::from([(FUNCTION_BINDINGS_META_KEY.to_string(), raw)]), + ); + + let reopened = + ArrowSchema::new_with_metadata(schema.fields().to_vec(), schema.metadata().clone()); + let bindings = function_bindings(&reopened).unwrap(); + assert_eq!(bindings, vec![binding.clone()]); + assert!(bindings[0].input_schema().is_some()); + assert!(bindings[0].output_schema().is_some()); + assert!(matches!( + computed_column_from_field(reopened.field(3)).unwrap().kind, + ComputedColumnKind::Function { + ref binding_id, + output_ordinal: 1, + } if binding_id == "fb_01K3TEXT" + )); + let err = plan_function_application(&reopened, &named_struct_application("{}"), None) + .unwrap_err(); + assert!(matches!(err, Error::NotSupported { .. })); + } + + #[test] + fn test_newer_binding_fields_remain_readable_but_fail_closed_on_mutation() { + let raw_binding: Value = serde_json::from_str(include_str!( + "../../tests/fixtures/first_class_functions/v1/remote_function_binding.json" + )) + .unwrap(); + let binding: FunctionBinding = serde_json::from_value(raw_binding.clone()).unwrap(); + assert_eq!(binding.binding_id(), "fb_01K3TEXT"); + + let schema = ArrowSchema::new_with_metadata( + Vec::::new(), + HashMap::from([( + FUNCTION_BINDINGS_META_KEY.to_string(), + serde_json::json!({ + "version": FUNCTION_BINDINGS_VERSION, + "bindings": [raw_binding], + }) + .to_string(), + )]), + ); + let err = ensure_supported_function_metadata(&schema).unwrap_err(); + assert!(matches!(err, Error::NotSupported { .. })); + } + + #[test] + fn test_named_struct_can_be_kept_as_one_nullable_physical_column() { + let application = named_struct_application("{}"); + let plan = + plan_function_application(&function_input_schema(), &application, Some("features")) + .unwrap(); + + assert_eq!(plan.outputs.len(), 1); + assert_eq!(plan.outputs[0].result_field, WHOLE_RESULT_FIELD); + assert_eq!(plan.output_schema.fields.len(), 1); + assert!(plan.output_schema.fields[0].nullable); + assert_eq!(plan.output_schema.fields[0].r#type.r#type, "struct"); + assert_eq!( + plan.output_schema.fields[0] + .r#type + .fields + .as_ref() + .unwrap() + .len(), + 2 + ); + } + + #[test] + fn test_function_mapping_and_sibling_collisions_fail_before_request() { + let unknown = named_struct_application(r#"{"missing":"renamed"}"#); + let err = plan_function_application(&function_input_schema(), &unknown, None).unwrap_err(); + assert!(matches!(&err, Error::InvalidInput { message } if message.contains("unknown"))); + + let duplicate = + named_struct_application(r#"{"normalized_text":"same","token_count":"same"}"#); + let err = + plan_function_application(&function_input_schema(), &duplicate, None).unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { message } if message.contains("destinations")) + ); + + let mut fields = function_input_schema().fields().to_vec(); + fields.push(Arc::new(ArrowField::new( + "token_count", + DataType::Int64, + true, + ))); + let collision_schema = ArrowSchema::new(fields); + let err = + plan_function_application(&collision_schema, &named_struct_application("{}"), None) + .unwrap_err(); + assert!(matches!(err, Error::ColumnAlreadyExists { name } if name == "token_count")); + } + + #[test] + fn test_unknown_and_mixed_version_function_contracts_fail_closed() { + let application = FunctionApplication::from_json( + r#"{ + "function":{"name":"f","version":"fv"}, + "inputs":[{"parameter":"title","kind":"future_source","value":{"path":"title"}}], + "output":{"kind":"scalar","arrow_type":"int64","nullable":false}, + "group_id":"fg" + }"#, + ) + .unwrap(); + let err = plan_function_application(&function_input_schema(), &application, Some("out")) + .unwrap_err(); + assert!(matches!(err, Error::NotSupported { .. })); + + let future_application = FunctionApplication::from_json( + r#"{ + "function":{"name":"f","version":"fv"}, + "inputs":[], + "output":{"kind":"scalar","arrow_type":"int64","nullable":false}, + "group_id":"fg", + "future_declaration":{"mode":"managed"} + }"#, + ) + .unwrap(); + let err = + plan_function_application(&function_input_schema(), &future_application, Some("out")) + .unwrap_err(); + assert!(matches!(err, Error::NotSupported { .. })); + + let nested_future_application = FunctionApplication::from_json( + r#"{ + "function":{"name":"f","version":"fv"}, + "inputs":[], + "output":{"kind":"scalar","arrow_type":"int64","nullable":false,"assignment":"cell_flag"}, + "group_id":"fg" + }"#, + ) + .unwrap(); + let err = plan_function_application( + &function_input_schema(), + &nested_future_application, + Some("out"), + ) + .unwrap_err(); + assert!(matches!(err, Error::NotSupported { .. })); + + let mixed_schema = ArrowSchema::new_with_metadata( + function_input_schema().fields().to_vec(), + HashMap::from([( + FUNCTION_BINDINGS_META_KEY.to_string(), + r#"{"version":2,"bindings":[]}"#.to_string(), + )]), + ); + let err = plan_function_application(&mixed_schema, &named_struct_application("{}"), None) + .unwrap_err(); + assert!(matches!(err, Error::NotSupported { .. })); + } + + #[test] + fn test_function_inputs_use_paths_and_cannot_be_computed() { + let mut schema = function_input_schema(); + let plan = + plan_function_application(&schema, &named_struct_application("{}"), None).unwrap(); + assert_eq!(plan.input_bindings[0].field_path, "title"); + assert_eq!(plan.input_bindings[1].field_path, "body"); + + let title = schema + .field(0) + .as_ref() + .clone() + .with_metadata(HashMap::from([ + (COMPUTED_COLUMN_META_KEY.to_string(), "true".to_string()), + (KIND_META_KEY.to_string(), SQL_KIND.to_string()), + (EXPRESSION_META_KEY.to_string(), "title".to_string()), + ])); + schema = ArrowSchema::new(vec![title, schema.field(1).as_ref().clone()]); + let err = + plan_function_application(&schema, &named_struct_application("{}"), None).unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { message } if message.contains("computed-on-computed")) + ); + } } diff --git a/rust/lancedb/src/table/merge.rs b/rust/lancedb/src/table/merge.rs index 3a5b6882d..ea68e99df 100644 --- a/rust/lancedb/src/table/merge.rs +++ b/rust/lancedb/src/table/merge.rs @@ -233,6 +233,10 @@ pub(crate) async fn execute_merge_insert( params: MergeInsertBuilder, new_data: Box, ) -> Result { + super::computed_columns::ensure_no_function_bindings_for_mutation( + table.schema().await?.as_ref(), + "merge_insert", + )?; match lsm::lsm_dispatch_decision(table, ¶ms).await? { lsm::LsmDispatch::Lsm(plan) => { let future = diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs index b29c97e98..e94f20f2b 100644 --- a/rust/lancedb/src/table/refresh.rs +++ b/rust/lancedb/src/table/refresh.rs @@ -169,6 +169,9 @@ fn declared_expression(dataset: &Dataset, column: &str) -> Result { })?; match declaration.kind { ComputedColumnKind::Sql { expression } => Ok(expression), + ComputedColumnKind::Function { .. } => Err(Error::NotSupported { + message: "registered Function columns are refreshed only by a remote server Job".into(), + }), ComputedColumnKind::Unrecognized { kind } => Err(Error::NotSupported { message: format!( "computed column '{column}' is defined by '{kind}', which this version of \ diff --git a/rust/lancedb/src/table/schema_evolution.rs b/rust/lancedb/src/table/schema_evolution.rs index 7503fd790..4e41f0e85 100644 --- a/rust/lancedb/src/table/schema_evolution.rs +++ b/rust/lancedb/src/table/schema_evolution.rs @@ -101,6 +101,10 @@ pub(crate) async fn execute_add_columns( transforms: NewColumnTransform, read_columns: Option>, ) -> Result { + computed_columns::ensure_no_function_bindings_for_mutation( + table.schema().await?.as_ref(), + "schema evolution", + )?; // Declarations are admitted only through [`execute_declare`]. match &transforms { NewColumnTransform::AllNulls(schema) => { @@ -124,6 +128,10 @@ pub(crate) async fn execute_declare( // checked against latest committed state, not this handle's snapshot. // The catch-up flag outlives unset and marks retained SSTable rows. table.checkout_latest().await?; + computed_columns::ensure_no_function_bindings_for_mutation( + table.schema().await?.as_ref(), + "schema evolution", + )?; let catchup = table.dataset.get().await?.manifest().reader_feature_flags & lance_table::feature_flags::FLAG_MEM_WAL_INDEX_CATCHUP != 0; @@ -163,6 +171,10 @@ pub(crate) async fn execute_alter_columns( // Nullability is not part of what an expression resolves against, so only // a rename or a retype can invalidate a binding. let schema = std::sync::Arc::new(ArrowSchema::from(dataset.schema())); + computed_columns::ensure_no_function_bindings_for_mutation( + schema.as_ref(), + "schema evolution", + )?; let rebinding = alterations .iter() .filter(|alteration| alteration.rename.is_some() || alteration.data_type.is_some()) @@ -190,6 +202,10 @@ pub(crate) async fn execute_drop_columns( ) -> Result { table.dataset.ensure_mutable()?; let mut dataset = (*table.dataset.get().await?).clone(); + computed_columns::ensure_no_function_bindings_for_mutation( + &ArrowSchema::from(dataset.schema()), + "schema evolution", + )?; computed_columns::ensure_not_an_input( &std::sync::Arc::new(ArrowSchema::from(dataset.schema())), columns, @@ -215,6 +231,7 @@ pub(crate) async fn execute_update_field_metadata( // binding out from under a refresh. A replace on a declared column would // silently erase it. let schema = ArrowSchema::from(dataset.schema()); + computed_columns::ensure_no_function_bindings_for_mutation(&schema, "schema evolution")?; let declared: Vec = computed_columns::computed_columns(&schema) .into_iter() .map(|declaration| declaration.name) diff --git a/rust/lancedb/src/table/update.rs b/rust/lancedb/src/table/update.rs index fd9fa6828..98050dfe8 100644 --- a/rust/lancedb/src/table/update.rs +++ b/rust/lancedb/src/table/update.rs @@ -82,6 +82,10 @@ pub(crate) async fn execute_update( // 1. Snapshot the current dataset let dataset = table.dataset.get().await?; + super::computed_columns::ensure_no_function_bindings_for_mutation( + &arrow_schema::Schema::from(dataset.schema()), + "update", + )?; super::computed_columns::ensure_not_written( &arrow_schema::Schema::from(dataset.schema()), update.columns.iter().map(|(name, _)| name.as_str()), diff --git a/rust/lancedb/tests/first_class_function_slice1.rs b/rust/lancedb/tests/first_class_function_slice1.rs index dab05fe48..dc565d309 100644 --- a/rust/lancedb/tests/first_class_function_slice1.rs +++ b/rust/lancedb/tests/first_class_function_slice1.rs @@ -88,6 +88,8 @@ fn application_and_binding_match_shared_remote_goldens() { assert_eq!(binding.function().version, "fv_01K3TEXT"); assert_eq!(binding.outputs()[0].output_ordinal, 0); assert_eq!(binding.outputs()[1].output_ordinal, 1); + assert!(binding.input_schema().is_some()); + assert!(binding.output_schema().is_some()); assert_eq!( binding.to_canonical_json().expect("canonical JSON"), fixture("remote_function_binding.canonical.json").trim() diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_binding.canonical.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_binding.canonical.json index c548c4a58..7bf93b8a8 100644 --- a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_binding.canonical.json +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_binding.canonical.json @@ -1 +1 @@ -{"binding_id":"fb_01K3TEXT","function":{"name":"text_features","version":"fv_01K3TEXT"},"group_id":"fg_01K3TEXT","inputs":[{"arrow_type":"utf8","field_id":11,"field_path":"title","nullable":true,"parameter":"title"},{"arrow_type":"utf8","field_id":12,"field_path":"body","nullable":true,"parameter":"body"}],"outputs":[{"arrow_type":"utf8","nullable":false,"output_field_id":21,"output_name":"search_text","output_ordinal":0,"result_field":"normalized_text"},{"arrow_type":"int64","nullable":false,"output_field_id":22,"output_name":"search_token_count","output_ordinal":1,"result_field":"token_count"}],"revision":3} +{"binding_id":"fb_01K3TEXT","function":{"name":"text_features","version":"fv_01K3TEXT"},"group_id":"fg_01K3TEXT","input_schema":{"fields":[{"name":"title","nullable":true,"type":{"type":"utf8"}},{"name":"body","nullable":true,"type":{"type":"utf8"}}]},"inputs":[{"arrow_type":"utf8","field_id":11,"field_path":"title","nullable":true,"parameter":"title"},{"arrow_type":"utf8","field_id":12,"field_path":"body","nullable":true,"parameter":"body"}],"output_schema":{"fields":[{"name":"search_text","nullable":true,"type":{"type":"utf8"}},{"name":"search_token_count","nullable":true,"type":{"type":"int64"}}]},"outputs":[{"arrow_type":"utf8","nullable":false,"output_field_id":21,"output_name":"search_text","output_ordinal":0,"result_field":"normalized_text"},{"arrow_type":"int64","nullable":false,"output_field_id":22,"output_name":"search_token_count","output_ordinal":1,"result_field":"token_count"}],"revision":3} diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_binding.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_binding.json index 5d8193eea..1a2053e42 100644 --- a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_binding.json +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_binding.json @@ -11,5 +11,17 @@ {"result_field": "normalized_text", "output_name": "search_text", "output_field_id": 21, "output_ordinal": 0, "arrow_type": "utf8", "nullable": false}, {"result_field": "token_count", "output_name": "search_token_count", "output_field_id": 22, "output_ordinal": 1, "arrow_type": "int64", "nullable": false} ], + "input_schema": { + "fields": [ + {"name": "title", "nullable": true, "type": {"type": "utf8"}}, + {"name": "body", "nullable": true, "type": {"type": "utf8"}} + ] + }, + "output_schema": { + "fields": [ + {"name": "search_text", "nullable": true, "type": {"type": "utf8"}}, + {"name": "search_token_count", "nullable": true, "type": {"type": "int64"}} + ] + }, "future_binding": {"metadata_revision": 1} } diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_grouped_declaration_request.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_grouped_declaration_request.json new file mode 100644 index 000000000..d0b42cc99 --- /dev/null +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_grouped_declaration_request.json @@ -0,0 +1,45 @@ +{ + "new_columns": [ + {"name": "search_text", "all_null": true}, + {"name": "token_count", "all_null": true} + ], + "function": { + "application": { + "function": {"name": "text_features", "version": "fv_01K3TEXT"}, + "inputs": [ + {"parameter": "title", "kind": "column", "value": {"path": "title"}}, + {"parameter": "body", "kind": "column", "value": {"path": "body"}} + ], + "output": { + "kind": "named_struct", + "fields": [ + {"name": "normalized_text", "arrow_type": "utf8", "nullable": false}, + {"name": "token_count", "arrow_type": "int64", "nullable": false} + ] + }, + "group_id": "fg_01K3TEXT", + "columns": {"normalized_text": "search_text"} + }, + "binding_metadata_version": 1, + "input_bindings": [ + {"parameter": "title", "field_path": "title", "arrow_type": "utf8", "nullable": true}, + {"parameter": "body", "field_path": "body", "arrow_type": "utf8", "nullable": true} + ], + "input_schema": { + "fields": [ + {"name": "title", "nullable": true, "type": {"type": "utf8"}}, + {"name": "body", "nullable": true, "type": {"type": "utf8"}} + ] + }, + "output_schema": { + "fields": [ + {"name": "search_text", "nullable": true, "type": {"type": "utf8"}}, + {"name": "token_count", "nullable": true, "type": {"type": "int64"}} + ] + }, + "outputs": [ + {"result_field": "normalized_text", "output_name": "search_text", "output_ordinal": 0}, + {"result_field": "token_count", "output_name": "token_count", "output_ordinal": 1} + ] + } +} diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_scalar_declaration_request.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_scalar_declaration_request.json new file mode 100644 index 000000000..0aaa0cf72 --- /dev/null +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_scalar_declaration_request.json @@ -0,0 +1,41 @@ +{ + "new_columns": [ + {"name": "embedding", "all_null": true} + ], + "function": { + "application": { + "function": {"name": "embed", "version": "fv_01K3EXACT"}, + "inputs": [ + {"parameter": "text", "kind": "column", "value": {"path": "description"}} + ], + "output": {"kind": "scalar", "arrow_type": "list", "nullable": false}, + "group_id": "fg_scalar" + }, + "binding_metadata_version": 1, + "input_bindings": [ + {"parameter": "text", "field_path": "description", "arrow_type": "utf8", "nullable": true} + ], + "input_schema": { + "fields": [ + {"name": "text", "nullable": true, "type": {"type": "utf8"}} + ] + }, + "output_schema": { + "fields": [ + { + "name": "embedding", + "nullable": true, + "type": { + "type": "list", + "fields": [ + {"name": "item", "nullable": false, "type": {"type": "float32"}} + ] + } + } + ] + }, + "outputs": [ + {"result_field": "$value", "output_name": "embedding", "output_ordinal": 0} + ] + } +}