From 685cb01d6d4aa354ba818b013f9ab002bbc6e8c1 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 21 Aug 2026 17:01:04 +0800 Subject: [PATCH 1/2] 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} + ] + } +} From a588208de68e1d887a17d02b3b326812b42f56fc Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 21 Aug 2026 17:19:13 +0800 Subject: [PATCH 2/2] feat: add scalar function authoring and catalog client (#3991) ## Problem The canonical Function wire values and typed remote Job contract do not yet provide a Python authoring surface or catalog client, so users cannot package a scalar callable, register it, or reopen the exact immutable Function version. ## Behavior This adds scalar-only `@udf` authoring with deterministic annotation or explicit Arrow schema validation, content-addressed Python artifacts, and an internal scalar-to-Arrow-batch adapter descriptor. Registration payloads model non-secret environment values and secret names only. Remote connections can submit `create_function_async` and receive a typed `Job`, then reopen that exact version by name and version ID. Synchronous connections can call `create_function` to submit and wait for the immutable version in one operation. Local Function catalog operations return a stable `NotSupported` error. Shared Rust/Python golden payloads and mocked catalog responses freeze the request, typed terminal result, and exact lookup contract. ## Validation - Rust formatting, remote check, clippy, and focused LDB-1/LDB-2 tests - Python formatting, lint, and focused LDB-1/LDB-2 tests - Python API documentation build --- docs/src/python/python.md | 12 + python/python/lancedb/__init__.py | 4 + python/python/lancedb/_lancedb.pyi | 9 + python/python/lancedb/db.py | 56 +- python/python/lancedb/functions.py | 553 +++++++++++++++++- python/python/lancedb/job.py | 53 +- python/python/lancedb/remote/db.py | 9 + .../tests/test_first_class_function_slice2.py | 254 ++++++++ python/src/connection.rs | 32 + python/src/job.rs | 53 ++ python/src/lib.rs | 1 + rust/lancedb/src/connection.rs | 27 + rust/lancedb/src/database.rs | 21 + rust/lancedb/src/function.rs | 50 ++ rust/lancedb/src/job.rs | 3 +- rust/lancedb/src/remote/db.rs | 88 +++ .../tests/first_class_function_slice2.rs | 81 +++ ...nction_registration_request.canonical.json | 1 + .../remote_function_registration_request.json | 46 ++ 19 files changed, 1326 insertions(+), 27 deletions(-) create mode 100644 python/python/tests/test_first_class_function_slice2.py create mode 100644 rust/lancedb/tests/first_class_function_slice2.rs create mode 100644 rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_registration_request.canonical.json create mode 100644 rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_registration_request.json diff --git a/docs/src/python/python.md b/docs/src/python/python.md index a99c0236a..70b2a7207 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -68,6 +68,18 @@ listing a storage directory. ::: lancedb.functions.PythonEnvironmentSpec +::: lancedb.functions.udf + +::: lancedb.functions.UdfDefinition + +::: lancedb.functions.FunctionRegistrationRequest + +::: lancedb.functions.FunctionArtifactRequest + +::: lancedb.functions.FunctionArtifactContent + +::: lancedb.functions.PythonAdapterSpec + ::: lancedb.functions.FunctionVersion ::: lancedb.functions.PythonRuntimeSpec diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index a8a336a6d..0ceda4558 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -23,10 +23,14 @@ from .expr import Expr, col, lit, func from .schema import blob, vector, BlobType from .job import AsyncJob, Job from .functions import ( + FunctionArtifactRequest as FunctionArtifactRequest, FunctionApplication as FunctionApplication, FunctionBinding as FunctionBinding, + FunctionRegistrationRequest as FunctionRegistrationRequest, FunctionVersion as FunctionVersion, PythonRuntimeSpec as PythonRuntimeSpec, + UdfDefinition as UdfDefinition, + udf as udf, ) from .table import AsyncTable, Table from .types import BaseTokenizerType diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index ea5d3e972..59537f45a 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -147,6 +147,8 @@ class Connection(object): limit: Optional[int], ) -> list[str]: ... # Deprecated: Use list_tables instead def job(self, job_id: str) -> Job: ... + async def create_function_async(self, request_json: str) -> FunctionJob: ... + async def get_function(self, name: str, version: str) -> str: ... async def list_jobs(self) -> List[JobInfo]: ... async def get_job(self, job_id: str) -> Optional[JobDescription]: ... async def cancel_job(self, job_id: str) -> bool: ... @@ -226,6 +228,13 @@ class Job: async def wait(self) -> None: ... async def cancel(self) -> None: ... +class FunctionJob: + @property + def id(self) -> Optional[str]: ... + async def status(self) -> str: ... + async def wait(self) -> str: ... + async def cancel(self) -> None: ... + class JobInfo: @property def job_id(self) -> str: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 14b6c0b0d..af18b6944 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -45,7 +45,8 @@ from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError from . import __version__ from ._lancedb import connect as lancedb_connect # type: ignore -from .job import AsyncJob, Job +from .functions import FunctionVersion, UdfDefinition +from .job import AsyncJob, Job, _function_job from .table import ( AsyncTable, LanceTable, @@ -616,6 +617,31 @@ class DBConnection(EnforceOverrides): """ raise NotImplementedError("serialize is not supported for this connection type") + def create_function(self, definition: UdfDefinition) -> FunctionVersion: + """Register a scalar Python UDF and wait for its immutable version. + + This is the blocking counterpart of :meth:`create_function_async`. + Local connections raise ``NotImplementedError``. + """ + return self.create_function_async(definition).wait() + + def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]: + """Register a scalar Python UDF through the remote Function catalog. + + Submission returns a typed job. The immutable Function version becomes + available only when :meth:`Job.wait` succeeds. Local connections raise + ``NotImplementedError``. + """ + raise NotImplementedError( + "Function catalog operations are not supported for this connection type" + ) + + def get_function(self, name: str, *, version: str) -> FunctionVersion: + """Open one exact immutable Function version from the remote catalog.""" + raise NotImplementedError( + "Function catalog operations are not supported for this connection type" + ) + def job(self, job_id: str) -> Job: """A [Job][lancedb.job.Job] handle for a server-side job by id. @@ -1256,6 +1282,15 @@ class LanceDBConnection(DBConnection): """ return Job(self._conn.job(job_id)) + @override + def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]: + job = LOOP.run(self._conn.create_function_async(definition)) + return Job(job) + + @override + def get_function(self, name: str, *, version: str) -> FunctionVersion: + return LOOP.run(self._conn.get_function(name, version=version)) + @override def list_jobs(self) -> List[JobInfo]: """List server-side jobs across the database's tables.""" @@ -2023,6 +2058,25 @@ class AsyncConnection(object): """ return AsyncJob(self._inner.job(job_id)) + async def create_function_async( + self, definition: UdfDefinition + ) -> AsyncJob[FunctionVersion]: + """Register a scalar Python UDF through the remote Function catalog. + + The returned typed job resolves to the immutable Function version. + Local connections raise ``NotImplementedError``. + """ + if not isinstance(definition, UdfDefinition): + raise TypeError("create_function_async requires a @udf definition") + inner = await self._inner.create_function_async( + definition.registration_request.to_canonical_json() + ) + return _function_job(inner) + + async def get_function(self, name: str, *, version: str) -> FunctionVersion: + """Open one exact immutable Function version from the remote catalog.""" + return FunctionVersion.from_json(await self._inner.get_function(name, version)) + async def list_jobs(self) -> List[JobInfo]: """List server-side jobs across the database's tables.""" return await self._inner.list_jobs() diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index df781f665..1613e03b4 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -9,10 +9,32 @@ environment bake, secret resolution, and execution are owned by Sophon. from __future__ import annotations +import ast +import base64 +import functools +import hashlib +import inspect import json +import math +import re +import sys +import textwrap +import types from collections.abc import Mapping -from typing import Any, Optional +from datetime import date, datetime +from typing import ( + Annotated, + Any, + Callable, + Optional, + Union, + get_args, + get_origin, + get_type_hints, + overload, +) +import pyarrow as pa from pydantic import ( BaseModel, ConfigDict, @@ -126,18 +148,10 @@ 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" + model_config = ConfigDict(extra="allow", frozen=True) def _unknown_field_names(self) -> set[str]: - if _PYDANTIC_V2: - return set((self.__pydantic_extra__ or {}).keys()) - return set(self.__dict__) - set(self.__fields__) + return set((self.__pydantic_extra__ or {}).keys()) class FunctionArtifact(_RemoteValue): @@ -148,6 +162,30 @@ class FunctionArtifact(_RemoteValue): entrypoint: str +class FunctionArtifactContent(_RemoteValue): + """Encoded artifact bytes uploaded during remote registration.""" + + encoding: str + data: str + + +class PythonAdapterSpec(_RemoteValue): + """Internal scalar-callable to Arrow-batch adapter selection.""" + + kind: str + version: _UInt32 + + +class FunctionArtifactRequest(_RemoteValue): + """Source artifact uploaded while registering a Function.""" + + kind: str + digest: str + entrypoint: str + content: FunctionArtifactContent + adapter: PythonAdapterSpec + + class FunctionParameter(_RemoteValue): name: str arrow_type: str @@ -228,6 +266,20 @@ class FunctionVersion(_RemoteValue): created_at: str +class FunctionRegistrationRequest(_RemoteValue): + """Stable remote registration envelope produced by :func:`udf`. + + Only secret names are represented. Secret values are resolved inside the + remote service and have no client request field. + """ + + name: str + artifact: FunctionArtifactRequest + signature: FunctionSignature + runtime: PythonRuntimeSpec + required_secrets: tuple[str, ...] = () + + class FunctionVersionRef(_OpenRemoteValue): name: str version: str @@ -361,13 +413,489 @@ class RefreshColumnResult(_RemoteValue): return self.published_version +_FUNCTION_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_.-]*$") +_SECRET_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +def _canonical_arrow_type(data_type: pa.DataType) -> str: + primitive_types = ( + (pa.bool_(), "bool"), + (pa.int8(), "int8"), + (pa.int16(), "int16"), + (pa.int32(), "int32"), + (pa.int64(), "int64"), + (pa.uint8(), "uint8"), + (pa.uint16(), "uint16"), + (pa.uint32(), "uint32"), + (pa.uint64(), "uint64"), + (pa.float16(), "float16"), + (pa.float32(), "float32"), + (pa.float64(), "float64"), + (pa.string(), "utf8"), + (pa.large_utf8(), "large_utf8"), + (pa.binary(), "binary"), + (pa.large_binary(), "large_binary"), + (pa.date32(), "date32"), + (pa.date64(), "date64"), + ) + for candidate, name in primitive_types: + if data_type == candidate: + return name + if pa.types.is_fixed_size_binary(data_type): + return f"fixed_size_binary[{data_type.byte_width}]" + if pa.types.is_list(data_type): + return f"list<{_canonical_arrow_type(data_type.value_type)}>" + if pa.types.is_large_list(data_type): + return f"large_list<{_canonical_arrow_type(data_type.value_type)}>" + if pa.types.is_fixed_size_list(data_type): + return ( + f"fixed_size_list<{_canonical_arrow_type(data_type.value_type)}>" + f"[{data_type.list_size}]" + ) + if pa.types.is_struct(data_type): + fields = ",".join( + f"{field.name}:{_canonical_arrow_type(field.type)}" for field in data_type + ) + return f"struct<{fields}>" + if pa.types.is_timestamp(data_type): + timezone = f",tz={data_type.tz}" if data_type.tz is not None else "" + return f"timestamp[{data_type.unit}{timezone}]" + if pa.types.is_time32(data_type) or pa.types.is_time64(data_type): + return f"time[{data_type.unit}]" + if pa.types.is_duration(data_type): + return f"duration[{data_type.unit}]" + if pa.types.is_decimal(data_type): + bit_width = data_type.bit_width + return f"decimal{bit_width}({data_type.precision},{data_type.scale})" + raise TypeError(f"unsupported Arrow type for Function signature: {data_type}") + + +def _annotation_type(annotation: Any) -> tuple[pa.DataType, bool]: + nullable = False + origin = get_origin(annotation) + if origin in (Union, types.UnionType): + arguments = get_args(annotation) + non_none = tuple( + argument for argument in arguments if argument is not type(None) + ) + if len(non_none) != 1 or len(non_none) == len(arguments): + raise TypeError(f"unsupported union annotation: {annotation!r}") + annotation = non_none[0] + nullable = True + + origin = get_origin(annotation) + if origin is Annotated: + base, *metadata = get_args(annotation) + arrow_types = [value for value in metadata if isinstance(value, pa.DataType)] + if len(arrow_types) != 1: + raise TypeError( + "Annotated Function types require exactly one PyArrow DataType" + ) + _, base_nullable = _annotation_type(base) + return arrow_types[0], nullable or base_nullable + + if isinstance(annotation, pa.DataType): + return annotation, nullable + if annotation is bool: + return pa.bool_(), nullable + if annotation is int: + return pa.int64(), nullable + if annotation is float: + return pa.float64(), nullable + if annotation is str: + return pa.string(), nullable + if annotation is bytes: + return pa.binary(), nullable + if annotation is date: + return pa.date32(), nullable + if annotation is datetime: + return pa.timestamp("us"), nullable + if get_origin(annotation) is list: + arguments = get_args(annotation) + if len(arguments) != 1: + raise TypeError(f"unsupported list annotation: {annotation!r}") + value_type, value_nullable = _annotation_type(arguments[0]) + if value_nullable: + raise TypeError("nullable Function list elements are not supported") + return pa.list_(value_type), nullable + raise TypeError(f"unsupported Function annotation: {annotation!r}") + + +def _callable_parameters(function: Callable[..., Any]) -> tuple[inspect.Parameter, ...]: + parameters = tuple(inspect.signature(function).parameters.values()) + for parameter in parameters: + if parameter.kind in ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + ): + raise TypeError("Function callables require named, non-variadic parameters") + if parameter.default is not inspect.Parameter.empty: + raise TypeError("Function callable defaults are not supported") + return parameters + + +def _function_output(output: pa.DataType | pa.Field | pa.Schema) -> FunctionOutput: + if isinstance(output, pa.Schema): + fields = tuple(output) + elif isinstance(output, pa.Field) and pa.types.is_struct(output.type): + if output.nullable: + raise ValueError("Function output must be non-nullable") + fields = tuple(output.type) + elif isinstance(output, pa.DataType) and pa.types.is_struct(output): + fields = tuple(output) + else: + field = ( + output + if isinstance(output, pa.Field) + else pa.field("result", output, nullable=False) + ) + if not isinstance(field, pa.Field): + raise TypeError( + "output_schema must be a PyArrow DataType, Field, or Schema" + ) + if field.nullable: + raise ValueError("Function output must be non-nullable") + return FunctionOutput( + kind="scalar", + arrow_type=_canonical_arrow_type(field.type), + nullable=False, + ) + + if not fields: + raise ValueError("named-struct Function output must contain at least one field") + if any(field.nullable for field in fields): + raise ValueError("Function output fields must be non-nullable") + names = [field.name for field in fields] + if len(set(names)) != len(names): + raise ValueError("Function output field names must be unique") + return FunctionOutput( + kind="named_struct", + fields=tuple( + FunctionResultField( + name=field.name, + arrow_type=_canonical_arrow_type(field.type), + nullable=False, + ) + for field in fields + ), + ) + + +def _infer_signature( + function: Callable[..., Any], + input_schema: Optional[pa.Schema], + output_schema: Optional[pa.DataType | pa.Field | pa.Schema], +) -> FunctionSignature: + parameters = _callable_parameters(function) + if (input_schema is None) != (output_schema is None): + raise ValueError("input_schema and output_schema must be provided together") + + if input_schema is not None: + if not isinstance(input_schema, pa.Schema): + raise TypeError("input_schema must be a PyArrow Schema") + expected = tuple(parameter.name for parameter in parameters) + actual = tuple(input_schema.names) + if actual != expected: + raise ValueError( + "input_schema fields must exactly match callable parameters in order: " + f"expected {expected!r}, got {actual!r}" + ) + inputs = tuple( + FunctionParameter( + name=field.name, + arrow_type=_canonical_arrow_type(field.type), + nullable=field.nullable, + ) + for field in input_schema + ) + return FunctionSignature(inputs=inputs, output=_function_output(output_schema)) + + try: + annotations = get_type_hints(function, include_extras=True) + except Exception as error: + raise TypeError(f"failed to resolve Function annotations: {error}") from error + missing = [ + parameter.name for parameter in parameters if parameter.name not in annotations + ] + if missing or "return" not in annotations: + names = missing + ([] if "return" in annotations else ["return"]) + raise TypeError(f"missing Function annotations: {names!r}") + inputs = [] + for parameter in parameters: + data_type, nullable = _annotation_type(annotations[parameter.name]) + inputs.append( + FunctionParameter( + name=parameter.name, + arrow_type=_canonical_arrow_type(data_type), + nullable=nullable, + ) + ) + output_type, output_nullable = _annotation_type(annotations["return"]) + if output_nullable: + raise ValueError("Function output must be non-nullable") + return FunctionSignature( + inputs=tuple(inputs), + output=_function_output(pa.field("result", output_type, nullable=False)), + ) + + +def _is_udf_decorator(node: ast.expr) -> bool: + if isinstance(node, ast.Call): + node = node.func + return (isinstance(node, ast.Name) and node.id == "udf") or ( + isinstance(node, ast.Attribute) and node.attr == "udf" + ) + + +def _literal_source(value: Any) -> str: + if value is None or type(value) in (bool, int, str, bytes): + return repr(value) + if type(value) is float and math.isfinite(value): + return repr(value) + if type(value) is tuple: + children = ", ".join(_literal_source(child) for child in value) + if len(value) == 1: + children += "," + return f"({children})" + raise TypeError( + "Function source references an unsupported global value of type " + f"{type(value).__name__}" + ) + + +def _package_source(function: Callable[..., Any]) -> bytes: + if not inspect.isfunction(function) or inspect.iscoroutinefunction(function): + raise TypeError("@udf requires a synchronous Python function") + try: + source = textwrap.dedent(inspect.getsource(function)) + except (OSError, TypeError) as error: + raise ValueError("@udf requires inspectable Python source") from error + module = ast.parse(source) + definitions = [ + node + for node in module.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and node.name == function.__name__ + ] + if len(definitions) != 1 or not isinstance(definitions[0], ast.FunctionDef): + raise ValueError("@udf source must contain exactly one synchronous function") + definition = definitions[0] + if any(not _is_udf_decorator(decorator) for decorator in definition.decorator_list): + raise ValueError("@udf cannot package additional Python decorators") + definition.decorator_list = [] + + closure = inspect.getclosurevars(function) + if closure.nonlocals: + raise ValueError("@udf cannot package functions that capture closure values") + if closure.unbound: + raise ValueError( + f"@udf source contains unresolved global names: {sorted(closure.unbound)!r}" + ) + globals_source = [] + for name, value in sorted(closure.globals.items()): + if isinstance(value, types.ModuleType): + globals_source.append(f"import {value.__name__} as {name}") + else: + globals_source.append(f"{name} = {_literal_source(value)}") + + function_source = ast.unparse(definition) + parts = ["from __future__ import annotations"] + if globals_source: + parts.extend(["", *globals_source]) + parts.extend(["", function_source, ""]) + return "\n".join(parts).encode("utf-8") + + +class UdfDefinition: + """A scalar Python callable prepared for remote Function registration. + + Instances are created with :func:`udf`. Calling an instance executes the + original scalar Python function, which keeps local unit testing ordinary. + Remote execution adapts that scalar callable to the internal Arrow batch + ABI described by the registration artifact. + """ + + def __init__( + self, + function: Callable[..., Any], + *, + name: Optional[str], + input_schema: Optional[pa.Schema], + output_schema: Optional[pa.DataType | pa.Field | pa.Schema], + pip: tuple[str, ...], + env: Mapping[str, str], + secrets: tuple[str, ...], + python_version: Optional[str], + ): + function_name = name or function.__name__ + if not _FUNCTION_NAME.fullmatch(function_name): + raise ValueError(f"invalid Function name: {function_name!r}") + packages = tuple(sorted(set(pip))) + if any(not package or package != package.strip() for package in packages): + raise ValueError("pip requirements must be non-empty and trimmed") + environment = dict(env) + if any( + not isinstance(key, str) or not isinstance(value, str) + for key, value in environment.items() + ): + raise TypeError("Function env keys and values must be strings") + required_secrets = tuple(sorted(set(secrets))) + invalid_secrets = [ + secret for secret in required_secrets if not _SECRET_NAME.fullmatch(secret) + ] + if invalid_secrets: + raise ValueError(f"invalid Function secret names: {invalid_secrets!r}") + overlap = set(environment) & set(required_secrets) + if overlap: + raise ValueError( + f"Function env and secret names must be disjoint: {sorted(overlap)!r}" + ) + + signature = _infer_signature(function, input_schema, output_schema) + source = _package_source(function) + digest = f"sha256:{hashlib.sha256(source).hexdigest()}" + runtime = PythonRuntimeSpec( + kind="python", + python_version=python_version + or f"{sys.version_info.major}.{sys.version_info.minor}", + environment=PythonEnvironmentSpec(kind="pip", packages=packages), + env=environment, + ) + self._function = function + self._request = FunctionRegistrationRequest( + name=function_name, + artifact=FunctionArtifactRequest( + kind="python_callable", + digest=digest, + entrypoint=function.__name__, + content=FunctionArtifactContent( + encoding="base64", + data=base64.b64encode(source).decode("ascii"), + ), + adapter=PythonAdapterSpec( + kind="scalar_to_arrow_batch", + version=1, + ), + ), + signature=signature, + runtime=runtime, + required_secrets=required_secrets, + ) + functools.update_wrapper(self, function) + + @property + def registration_request(self) -> FunctionRegistrationRequest: + """The immutable request sent by ``create_function_async``.""" + return self._request + + def __call__(self, *args, **kwargs): + return self._function(*args, **kwargs) + + +@overload +def udf(function: Callable[..., Any]) -> UdfDefinition: ... + + +@overload +def udf( + function: None = None, + *, + name: Optional[str] = None, + input_schema: Optional[pa.Schema] = None, + output_schema: Optional[pa.DataType | pa.Field | pa.Schema] = None, + pip: tuple[str, ...] | list[str] = (), + env: Optional[Mapping[str, str]] = None, + secrets: tuple[str, ...] | list[str] = (), + python_version: Optional[str] = None, +) -> Callable[[Callable[..., Any]], UdfDefinition]: ... + + +def udf( + function: Optional[Callable[..., Any]] = None, + *, + name: Optional[str] = None, + input_schema: Optional[pa.Schema] = None, + output_schema: Optional[pa.DataType | pa.Field | pa.Schema] = None, + pip: tuple[str, ...] | list[str] = (), + env: Optional[Mapping[str, str]] = None, + secrets: tuple[str, ...] | list[str] = (), + python_version: Optional[str] = None, +): + """Prepare a scalar Python callable for remote Function registration. + + Input and output signatures are inferred from supported annotations. For + Arrow types annotations cannot express precisely, pass ``input_schema`` + and ``output_schema`` together. Nullable outputs are rejected because V1 + uses physical NULL to represent unassigned computed-column rows. + + Parameters + ---------- + function : Callable, optional + The synchronous scalar callable to package. + name : str, optional + The remote Function name. Defaults to the callable name. + input_schema : pyarrow.Schema, optional + Explicit input fields in the exact order of the callable parameters. + Must be provided together with ``output_schema``. + output_schema : pyarrow.DataType, pyarrow.Field, or pyarrow.Schema, optional + Explicit scalar or named-struct output. Must be non-nullable and be + provided together with ``input_schema``. + pip : sequence of str, optional + Pip requirements for the remote environment. + env : mapping of str to str, optional + Non-secret environment variables. Use ``secrets`` for credentials. + secrets : sequence of str, optional + Names of secrets resolved by the remote service. Secret values are not + accepted by this API or included in the registration request. + python_version : str, optional + Remote Python major/minor version. Defaults to the client version. + + Returns + ------- + UdfDefinition + A callable definition accepted by + :meth:`lancedb.db.DBConnection.create_function`, + :meth:`lancedb.db.AsyncConnection.create_function_async` and + :meth:`lancedb.db.DBConnection.create_function_async`. + + Examples + -------- + >>> from lancedb import udf + >>> @udf(pip=["numpy==2.2.0"], secrets=["MODEL_TOKEN"]) + ... def score(value: float) -> float: + ... return value * 2 + >>> score(1.5) + 3.0 + """ + + def decorate(target: Callable[..., Any]) -> UdfDefinition: + return UdfDefinition( + target, + name=name, + input_schema=input_schema, + output_schema=output_schema, + pip=tuple(pip), + env={} if env is None else env, + secrets=tuple(secrets), + python_version=python_version, + ) + + if function is None: + return decorate + return decorate(function) + + __all__ = [ "ApplicationInput", "FunctionApplication", "FunctionArtifact", + "FunctionArtifactContent", + "FunctionArtifactRequest", "FunctionBinding", "FunctionOutput", "FunctionParameter", + "FunctionRegistrationRequest", "FunctionResultField", "FunctionSignature", "FunctionVersion", @@ -375,6 +903,9 @@ __all__ = [ "InputBinding", "OutputMapping", "PythonEnvironmentSpec", + "PythonAdapterSpec", "PythonRuntimeSpec", "RefreshColumnResult", + "UdfDefinition", + "udf", ] diff --git a/python/python/lancedb/job.py b/python/python/lancedb/job.py index d33b62cbf..7bd600a74 100644 --- a/python/python/lancedb/job.py +++ b/python/python/lancedb/job.py @@ -5,20 +5,23 @@ import asyncio from datetime import timedelta -from typing import Optional +from typing import Any, Generic, Optional, TypeVar, cast from lancedb.background_loop import LOOP from . import _lancedb +from .functions import FunctionVersion + +T = TypeVar("T") -class AsyncJob: +class AsyncJob(Generic[T]): """A handle to an operation that may still be running. The operation may already be complete when the handle is created. """ - def __init__(self, inner: Optional["_lancedb.Job"]): + def __init__(self, inner: Optional[Any]): self._inner = inner @property @@ -44,18 +47,20 @@ class AsyncJob: return "finished" return await self._inner.status() - async def wait(self, timeout: Optional[timedelta] = None): + async def wait(self, timeout: Optional[timedelta] = None) -> T: """Wait until the operation reaches a terminal state. Raises `JobFailedError` if the operation failed, `JobCancelledError` if it was cancelled, and `TimeoutError` if `timeout` elapses first. """ if self._inner is None: - return + return cast(T, None) if timeout is None: - await self._inner.wait() - else: - await asyncio.wait_for(self._inner.wait(), timeout.total_seconds()) + return cast(T, await self._inner.wait()) + return cast( + T, + await asyncio.wait_for(self._inner.wait(), timeout.total_seconds()), + ) async def cancel(self): """Request cancellation. Cancelling a finished operation is a no-op.""" @@ -64,10 +69,10 @@ class AsyncJob: await self._inner.cancel() -class Job: +class Job(Generic[T]): """Synchronous counterpart of `AsyncJob`.""" - def __init__(self, inner: Optional[AsyncJob]): + def __init__(self, inner: Optional[AsyncJob[T]]): self._inner = inner @property @@ -88,18 +93,40 @@ class Job: return "finished" return LOOP.run(self._inner.status()) - def wait(self, timeout: Optional[timedelta] = None): + def wait(self, timeout: Optional[timedelta] = None) -> T: """Block until the operation reaches a terminal state. Raises `JobFailedError` if the operation failed, `JobCancelledError` if it was cancelled, and `TimeoutError` if `timeout` elapses first. """ if self._inner is None: - return - LOOP.run(self._inner.wait(timeout)) + return cast(T, None) + return LOOP.run(self._inner.wait(timeout)) def cancel(self): """Request cancellation. Cancelling a finished operation is a no-op.""" if self._inner is None: return LOOP.run(self._inner.cancel()) + + +class _FunctionJobAdapter: + def __init__(self, inner: "_lancedb.FunctionJob"): + self._inner = inner + + @property + def id(self) -> Optional[str]: + return self._inner.id + + async def status(self) -> str: + return await self._inner.status() + + async def wait(self) -> FunctionVersion: + return FunctionVersion.from_json(await self._inner.wait()) + + async def cancel(self): + await self._inner.cancel() + + +def _function_job(inner: "_lancedb.FunctionJob") -> AsyncJob[FunctionVersion]: + return AsyncJob(_FunctionJobAdapter(inner)) diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 16ad65dcb..822756a34 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -23,6 +23,7 @@ import pyarrow as pa from ..common import DATA from ..db import DBConnection, LOOP +from ..functions import FunctionVersion, UdfDefinition from ..job import AsyncJob, Job if TYPE_CHECKING: @@ -713,6 +714,14 @@ class RemoteDBConnection(DBConnection): """ return Job(self._conn.job(job_id)) + @override + def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]: + return Job(LOOP.run(self._conn.create_function_async(definition))) + + @override + def get_function(self, name: str, *, version: str) -> FunctionVersion: + return LOOP.run(self._conn.get_function(name, version=version)) + @override def list_jobs(self) -> List["JobInfo"]: """List server-side jobs across the database's tables.""" diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py new file mode 100644 index 000000000..612a711af --- /dev/null +++ b/python/python/tests/test_first_class_function_slice2.py @@ -0,0 +1,254 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +from __future__ import annotations + +import contextlib +import http.server +import json +from pathlib import Path +import threading +from typing import Optional + +import pyarrow as pa +import pytest + +import lancedb +from lancedb.functions import UdfDefinition, udf + + +FIXTURES = ( + Path(__file__).parents[3] + / "rust" + / "lancedb" + / "tests" + / "fixtures" + / "first_class_functions" + / "v1" +) + + +@udf( + pip=["numpy>=2"], + env={"MODE": "test"}, + secrets=["API_TOKEN"], + python_version="3.12", +) +def normalize_score(value: float) -> float: + return value / 100.0 + + +def _assert_no_secret_values(value): + if isinstance(value, dict): + for key, child in value.items(): + assert key not in { + "secret_value", + "secret_values", + "resolved_secret", + "resolved_secrets", + } + _assert_no_secret_values(child) + elif isinstance(value, list): + for child in value: + _assert_no_secret_values(child) + + +def test_scalar_udf_matches_shared_registration_golden_and_remains_callable(): + assert isinstance(normalize_score, UdfDefinition) + assert normalize_score(25.0) == 0.25 + assert ( + normalize_score.registration_request.to_canonical_json() + == (FIXTURES / "remote_function_registration_request.canonical.json") + .read_text() + .strip() + ) + request = json.loads(normalize_score.registration_request.to_canonical_json()) + assert request["artifact"]["adapter"] == { + "kind": "scalar_to_arrow_batch", + "version": 1, + } + assert request["required_secrets"] == ["API_TOKEN"] + _assert_no_secret_values(request) + + +def test_explicit_arrow_schema_is_deterministic(): + input_schema = pa.schema([pa.field("value", pa.float32(), nullable=True)]) + output_schema = pa.field("embedding", pa.list_(pa.float32(), 3), nullable=False) + + @udf(input_schema=input_schema, output_schema=output_schema) + def explicit(value): + return [value, value, value] + + signature = explicit.registration_request.signature + assert signature.inputs[0].arrow_type == "float32" + assert signature.inputs[0].nullable is True + assert signature.output.arrow_type == "fixed_size_list[3]" + assert signature.output.nullable is False + + +def test_annotation_and_explicit_schema_validation_fail_closed(): + with pytest.raises(TypeError, match="missing Function annotations"): + + @udf + def missing(value): + return value + + with pytest.raises(TypeError, match="unsupported Function annotation"): + + @udf + def unsupported(value: set[str]) -> str: + return "" + + with pytest.raises(ValueError, match="output must be non-nullable"): + + @udf + def nullable_output(value: int) -> Optional[int]: + return value + + with pytest.raises(ValueError, match="provided together"): + + @udf(input_schema=pa.schema([pa.field("value", pa.int64())])) + def partial_schema(value): + return value + + with pytest.raises(ValueError, match="exactly match callable parameters"): + + @udf( + input_schema=pa.schema([pa.field("other", pa.int64())]), + output_schema=pa.int64(), + ) + def wrong_name(value): + return value + + with pytest.raises(ValueError, match="output must be non-nullable"): + + @udf( + input_schema=pa.schema([pa.field("value", pa.int64())]), + output_schema=pa.field("result", pa.int64(), nullable=True), + ) + def nullable_explicit(value): + return value + + +def test_environment_rejects_secret_value_overlap(): + with pytest.raises(ValueError, match="must be disjoint"): + + @udf(env={"TOKEN": "plaintext"}, secrets=["TOKEN"]) + def overlapping(value: int) -> int: + return value + + +def test_local_function_catalog_operations_are_not_supported(tmp_path): + db = lancedb.connect(tmp_path) + message = "Function catalog operations are not supported by this database" + with pytest.raises(NotImplementedError, match=message): + db.create_function(normalize_score) + with pytest.raises(NotImplementedError, match=message): + db.create_function_async(normalize_score) + with pytest.raises(NotImplementedError, match=message): + db.get_function("normalize_score", version="fv_exact") + + +@contextlib.contextmanager +def _mock_remote_function_catalog(): + state = {"requests": [], "version": None} + + class Handler(http.server.BaseHTTPRequestHandler): + def log_message(self, *args): + pass + + def do_POST(self): + length = int(self.headers.get("Content-Length", "0")) + body = json.loads(self.rfile.read(length) or b"{}") + state["requests"].append((self.path, body)) + status = 200 + if self.path == "/v1/function/create": + state["version"] = { + "name": body["name"], + "version": "fv_exact", + "artifact": { + key: body["artifact"][key] + for key in ("kind", "digest", "entrypoint") + }, + "signature": body["signature"], + "runtime": body["runtime"], + "runtime_digest": "sha256:runtime", + "environment_digest": "sha256:environment", + "required_secrets": body.get("required_secrets", []), + "created_at": "2026-08-21T00:00:00Z", + } + response = {"job_id": "job-register"} + status = 202 + elif self.path == "/v1/jobs/describe": + assert body == {"job_id": "job-register"} + response = { + "job_id": "job-register", + "job_type": "create_function", + "job_state": "DONE", + "result": state["version"], + } + elif self.path == "/v1/function/describe": + assert body == { + "name": "normalize_score", + "version": "fv_exact", + } + response = state["version"] + else: + status = 404 + response = {"error": "not found"} + encoded = json.dumps(response).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(encoded))) + self.end_headers() + self.wfile.write(encoded) + + with http.server.HTTPServer(("localhost", 0), Handler) as server: + thread = threading.Thread(target=server.serve_forever) + thread.start() + try: + yield f"http://localhost:{server.server_address[1]}", state + finally: + server.shutdown() + thread.join() + + +def test_remote_registration_job_and_exact_version_reopen_round_trip(): + with _mock_remote_function_catalog() as (host, state): + db = lancedb.connect( + "db://dev", + api_key="fake", + host_override=host, + client_config={"retry_config": {"retries": 0}}, + ) + registration = db.create_function_async(normalize_score) + assert registration.id == "job-register" + created = registration.wait() + reopened = db.get_function("normalize_score", version=created.version) + + assert created == reopened + assert reopened.name == "normalize_score" + assert reopened.version == "fv_exact" + create_request = state["requests"][0][1] + assert create_request == json.loads( + normalize_score.registration_request.to_canonical_json() + ) + _assert_no_secret_values(create_request) + + +def test_blocking_remote_registration_returns_function_version(): + with _mock_remote_function_catalog() as (host, state): + db = lancedb.connect( + "db://dev", + api_key="fake", + host_override=host, + client_config={"retry_config": {"retries": 0}}, + ) + created = db.create_function(normalize_score) + + assert created.name == "normalize_score" + assert created.version == "fv_exact" + assert [path for path, _ in state["requests"]] == [ + "/v1/function/create", + "/v1/jobs/describe", + ] diff --git a/python/src/connection.rs b/python/src/connection.rs index dbda29ba6..87870b800 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -563,6 +563,38 @@ impl Connection { Ok(crate::job::Job::new(inner.job(job_id).infer_error()?)) } + pub fn create_function_async( + self_: PyRef<'_, Self>, + request_json: String, + ) -> PyResult> { + let inner = self_.get_inner()?.clone(); + let request = lancedb::function::FunctionRegistrationRequest::from_json(&request_json) + .infer_error()?; + future_into_py(self_.py(), async move { + inner + .create_function_async(request) + .await + .infer_error() + .map(crate::job::FunctionJob::new) + }) + } + + pub fn get_function( + self_: PyRef<'_, Self>, + name: String, + version: String, + ) -> PyResult> { + let inner = self_.get_inner()?.clone(); + future_into_py(self_.py(), async move { + inner + .get_function(name, version) + .await + .infer_error()? + .to_canonical_json() + .infer_error() + }) + } + pub fn list_jobs(self_: PyRef<'_, Self>) -> PyResult> { let inner = self_.get_inner()?.clone(); future_into_py(self_.py(), async move { diff --git a/python/src/job.rs b/python/src/job.rs index 56ee211f4..2755a28c5 100644 --- a/python/src/job.rs +++ b/python/src/job.rs @@ -13,6 +13,23 @@ pub struct Job { inner: Arc, } +/// Python bridge for a typed remote Function registration job. +/// +/// The public Python layer decodes the canonical JSON returned by `wait` +/// into its immutable `FunctionVersion` model. +#[pyclass] +pub struct FunctionJob { + inner: Arc>, +} + +impl FunctionJob { + pub(crate) fn new(inner: lancedb::Job) -> Self { + Self { + inner: Arc::new(inner), + } + } +} + impl Job { pub(crate) fn new(inner: lancedb::Job) -> Self { Self { @@ -21,6 +38,42 @@ impl Job { } } +#[pymethods] +impl FunctionJob { + #[getter] + pub fn id(&self) -> Option { + self.inner.id().map(str::to_string) + } + + pub fn status(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner.clone(); + future_into_py( + self_.py(), + async move { inner.status().await.infer_error() }, + ) + } + + pub fn wait(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + inner + .wait() + .await + .infer_error()? + .to_canonical_json() + .infer_error() + }) + } + + pub fn cancel(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + inner.cancel().await.infer_error()?; + Ok(()) + }) + } +} + #[pymethods] impl Job { #[getter] diff --git a/python/src/lib.rs b/python/src/lib.rs index a19bf172d..756b3557f 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -47,6 +47,7 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 12ca306b8..8935855a8 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -496,6 +496,33 @@ impl Connection { ) } + /// Register a Python callable as a new immutable Function version. + /// + /// Registration is remote-only and always asynchronous. Waiting on the + /// returned typed job yields the durable [`crate::function::FunctionVersion`]. + /// Local databases return [`Error::NotSupported`]. + pub async fn create_function_async( + &self, + request: crate::function::FunctionRegistrationRequest, + ) -> Result> { + self.internal.create_function_async(request).await + } + + /// Look up one exact immutable Function version in the remote catalog. + /// + /// Both the logical name and server-assigned version id are required; + /// mutable aliases and latest-version lookup are intentionally absent. + /// Local databases return [`Error::NotSupported`]. + pub async fn get_function( + &self, + name: impl AsRef, + version: impl AsRef, + ) -> Result { + self.internal + .get_function(name.as_ref(), version.as_ref()) + .await + } + /// Rename a table in the database. /// /// This is only supported in LanceDB Cloud. diff --git a/rust/lancedb/src/database.rs b/rust/lancedb/src/database.rs index f52c02439..6c4537972 100644 --- a/rust/lancedb/src/database.rs +++ b/rust/lancedb/src/database.rs @@ -241,6 +241,12 @@ fn job_op_not_supported(what: &str) -> Result { }) } +fn function_catalog_not_supported() -> Result { + Err(crate::error::Error::NotSupported { + message: "Function catalog operations are not supported by this database".to_string(), + }) +} + /// The `Database` trait defines the interface for database implementations. /// /// A database is responsible for managing tables and their metadata. @@ -286,6 +292,21 @@ pub trait Database: /// /// See [`CloneTableRequest`] for detailed documentation and examples. async fn clone_table(&self, request: CloneTableRequest) -> Result>; + /// Register an immutable Function version through the remote catalog. + async fn create_function_async( + &self, + _request: crate::function::FunctionRegistrationRequest, + ) -> Result> { + function_catalog_not_supported() + } + /// Look up one exact immutable Function version. + async fn get_function( + &self, + _name: &str, + _version: &str, + ) -> Result { + function_catalog_not_supported() + } /// A [`crate::job::Job`] handle for a server-side job by id, suitable for /// waiting on or cancelling the job. The handle is constructed without a /// server round trip; an unknown id surfaces when the handle is used. diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index fe91f1680..835fca9a2 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -369,6 +369,56 @@ impl FunctionVersion { impl_json!(FunctionVersion); +/// Encoded artifact bytes uploaded with a Function registration request. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct FunctionArtifactContent { + /// Encoding of `data`. V1 Python authoring uses `base64`. + pub encoding: String, + pub data: String, +} + +/// Internal execution adapter selected for a Python callable artifact. +/// +/// The adapter converts the public scalar callable to the Arrow batch ABI +/// used by the remote executor. It is part of the request envelope, not a +/// public batch-UDF authoring mode. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct PythonAdapterSpec { + pub kind: String, + pub version: u32, +} + +/// Python artifact uploaded while registering a Function. +/// +/// Unlike [`FunctionArtifact`], which is the durable artifact identity +/// returned by the catalog, this request value contains the encoded source +/// bytes that Sophon must durably bake before publishing a FunctionVersion. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct FunctionArtifactRequest { + pub kind: String, + pub digest: String, + pub entrypoint: String, + pub content: FunctionArtifactContent, + pub adapter: PythonAdapterSpec, +} + +/// Stable request envelope for remote immutable Function registration. +/// +/// Secret values deliberately have no field in this model. The only secret +/// material the client may send is the ordered set of names Sophon resolves +/// inside the remote runtime. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct FunctionRegistrationRequest { + pub name: String, + pub artifact: FunctionArtifactRequest, + pub signature: FunctionSignature, + pub runtime: PythonRuntimeSpec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub required_secrets: Vec, +} + +impl_json!(FunctionRegistrationRequest); + /// Exact FunctionVersion reference embedded in applications and bindings. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct FunctionVersionRef { diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index 0f880e398..1a76c7683 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -126,8 +126,7 @@ impl Job where T: Clone + DeserializeOwned + Send + Sync + 'static, { - /// Construct a typed remote Job before result-specific submit APIs are added. - #[allow(dead_code)] + /// Construct a typed remote Job for a result-specific submit API. pub(crate) fn new_typed(handle: Box) -> Self { Self { inner: JobInner::Handle { diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 03a13cb4e..08f71368c 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -24,6 +24,7 @@ use crate::database::{ JobDescription, JobInfo, OpenTableRequest, ReadConsistency, TableNamesRequest, }; use crate::error::Result; +use crate::function::{FunctionRegistrationRequest, FunctionVersion}; use crate::job::Job; use crate::remote::job::{DescribeJobResponse, RemoteJob, job_state_to_client}; use crate::remote::util::stream_as_body; @@ -489,6 +490,39 @@ impl Database for RemoteDatabase { }) } + async fn create_function_async( + &self, + request: FunctionRegistrationRequest, + ) -> Result> { + let req = self.client.post("/v1/function/create").json(&request); + let (request_id, response) = self.client.send(req).await?; + let response = self.client.check_response(&request_id, response).await?; + let status = response.status(); + let body = response.text().await.err_to_http(request_id.clone())?; + let job_id = extract_job_id(&body).ok_or_else(|| Error::Http { + source: "Function registration response did not contain a valid job_id".into(), + request_id, + status_code: Some(status), + })?; + Ok(Job::new_typed(Box::new(RemoteJob::new( + self.client.clone(), + job_id, + )))) + } + + async fn get_function(&self, name: &str, version: &str) -> Result { + let req = self + .client + .post("/v1/function/describe") + .json(&serde_json::json!({ + "name": name, + "version": version, + })); + let (request_id, response) = self.client.send(req).await?; + let response = self.client.check_response(&request_id, response).await?; + response.json().await.err_to_http(request_id) + } + fn job(&self, job_id: &str) -> Result { Ok(crate::job::Job::new(Box::new(super::job::RemoteJob::new( self.client.clone(), @@ -2446,6 +2480,60 @@ mod tests { assert_eq!(batches[0].num_rows(), 2); } + #[tokio::test] + async fn test_create_function_async_sends_canonical_request_and_decodes_typed_job() { + const REQUEST: &str = include_str!( + "../../tests/fixtures/first_class_functions/v1/remote_function_registration_request.json" + ); + const FUNCTION_JOB: &str = + include_str!("../../tests/fixtures/first_class_functions/v1/remote_function_job.json"); + let expected: serde_json::Value = serde_json::from_str(REQUEST).unwrap(); + let conn = Connection::new_with_handler(move |request| match request.url().path() { + "/v1/function/create" => { + assert_eq!(request.method(), &reqwest::Method::POST); + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!(body, expected); + http::Response::builder() + .status(202) + .body(r#"{"job_id":"job-function-1"}"#) + .unwrap() + } + "/v1/jobs/describe" => http::Response::builder() + .status(200) + .body(FUNCTION_JOB) + .unwrap(), + path => panic!("unexpected path: {path}"), + }); + let request = crate::function::FunctionRegistrationRequest::from_json(REQUEST).unwrap(); + let job = conn.create_function_async(request).await.unwrap(); + assert_eq!(job.id(), Some("job-function-1")); + let version = job.wait().await.unwrap(); + assert_eq!(version.name(), "embed"); + assert_eq!(version.version(), "fv_01K3EXACT"); + } + + #[tokio::test] + async fn test_get_function_requires_and_sends_exact_version() { + const VERSION: &str = include_str!( + "../../tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json" + ); + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().path(), "/v1/function/describe"); + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!( + body, + serde_json::json!({"name": "embed", "version": "fv_01K3EXACT"}) + ); + http::Response::builder().status(200).body(VERSION).unwrap() + }); + let version = conn.get_function("embed", "fv_01K3EXACT").await.unwrap(); + assert_eq!(version.name(), "embed"); + assert_eq!(version.version(), "fv_01K3EXACT"); + } + #[tokio::test] async fn test_conn_job_waits_to_done() { let polls = Arc::new(AtomicUsize::new(0)); diff --git a/rust/lancedb/tests/first_class_function_slice2.rs b/rust/lancedb/tests/first_class_function_slice2.rs new file mode 100644 index 000000000..3bae57122 --- /dev/null +++ b/rust/lancedb/tests/first_class_function_slice2.rs @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use std::fs; +use std::path::PathBuf; + +use lancedb::Error; +use lancedb::function::FunctionRegistrationRequest; +use serde_json::Value; + +fn fixture(name: &str) -> String { + let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("tests/fixtures/first_class_functions/v1") + .join(name); + fs::read_to_string(path).expect("fixture must be readable") +} + +fn assert_no_secret_values(value: &Value) { + match value { + Value::Object(values) => { + for (key, value) in values { + assert!( + !matches!( + key.as_str(), + "secret_value" | "secret_values" | "resolved_secret" | "resolved_secrets" + ), + "registration requests must not model resolved secret material" + ); + assert_no_secret_values(value); + } + } + Value::Array(values) => values.iter().for_each(assert_no_secret_values), + _ => {} + } +} + +#[test] +fn registration_request_matches_shared_canonical_golden() { + let request = FunctionRegistrationRequest::from_json(&fixture( + "remote_function_registration_request.json", + )) + .expect("registration request"); + assert_eq!(request.name, "normalize_score"); + assert_eq!(request.artifact.adapter.kind, "scalar_to_arrow_batch"); + assert_eq!(request.required_secrets, ["API_TOKEN"]); + assert_eq!( + request.to_canonical_json().expect("canonical request"), + fixture("remote_function_registration_request.canonical.json").trim() + ); + + let value: Value = + serde_json::from_str(&request.to_canonical_json().expect("canonical request")) + .expect("request JSON"); + assert_no_secret_values(&value); +} + +#[tokio::test] +async fn local_function_catalog_operations_return_stable_not_supported() { + let directory = tempfile::tempdir().unwrap(); + let connection = lancedb::connect(directory.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let request = FunctionRegistrationRequest::from_json(&fixture( + "remote_function_registration_request.json", + )) + .unwrap(); + + let create_error = connection.create_function_async(request).await.unwrap_err(); + let lookup_error = connection + .get_function("normalize_score", "fv_exact") + .await + .unwrap_err(); + for error in [create_error, lookup_error] { + assert!(matches!( + error, + Error::NotSupported { message } + if message == "Function catalog operations are not supported by this database" + )); + } +} diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_registration_request.canonical.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_registration_request.canonical.json new file mode 100644 index 000000000..24fa2cf30 --- /dev/null +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_registration_request.canonical.json @@ -0,0 +1 @@ +{"artifact":{"adapter":{"kind":"scalar_to_arrow_batch","version":1},"content":{"data":"ZnJvbSBfX2Z1dHVyZV9fIGltcG9ydCBhbm5vdGF0aW9ucwoKZGVmIG5vcm1hbGl6ZV9zY29yZSh2YWx1ZTogZmxvYXQpIC0+IGZsb2F0OgogICAgcmV0dXJuIHZhbHVlIC8gMTAwLjAK","encoding":"base64"},"digest":"sha256:760784bdcef57b802f389b97804cc0b618aae39e86044733451bcd0b13089a7f","entrypoint":"normalize_score","kind":"python_callable"},"name":"normalize_score","required_secrets":["API_TOKEN"],"runtime":{"env":{"MODE":"test"},"environment":{"kind":"pip","packages":["numpy>=2"]},"kind":"python","python_version":"3.12"},"signature":{"inputs":[{"arrow_type":"float64","name":"value","nullable":false}],"output":{"arrow_type":"float64","kind":"scalar","nullable":false}}} diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_registration_request.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_registration_request.json new file mode 100644 index 000000000..bbfec3169 --- /dev/null +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_registration_request.json @@ -0,0 +1,46 @@ +{ + "name": "normalize_score", + "artifact": { + "kind": "python_callable", + "digest": "sha256:760784bdcef57b802f389b97804cc0b618aae39e86044733451bcd0b13089a7f", + "entrypoint": "normalize_score", + "content": { + "encoding": "base64", + "data": "ZnJvbSBfX2Z1dHVyZV9fIGltcG9ydCBhbm5vdGF0aW9ucwoKZGVmIG5vcm1hbGl6ZV9zY29yZSh2YWx1ZTogZmxvYXQpIC0+IGZsb2F0OgogICAgcmV0dXJuIHZhbHVlIC8gMTAwLjAK" + }, + "adapter": { + "kind": "scalar_to_arrow_batch", + "version": 1 + } + }, + "signature": { + "inputs": [ + { + "name": "value", + "arrow_type": "float64", + "nullable": false + } + ], + "output": { + "kind": "scalar", + "arrow_type": "float64", + "nullable": false + } + }, + "runtime": { + "kind": "python", + "python_version": "3.12", + "environment": { + "kind": "pip", + "packages": [ + "numpy>=2" + ] + }, + "env": { + "MODE": "test" + } + }, + "required_secrets": [ + "API_TOKEN" + ] +}