diff --git a/.bumpversion.toml b/.bumpversion.toml index 3d57dc0fe..763df001b 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -1,5 +1,5 @@ [tool.bumpversion] -current_version = "0.38.0-beta.12" +current_version = "0.38.0-beta.14" parse = """(?x) (?P0|[1-9]\\d*)\\. (?P0|[1-9]\\d*)\\. diff --git a/Cargo.lock b/Cargo.lock index 6d973b073..69c33c587 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5402,7 +5402,7 @@ dependencies = [ [[package]] name = "lancedb" -version = "0.38.0-beta.12" +version = "0.38.0-beta.14" dependencies = [ "ahash", "anyhow", @@ -5490,7 +5490,7 @@ dependencies = [ [[package]] name = "lancedb-nodejs" -version = "0.38.0-beta.12" +version = "0.38.0-beta.14" dependencies = [ "arrow-array", "arrow-buffer", @@ -5515,7 +5515,7 @@ dependencies = [ [[package]] name = "lancedb-python" -version = "0.38.0-beta.12" +version = "0.38.0-beta.14" dependencies = [ "arrow", "async-trait", diff --git a/docs/src/java/java.md b/docs/src/java/java.md index e19880e29..b66a8db2f 100644 --- a/docs/src/java/java.md +++ b/docs/src/java/java.md @@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`: com.lancedb lancedb-core - 0.38.0-beta.12 + 0.38.0-beta.14 ``` diff --git a/java/lancedb-core/pom.xml b/java/lancedb-core/pom.xml index 3864ed127..6b4f2ea82 100644 --- a/java/lancedb-core/pom.xml +++ b/java/lancedb-core/pom.xml @@ -8,7 +8,7 @@ com.lancedb lancedb-parent - 0.38.0-beta.12 + 0.38.0-beta.14 ../pom.xml diff --git a/java/pom.xml b/java/pom.xml index 87e6a2bae..9a752b3ef 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -6,7 +6,7 @@ com.lancedb lancedb-parent - 0.38.0-beta.12 + 0.38.0-beta.14 pom ${project.artifactId} LanceDB Java SDK Parent POM diff --git a/nodejs/Cargo.toml b/nodejs/Cargo.toml index c3b69424f..accd9e3bf 100644 --- a/nodejs/Cargo.toml +++ b/nodejs/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "lancedb-nodejs" edition.workspace = true -version = "0.38.0-beta.12" +version = "0.38.0-beta.14" publish = false license.workspace = true description.workspace = true diff --git a/nodejs/npm/darwin-arm64/package.json b/nodejs/npm/darwin-arm64/package.json index 2b8d43c3d..8ae5be4ef 100644 --- a/nodejs/npm/darwin-arm64/package.json +++ b/nodejs/npm/darwin-arm64/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-darwin-arm64", - "version": "0.38.0-beta.12", + "version": "0.38.0-beta.14", "os": ["darwin"], "cpu": ["arm64"], "main": "lancedb.darwin-arm64.node", diff --git a/nodejs/npm/linux-arm64-gnu/package.json b/nodejs/npm/linux-arm64-gnu/package.json index 6656fdd5c..b923bf224 100644 --- a/nodejs/npm/linux-arm64-gnu/package.json +++ b/nodejs/npm/linux-arm64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-gnu", - "version": "0.38.0-beta.12", + "version": "0.38.0-beta.14", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-gnu.node", diff --git a/nodejs/npm/linux-arm64-musl/package.json b/nodejs/npm/linux-arm64-musl/package.json index f8f3e151f..2e2bd92c5 100644 --- a/nodejs/npm/linux-arm64-musl/package.json +++ b/nodejs/npm/linux-arm64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-musl", - "version": "0.38.0-beta.12", + "version": "0.38.0-beta.14", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-musl.node", diff --git a/nodejs/npm/linux-x64-gnu/package.json b/nodejs/npm/linux-x64-gnu/package.json index 20efa860a..e0b918b2f 100644 --- a/nodejs/npm/linux-x64-gnu/package.json +++ b/nodejs/npm/linux-x64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-gnu", - "version": "0.38.0-beta.12", + "version": "0.38.0-beta.14", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-gnu.node", diff --git a/nodejs/npm/linux-x64-musl/package.json b/nodejs/npm/linux-x64-musl/package.json index 4b735a687..da6de64d1 100644 --- a/nodejs/npm/linux-x64-musl/package.json +++ b/nodejs/npm/linux-x64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-musl", - "version": "0.38.0-beta.12", + "version": "0.38.0-beta.14", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-musl.node", diff --git a/nodejs/npm/win32-arm64-msvc/package.json b/nodejs/npm/win32-arm64-msvc/package.json index 35fee5ee0..5f7c7d6c2 100644 --- a/nodejs/npm/win32-arm64-msvc/package.json +++ b/nodejs/npm/win32-arm64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-arm64-msvc", - "version": "0.38.0-beta.12", + "version": "0.38.0-beta.14", "os": [ "win32" ], diff --git a/nodejs/npm/win32-x64-msvc/package.json b/nodejs/npm/win32-x64-msvc/package.json index 925211fbe..d2a808411 100644 --- a/nodejs/npm/win32-x64-msvc/package.json +++ b/nodejs/npm/win32-x64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-x64-msvc", - "version": "0.38.0-beta.12", + "version": "0.38.0-beta.14", "os": ["win32"], "cpu": ["x64"], "main": "lancedb.win32-x64-msvc.node", diff --git a/nodejs/package.json b/nodejs/package.json index a4a2286b7..38536d7f5 100644 --- a/nodejs/package.json +++ b/nodejs/package.json @@ -11,7 +11,7 @@ "ann" ], "private": false, - "version": "0.38.0-beta.12", + "version": "0.38.0-beta.14", "main": "dist/index.js", "exports": { ".": "./dist/index.js", diff --git a/python/Cargo.toml b/python/Cargo.toml index 0ee561977..972c86ede 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb-python" -version = "0.38.0-beta.12" +version = "0.38.0-beta.14" publish = false edition.workspace = true description = "Python bindings for LanceDB" diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 8a19a9d37..9be63a558 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -54,6 +54,18 @@ _UInt32 = conint(strict=True, ge=0, le=2**32 - 1) _UInt64 = conint(strict=True, ge=0, le=2**64 - 1) +def _validate_gpu_wire_marker(value: Any) -> bool: + if value is not True: + raise ValueError("runtime.gpu must be true") + return True + + +def _normalize_gpu_marker(value: bool) -> Optional[bool]: + if not isinstance(value, bool): + raise ValueError("gpu must be a boolean") + return True if value else None + + class _FrozenDict(dict): def _immutable(self, *args, **kwargs): raise TypeError("remote canonical values are immutable") @@ -239,6 +251,23 @@ class PythonRuntimeSpec(_RemoteValue): python_version: Optional[str] = None environment: Optional[PythonEnvironmentSpec] = None env: Optional[Mapping[str, str]] = None + gpu: Optional[bool] = None + + @model_validator(mode="before") + @classmethod + def _discard_unknown_runtime_payload(cls, value): + if isinstance(value, Mapping): + kind = value.get("kind") + if isinstance(kind, str) and kind not in {"python", "python_v2"}: + return {"kind": kind} + return value + + @field_validator("gpu", mode="before") + @classmethod + def _validate_gpu_marker(cls, value): + if value is None: + return None + return _validate_gpu_wire_marker(value) @model_validator(mode="after") def _validate_runtime_kind(self): @@ -247,18 +276,28 @@ class PythonRuntimeSpec(_RemoteValue): raise ValueError("python runtime requires python_version") if self.environment is None: raise ValueError("python runtime requires environment") + if self.gpu is not None: + raise ValueError("python runtime with gpu requires kind='python_v2'") + elif self.kind == "python_v2": + if self.python_version is None: + raise ValueError("python_v2 runtime requires python_version") + if self.environment is None: + raise ValueError("python_v2 runtime requires environment") + if self.gpu is None: + raise ValueError("python_v2 runtime requires gpu") else: object.__setattr__(self, "python_version", None) object.__setattr__(self, "environment", None) object.__setattr__(self, "env", None) + object.__setattr__(self, "gpu", None) return self class FunctionVersion(_RemoteValue): """An exact immutable Function version returned by Enterprise. - Scheduling resources, priority, concurrency, and retry policy belong to - the submitting Job and are not part of this identity. + The GPU execution requirement is part of this identity. CPU and memory sizing, + priority, concurrency, and retry policy belong to the execution platform. """ name: str @@ -502,31 +541,107 @@ _GRAMMAR_PRIMITIVES = ( def _canonical_arrow_type(data_type: pa.DataType) -> str: - """The server's V1 Function type grammar. Anything outside it is rejected - here rather than at registration.""" + """The compact Function grammar, or canonical exact JSON for nested types.""" + grammar = _grammar_arrow_type(data_type) + if grammar is not None: + return grammar + exact = _exact_arrow_type(data_type) + return json.dumps(exact, ensure_ascii=False, sort_keys=True, separators=(",", ":")) + + +def _grammar_arrow_type(data_type: pa.DataType) -> Optional[str]: for candidate, name in _GRAMMAR_PRIMITIVES: if data_type == candidate: return name if pa.types.is_list(data_type) or pa.types.is_large_list(data_type): + item = _grammar_list_item(data_type) + if item is None: + return None prefix = "list" if pa.types.is_list(data_type) else "large_list" - return f"{prefix}<{_canonical_list_item(data_type)}>" + return f"{prefix}<{item}>" if pa.types.is_fixed_size_list(data_type) and data_type.list_size > 0: - return ( - f"fixed_size_list<{_canonical_list_item(data_type)}, {data_type.list_size}>" - ) - raise TypeError(f"unsupported Arrow type for Function signature: {data_type}") + item = _grammar_list_item(data_type) + if item is not None: + return f"fixed_size_list<{item}, {data_type.list_size}>" + return None -def _canonical_list_item(data_type: pa.DataType) -> str: +def _grammar_list_item(data_type: pa.DataType) -> Optional[str]: """The grammar names only the item type; it always means a non-nullable - child called `item`, so any other child metadata cannot be represented.""" + child called `item`, so other child properties require exact JSON.""" child = data_type.value_field if child.name != "item" or child.nullable or child.metadata: + return None + return _grammar_arrow_type(child.type) + + +def _validate_exact_arrow_field(field: pa.Field) -> None: + if not field.name: raise TypeError( - "unsupported Arrow type for Function signature: list items must be a " - f"non-nullable field named 'item', got {child}" + "unsupported Arrow type for Function signature: field names " + "must not be empty" ) - return _canonical_arrow_type(child.type) + if field.metadata: + raise TypeError( + "unsupported Arrow type for Function signature: field metadata " + f"is not supported, got {field}" + ) + + +def _exact_arrow_field(field: pa.Field) -> dict[str, Any]: + _validate_exact_arrow_field(field) + return { + "name": field.name, + "nullable": field.nullable, + "type": _exact_arrow_type(field.type), + } + + +def _exact_arrow_type(data_type: pa.DataType) -> dict[str, Any]: + for candidate, name in _GRAMMAR_PRIMITIVES: + if data_type == candidate: + return {"type": name} + if pa.types.is_struct(data_type): + fields = list(data_type) + names = [field.name for field in fields] + if not fields or len(set(names)) != len(names): + raise TypeError( + "unsupported Arrow type for Function signature: structs must have " + "non-empty, uniquely named fields" + ) + return { + "type": "struct", + "fields": [_exact_arrow_field(field) for field in fields], + } + if ( + pa.types.is_list(data_type) + or pa.types.is_large_list(data_type) + or pa.types.is_fixed_size_list(data_type) + ): + if pa.types.is_fixed_size_list(data_type): + if data_type.value_field.name != "item": + raise TypeError( + "unsupported Arrow type for Function signature: fixed-size list " + "items must be named 'item'" + ) + if data_type.list_size <= 0: + raise TypeError( + f"unsupported Arrow type for Function signature: {data_type}" + ) + value: dict[str, Any] = { + "type": ( + "list" + if pa.types.is_list(data_type) + else "large_list" + if pa.types.is_large_list(data_type) + else "fixed_size_list" + ), + "fields": [_exact_arrow_field(data_type.value_field)], + } + if pa.types.is_fixed_size_list(data_type): + value["length"] = data_type.list_size + return value + raise TypeError(f"unsupported Arrow type for Function signature: {data_type}") def _list_of(item: pa.DataType) -> pa.DataType: @@ -600,8 +715,11 @@ def _callable_parameters(function: Callable[..., Any]) -> tuple[inspect.Paramete def _function_output(output: pa.DataType | pa.Field | pa.Schema) -> FunctionOutput: if isinstance(output, pa.Schema): + if output.metadata: + raise TypeError("Function output schema metadata is not supported") fields = tuple(output) elif isinstance(output, pa.Field) and pa.types.is_struct(output.type): + _validate_exact_arrow_field(output) if output.nullable: raise ValueError("Function output must be non-nullable") fields = tuple(output.type) @@ -617,6 +735,7 @@ def _function_output(output: pa.DataType | pa.Field | pa.Schema) -> FunctionOutp raise TypeError( "output_schema must be a PyArrow DataType, Field, or Schema" ) + _validate_exact_arrow_field(field) if field.nullable: raise ValueError("Function output must be non-nullable") return FunctionOutput( @@ -629,6 +748,8 @@ def _function_output(output: pa.DataType | pa.Field | pa.Schema) -> FunctionOutp 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") + for field in fields: + _validate_exact_arrow_field(field) names = [field.name for field in fields] if len(set(names)) != len(names): raise ValueError("Function output field names must be unique") @@ -657,6 +778,10 @@ def _infer_signature( if input_schema is not None: if not isinstance(input_schema, pa.Schema): raise TypeError("input_schema must be a PyArrow Schema") + if input_schema.metadata: + raise TypeError("Function input schema metadata is not supported") + for field in input_schema: + _validate_exact_arrow_field(field) expected = tuple(parameter.name for parameter in parameters) actual = tuple(input_schema.names) if actual != expected: @@ -910,6 +1035,7 @@ class UdfDefinition: pip: tuple[str, ...], env: Mapping[str, str], python_version: Optional[str], + gpu: bool = False, conda: tuple[str, ...] = (), conda_channels: tuple[str, ...] = (), ): @@ -938,12 +1064,14 @@ class UdfDefinition: signature = _infer_signature(function, input_schema, output_schema) source = _package_source(function) digest = f"sha256:{hashlib.sha256(source).hexdigest()}" + gpu_marker = _normalize_gpu_marker(gpu) runtime = PythonRuntimeSpec( - kind="python", + kind="python_v2" if gpu_marker is not None else "python", python_version=python_version or f"{sys.version_info.major}.{sys.version_info.minor}", environment=environment_spec, env=environment, + gpu=gpu_marker, ) self._function = function self._request = FunctionRegistrationRequest( @@ -989,6 +1117,7 @@ def udf( pip: tuple[str, ...] | list[str] = (), env: Optional[Mapping[str, str]] = None, python_version: Optional[str] = None, + gpu: bool = False, conda: tuple[str, ...] | list[str] = (), conda_channels: tuple[str, ...] | list[str] = (), ) -> Callable[[Callable[..., Any]], UdfDefinition]: ... @@ -1003,6 +1132,7 @@ def udf( pip: tuple[str, ...] | list[str] = (), env: Optional[Mapping[str, str]] = None, python_version: Optional[str] = None, + gpu: bool = False, conda: tuple[str, ...] | list[str] = (), conda_channels: tuple[str, ...] | list[str] = (), ): @@ -1035,6 +1165,10 @@ def udf( Environment variables included in the Function definition. python_version : str, optional Remote Python major/minor version. Defaults to the client version. + gpu : bool, default False + Whether every remote execution requires a GPU. The execution platform + selects one compatible GPU for each worker. The requirement is part of + the immutable Function version. The packaged artifact is a snapshot: the function source plus exactly the module-level names it references (modules as imports, importable @@ -1059,6 +1193,11 @@ def udf( ... return value * 2 >>> score(1.5) 3.0 + >>> @udf(pip=["cupy-cuda12x"], gpu=True) + ... def gpu_score(value: int) -> int: + ... return value * 2 + >>> gpu_score.registration_request.runtime.gpu + True """ def decorate(target: Callable[..., Any]) -> UdfDefinition: @@ -1070,6 +1209,7 @@ def udf( pip=tuple(pip), env={} if env is None else env, python_version=python_version, + gpu=gpu, conda=tuple(conda), conda_channels=tuple(conda_channels), ) diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 7ce6b6b91..cf1542b55 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -19,7 +19,7 @@ import pyarrow as pa import pytest import lancedb -from lancedb.functions import UdfDefinition, udf +from lancedb.functions import PythonRuntimeSpec, UdfDefinition, udf THRESHOLD = 20 _CACHE = None @@ -89,6 +89,58 @@ def test_udf_conda_environment(): udf(name="channels", conda_channels=["conda-forge"])(lambda value: value) +def test_udf_gpu_marker_uses_gpu_runtime(): + @udf(pip=["cupy-cuda12x"], gpu=True) + def double_on_gpu(value: int) -> int: + return value * 2 + + request = json.loads(double_on_gpu.registration_request.to_canonical_json()) + assert request["runtime"]["kind"] == "python_v2" + assert request["runtime"]["gpu"] is True + + @udf(pip=["pyarrow"]) + def cpu_function(value: int) -> int: + return value + + cpu_runtime = json.loads(cpu_function.registration_request.to_canonical_json())[ + "runtime" + ] + assert cpu_runtime["kind"] == "python" + assert "gpu" not in cpu_runtime + + def identity(value: int) -> int: + return value + + for invalid in [None, 0, 1, -1, 1.5, "", "true", "1", "H100"]: + with pytest.raises(ValueError, match="gpu must be a boolean"): + udf(name="invalid_gpu", gpu=invalid)(identity) + + base_runtime = { + "kind": "python_v2", + "python_version": "3.12", + "environment": {"kind": "pip"}, + } + runtime = PythonRuntimeSpec.model_validate({**base_runtime, "gpu": True}) + assert runtime.gpu is True + for invalid in [False, 1, 0, "", "true", "1", "H100"]: + with pytest.raises(ValueError, match="runtime.gpu must be true"): + PythonRuntimeSpec.model_validate({**base_runtime, "gpu": invalid}) + + +def test_unknown_runtime_discards_payload_before_known_field_validation(): + for payload in [ + {"kind": "python_v3", "gpu": {"model": "H100"}}, + {"kind": "python_v3", "resources": []}, + { + "kind": "python_v3", + "environment": {"kind": []}, + "python_version": 3.15, + }, + ]: + runtime = PythonRuntimeSpec.model_validate(payload) + assert runtime.to_canonical_json() == '{"kind":"python_v3"}' + + def test_udf_packages_attribute_access_and_body_imports(): @udf def word_norm(body: str) -> float: @@ -168,7 +220,7 @@ def test_udf_resolves_module_globals_before_builtins(tmp_path): udf(module.uses_callable_shadow) -def test_canonical_arrow_type_is_exactly_the_grammar(): +def test_canonical_arrow_type_prefers_the_compact_grammar(): from lancedb.functions import _GRAMMAR_PRIMITIVES, _canonical_arrow_type golden = json.loads( @@ -181,6 +233,13 @@ def test_canonical_arrow_type_is_exactly_the_grammar(): case["arrow_type"] for case in golden["valid"] if "<" not in case["arrow_type"] ] assert [name for _, name in _GRAMMAR_PRIMITIVES] == primitives + assert _canonical_arrow_type(pa.list_(pa.field("item", pa.float32(), False))) == ( + "list" + ) + assert ( + _canonical_arrow_type(pa.large_list(pa.field("item", pa.float32(), False))) + == "large_list" + ) for outside in [ pa.timestamp("us"), pa.decimal128(10, 2), @@ -188,7 +247,6 @@ def test_canonical_arrow_type_is_exactly_the_grammar(): pa.large_binary(), pa.binary(4), pa.duration("s"), - pa.struct([pa.field("a", pa.int32())]), pa.list_(pa.float32(), 0), pa.list_(pa.timestamp("us")), ]: @@ -378,14 +436,29 @@ def test_udf_recursion_versus_a_rebound_module_name(tmp_path): udf(raw_fact) -def test_canonical_arrow_type_rejects_unrepresentable_list_children(): +def test_canonical_arrow_type_uses_exact_json_for_list_child_properties(): from lancedb.functions import _canonical_arrow_type + nullable = pa.list_(pa.float32()) + assert json.loads(_canonical_arrow_type(nullable)) == { + "type": "list", + "fields": [ + { + "name": "item", + "nullable": True, + "type": {"type": "float32"}, + } + ], + } + named = pa.list_(pa.field("custom", pa.float32(), nullable=False)) + assert json.loads(_canonical_arrow_type(named))["fields"][0]["name"] == "custom" for outside in [ - pa.list_(pa.float32()), # pyarrow default: nullable child - pa.list_(pa.field("custom", pa.float32(), nullable=False)), pa.list_(pa.field("item", pa.float32(), nullable=False, metadata={"k": "v"})), pa.list_(pa.field("item", pa.float32(), nullable=False), 0), + pa.list_( + pa.field("item", pa.float32(), nullable=False, metadata={"k": "v"}), 3 + ), + pa.list_(pa.field("custom", pa.float32(), nullable=False), 3), ]: with pytest.raises(TypeError, match="unsupported Arrow type"): _canonical_arrow_type(outside) @@ -395,6 +468,29 @@ def test_canonical_arrow_type_rejects_unrepresentable_list_children(): ) == "fixed_size_list" ) + fixed = json.loads(_canonical_arrow_type(pa.list_(pa.float32(), 3))) + assert fixed == { + "type": "fixed_size_list", + "fields": [ + { + "name": "item", + "nullable": True, + "type": {"type": "float32"}, + } + ], + "length": 3, + } + large = json.loads(_canonical_arrow_type(pa.large_list(pa.float32()))) + assert large["type"] == "large_list" + assert large["fields"][0]["nullable"] is True + + for invalid_struct in [ + pa.struct([]), + pa.struct([pa.field("a", pa.int32()), pa.field("a", pa.int64())]), + pa.struct([pa.field("", pa.int32())]), + ]: + with pytest.raises(TypeError, match="unsupported Arrow type"): + _canonical_arrow_type(invalid_struct) def _calls_missing(value: int) -> int: @@ -482,6 +578,105 @@ def test_explicit_arrow_schema_is_deterministic(): assert signature.output.nullable is False +def test_nested_struct_output_uses_canonical_exact_json(): + token = pa.struct( + [ + pa.field("position", pa.int32(), nullable=False), + pa.field("value", pa.string(), nullable=False), + pa.field("length", pa.int32(), nullable=False), + ] + ) + analysis = pa.struct( + [ + pa.field("normalized_text", pa.string(), nullable=False), + pa.field("has_content", pa.bool_(), nullable=False), + pa.field( + "metrics", + pa.struct( + [ + pa.field("character_count", pa.int64(), nullable=False), + pa.field("word_count", pa.int32(), nullable=False), + pa.field("average_word_length", pa.float64(), nullable=False), + ] + ), + nullable=False, + ), + pa.field( + "diagnostics", + pa.struct( + [ + pa.field("status", pa.string(), nullable=False), + pa.field( + "normalization", + pa.struct( + [ + pa.field("changed", pa.bool_(), nullable=False), + pa.field( + "original_length", pa.int64(), nullable=False + ), + ] + ), + nullable=False, + ), + ] + ), + nullable=False, + ), + pa.field( + "token_preview", + pa.list_(pa.field("item", token, nullable=False)), + nullable=False, + ), + ] + ) + + @udf( + input_schema=pa.schema([pa.field("text", pa.string(), nullable=False)]), + output_schema=pa.field("analysis", analysis, nullable=False), + ) + def analyze(text): + return {"normalized_text": text} + + output = analyze.registration_request.signature.output + assert output.kind == "named_struct" + assert [field.name for field in output.fields] == [ + "normalized_text", + "has_content", + "metrics", + "diagnostics", + "token_preview", + ] + metrics = json.loads(output.fields[2].arrow_type) + assert metrics == { + "type": "struct", + "fields": [ + { + "name": "character_count", + "nullable": False, + "type": {"type": "int64"}, + }, + { + "name": "word_count", + "nullable": False, + "type": {"type": "int32"}, + }, + { + "name": "average_word_length", + "nullable": False, + "type": {"type": "float64"}, + }, + ], + } + preview = json.loads(output.fields[4].arrow_type) + assert preview["type"] == "list" + assert preview["fields"][0]["type"]["type"] == "struct" + assert [field["name"] for field in preview["fields"][0]["type"]["fields"]] == [ + "position", + "value", + "length", + ] + + def test_annotation_and_explicit_schema_validation_fail_closed(): with pytest.raises(TypeError, match="missing Function annotations"): @@ -525,6 +720,72 @@ def test_annotation_and_explicit_schema_validation_fail_closed(): def nullable_explicit(value): return value + for invalid_field in [ + pa.field("", pa.int32(), nullable=False), + pa.field("result", pa.int32(), nullable=False, metadata={"k": "v"}), + ]: + with pytest.raises(TypeError, match="unsupported Arrow type"): + + @udf( + input_schema=pa.schema([pa.field("value", pa.int64())]), + output_schema=pa.schema([invalid_field]), + ) + def invalid_explicit_field(value): + return value + + with pytest.raises(TypeError, match="unsupported Arrow type"): + + @udf( + input_schema=pa.schema( + [pa.field("value", pa.int64(), metadata={"k": "v"})] + ), + output_schema=pa.int64(), + ) + def input_field_metadata(value): + return value + + with pytest.raises(TypeError, match="unsupported Arrow type"): + + @udf( + input_schema=pa.schema([pa.field("value", pa.int64())]), + output_schema=pa.field( + "result", pa.int64(), nullable=False, metadata={"k": "v"} + ), + ) + def scalar_output_field_metadata(value): + return value + + struct_type = pa.struct([pa.field("value", pa.int64(), nullable=False)]) + with pytest.raises(TypeError, match="unsupported Arrow type"): + + @udf( + input_schema=pa.schema([pa.field("value", pa.int64())]), + output_schema=pa.field( + "result", struct_type, nullable=False, metadata={"k": "v"} + ), + ) + def struct_output_field_metadata(value): + return {"value": value} + + for input_schema, output_schema in [ + ( + pa.schema([pa.field("value", pa.int64())], metadata={"k": "v"}), + pa.int64(), + ), + ( + pa.schema([pa.field("value", pa.int64())]), + pa.schema( + [pa.field("result", pa.int64(), nullable=False)], + metadata={"k": "v"}, + ), + ), + ]: + with pytest.raises(TypeError, match="schema metadata"): + + @udf(input_schema=input_schema, output_schema=output_schema) + def schema_metadata(value): + return value + def test_local_function_catalog_operations_are_not_supported(tmp_path): db = lancedb.connect(tmp_path) diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index 881a5017e..21babc252 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb" -version = "0.38.0-beta.12" +version = "0.38.0-beta.14" edition.workspace = true description = "LanceDB: A serverless, low-latency vector database for AI applications" license.workspace = true diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index 5366d984e..4b31a4376 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -207,6 +207,33 @@ pub enum PythonRuntimeSpec { environment: PythonEnvironmentSpec, env: BTreeMap, }, + /// The GPU-enabled Sophon-managed Python runtime. + /// + /// # Examples + /// + /// ``` + /// use std::collections::BTreeMap; + /// use lancedb::function::{PythonEnvironmentSpec, PythonRuntimeSpec}; + /// + /// let runtime = PythonRuntimeSpec::PythonV2 { + /// python_version: "3.12".to_string(), + /// environment: PythonEnvironmentSpec { + /// kind: "pip".to_string(), + /// packages: vec!["cupy-cuda12x".to_string()], + /// channels: Vec::new(), + /// path: None, + /// modules: Vec::new(), + /// image: None, + /// }, + /// env: BTreeMap::new(), + /// }; + /// assert!(runtime.requires_gpu()); + /// ``` + PythonV2 { + python_version: String, + environment: PythonEnvironmentSpec, + env: BTreeMap, + }, /// A runtime kind introduced by a newer server. /// /// Unknown payload fields are intentionally not retained because the @@ -219,22 +246,27 @@ impl PythonRuntimeSpec { pub fn kind(&self) -> &str { match self { Self::Python { .. } => "python", + Self::PythonV2 { .. } => "python_v2", Self::Unrecognized { kind } => kind, } } - /// The Python version for the V1 runtime, or `None` for an unknown kind. + /// The Python version for a known Python runtime, or `None` for an unknown kind. pub fn python_version(&self) -> Option<&str> { match self { - Self::Python { python_version, .. } => Some(python_version), + Self::Python { python_version, .. } | Self::PythonV2 { python_version, .. } => { + Some(python_version) + } Self::Unrecognized { .. } => None, } } - /// The Python environment for the V1 runtime, or `None` for an unknown kind. + /// The Python environment for a known Python runtime, or `None` for an unknown kind. pub fn environment(&self) -> Option<&PythonEnvironmentSpec> { match self { - Self::Python { environment, .. } => Some(environment), + Self::Python { environment, .. } | Self::PythonV2 { environment, .. } => { + Some(environment) + } Self::Unrecognized { .. } => None, } } @@ -242,38 +274,73 @@ impl PythonRuntimeSpec { /// Environment variables, or `None` for an unknown kind. pub fn env(&self) -> Option<&BTreeMap> { match self { - Self::Python { env, .. } => Some(env), + Self::Python { env, .. } | Self::PythonV2 { env, .. } => Some(env), Self::Unrecognized { .. } => None, } } + + /// Whether the runtime requires a GPU selected by the execution platform. + pub fn requires_gpu(&self) -> bool { + matches!(self, Self::PythonV2 { .. }) + } } #[derive(Deserialize)] -struct PythonRuntimeWire { - kind: String, - #[serde(default)] - python_version: Option, - #[serde(default)] - environment: Option, +struct PythonRuntimeV1Wire { + python_version: String, + environment: PythonEnvironmentSpec, #[serde(default)] env: BTreeMap, + #[serde(default)] + gpu: Option, +} + +#[derive(Deserialize)] +struct PythonRuntimeV2Wire { + python_version: String, + environment: PythonEnvironmentSpec, + #[serde(default)] + env: BTreeMap, + gpu: bool, } impl<'de> Deserialize<'de> for PythonRuntimeSpec { fn deserialize>(deserializer: D) -> std::result::Result { - let wire = PythonRuntimeWire::deserialize(deserializer)?; - if wire.kind == "python" { - Ok(Self::Python { - python_version: wire - .python_version - .ok_or_else(|| de::Error::missing_field("python_version"))?, - environment: wire - .environment - .ok_or_else(|| de::Error::missing_field("environment"))?, - env: wire.env, - }) - } else { - Ok(Self::Unrecognized { kind: wire.kind }) + let value = Value::deserialize(deserializer)?; + let kind = value + .get("kind") + .ok_or_else(|| de::Error::missing_field("kind"))? + .as_str() + .ok_or_else(|| de::Error::custom("runtime.kind must be a string"))? + .to_string(); + match kind.as_str() { + "python" => { + let wire: PythonRuntimeV1Wire = + serde_json::from_value(value).map_err(de::Error::custom)?; + if wire.gpu.is_some() { + return Err(de::Error::custom( + "python runtime with gpu requires kind='python_v2'", + )); + } + Ok(Self::Python { + python_version: wire.python_version, + environment: wire.environment, + env: wire.env, + }) + } + "python_v2" => { + let wire: PythonRuntimeV2Wire = + serde_json::from_value(value).map_err(de::Error::custom)?; + if !wire.gpu { + return Err(de::Error::custom("runtime.gpu must be true")); + } + Ok(Self::PythonV2 { + python_version: wire.python_version, + environment: wire.environment, + env: wire.env, + }) + } + _ => Ok(Self::Unrecognized { kind }), } } } @@ -287,6 +354,8 @@ impl Serialize for PythonRuntimeSpec { environment: &'a PythonEnvironmentSpec, #[serde(skip_serializing_if = "BTreeMap::is_empty")] env: &'a BTreeMap, + #[serde(skip_serializing_if = "Option::is_none")] + gpu: Option, } #[derive(Serialize)] @@ -304,6 +373,19 @@ impl Serialize for PythonRuntimeSpec { python_version, environment, env, + gpu: None, + } + .serialize(serializer), + Self::PythonV2 { + python_version, + environment, + env, + } => PythonRuntimeRef { + kind: "python_v2", + python_version, + environment, + env, + gpu: Some(true), } .serialize(serializer), Self::Unrecognized { kind } => UnrecognizedRuntimeRef { kind }.serialize(serializer), @@ -313,8 +395,8 @@ impl Serialize for PythonRuntimeSpec { /// Immutable Function version returned by the Enterprise catalog. /// -/// Scheduling resources, priority, concurrency, and retry policy belong to -/// the submitting Job and are not part of this identity. +/// The GPU execution requirement is part of this identity. CPU and memory sizing, +/// priority, concurrency, and retry policy belong to the execution platform. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct FunctionVersion { name: String, @@ -589,7 +671,7 @@ impl_json!(RefreshColumnResult); #[cfg(test)] mod conda_environment_tests { - use super::PythonEnvironmentSpec; + use super::{PythonEnvironmentSpec, PythonRuntimeSpec}; #[test] fn conda_channels_round_trip_and_pip_stays_bare() { @@ -608,4 +690,45 @@ mod conda_environment_tests { serde_json::from_str(r#"{"kind":"pip","packages":["numpy"]}"#).unwrap(); assert!(!serde_json::to_string(&pip).unwrap().contains("channels")); } + + #[test] + fn gpu_python_runtime_marker_round_trips_and_validates() { + let runtime: PythonRuntimeSpec = serde_json::from_str( + r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"},"gpu":true}"#, + ) + .unwrap(); + assert_eq!(runtime.kind(), "python_v2"); + assert!(runtime.requires_gpu()); + assert_eq!( + super::canonical_json(&runtime).unwrap(), + r#"{"environment":{"kind":"pip"},"gpu":true,"kind":"python_v2","python_version":"3.12"}"# + ); + + for invalid in [ + r#"{"kind":"python","python_version":"3.12","environment":{"kind":"pip"},"gpu":true}"#, + r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"}}"#, + r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"},"gpu":1}"#, + r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"},"gpu":false}"#, + r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"},"gpu":"true"}"#, + r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"},"gpu":"H100"}"#, + ] { + assert!(serde_json::from_str::(invalid).is_err()); + } + } + + #[test] + fn unknown_runtime_discards_payload_before_known_field_validation() { + for encoded in [ + r#"{"kind":"python_v3","gpu":{"model":"H100"}}"#, + r#"{"kind":"python_v3","resources":[]}"#, + r#"{"kind":"python_v3","python_version":3.15,"environment":{"kind":[]}}"#, + ] { + let runtime: PythonRuntimeSpec = serde_json::from_str(encoded).unwrap(); + assert_eq!(runtime.kind(), "python_v3"); + assert_eq!( + super::canonical_json(&runtime).unwrap(), + r#"{"kind":"python_v3"}"# + ); + } + } } diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 5ce886369..d372a6f56 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -7464,6 +7464,93 @@ mod tests { assert_eq!(result.version, 8); } + #[tokio::test] + async fn test_add_function_column_allows_an_existing_binding() { + let binding = crate::function::FunctionBinding::from_json(include_str!( + "../../tests/fixtures/first_class_functions/v1/remote_function_binding.json" + )) + .unwrap(); + let binding_metadata = crate::table::computed_columns::function_bindings_metadata( + std::slice::from_ref(&binding), + ) + .unwrap(); + let mut fields = vec![ + Field::new("title", DataType::Utf8, true), + Field::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}"), + }; + Field::new(&output.output_name, data_type, true).with_metadata( + crate::table::computed_columns::function_computed_column_metadata( + binding.binding_id(), + output.output_ordinal, + &["title".into(), "body".into()], + ), + ) + })); + let schema = Schema::new_with_metadata( + fields, + HashMap::from([( + crate::table::computed_columns::FUNCTION_BINDINGS_META_KEY.to_string(), + binding_metadata, + )]), + ); + let table = + Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap(), + "/v1/table/my_table/add_columns/" => { + let actual: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) + .unwrap(); + assert_eq!( + actual["new_columns"], + serde_json::json!([ + {"name":"secondary_text","all_null":true}, + {"name":"secondary_token_count","all_null":true} + ]) + ); + http::Response::builder() + .status(200) + .body(r#"{"version":10}"#.to_string()) + .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} + ]}, + "columns":{ + "normalized_text":"secondary_text", + "token_count":"secondary_token_count" + } + }"#, + ) + .unwrap(); + + let result = table + .add_columns() + .function(application) + .execute() + .await + .unwrap(); + assert_eq!(result.version, 10); + } + #[tokio::test] async fn test_add_fixed_size_list_function_column_declares_the_vector_type() { let table = Table::new_with_handler("my_table", |request| { diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index 70fae66c6..4197f1fd2 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -547,7 +547,12 @@ fn ensure_known_binding_shape(value: &Value) -> Result<()> { Ok(()) } -fn resolve_field_path<'a>(schema: &'a ArrowSchema, path: &str) -> Result<&'a ArrowField> { +struct ResolvedFieldPath<'a> { + root: &'a ArrowField, + leaf: &'a ArrowField, +} + +fn resolve_field_path<'a>(schema: &'a ArrowSchema, path: &str) -> Result> { let parts = lance_core::datatypes::parse_field_path(path).map_err(|e| { invalid_function(format!("invalid Function input field path '{path}': {e}")) })?; @@ -556,22 +561,23 @@ fn resolve_field_path<'a>(schema: &'a ArrowSchema, path: &str) -> Result<&'a Arr "Function input field path cannot be empty", )); }; - let mut field = schema + let root = schema .field_with_name(root) .map_err(|_| invalid_function(format!("unknown Function input column '{path}'")))?; + let mut leaf = root; for child in children { - let DataType::Struct(fields) = field.data_type() else { + let DataType::Struct(fields) = leaf.data_type() else { return Err(invalid_function(format!( "Function input field path '{path}' traverses a non-struct field" ))); }; - field = fields + leaf = fields .iter() .find(|field| field.name() == child) .map(AsRef::as_ref) .ok_or_else(|| invalid_function(format!("unknown Function input column '{path}'")))?; } - Ok(field) + Ok(ResolvedFieldPath { root, leaf }) } fn canonical_input_arrow_type(field: &ArrowField) -> Result { @@ -757,7 +763,8 @@ fn function_fields_equivalent(actual: &ArrowField, expected: &ArrowField) -> boo 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)?; + let resolved = resolve_field_path(schema, &input.field_path)?; + let field = resolved.leaf; if field .metadata() .get(COMPUTED_COLUMN_META_KEY) @@ -807,6 +814,11 @@ fn ensure_binding_matches_schema(schema: &ArrowSchema, binding: &FunctionBinding ))); } + let expected_inputs = binding + .inputs() + .iter() + .map(|input| input.field_path.clone()) + .collect::>(); 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(|_| { @@ -833,6 +845,28 @@ fn ensure_binding_matches_schema(schema: &ArrowSchema, binding: &FunctionBinding binding.binding_id() ))); } + let metadata = field.metadata(); + let declared_inputs = metadata + .get(INPUTS_META_KEY) + .and_then(|raw| serde_json::from_str::>(raw).ok()); + if metadata.get(COMPUTED_COLUMN_META_KEY).map(String::as_str) != Some("true") + || metadata.get(KIND_META_KEY).map(String::as_str) != Some(FUNCTION_KIND) + || metadata + .get(FUNCTION_BINDING_ID_META_KEY) + .map(String::as_str) + != Some(binding.binding_id()) + || metadata + .get(FUNCTION_OUTPUT_ORDINAL_META_KEY) + .and_then(|value| value.parse::().ok()) + != Some(output.output_ordinal) + || declared_inputs.as_deref() != Some(expected_inputs.as_slice()) + { + return Err(invalid_function(format!( + "Function output '{}' declaration metadata does not match binding '{}'", + output.output_name, + binding.binding_id() + ))); + } output_fields.push(expected_field); } let output_schema = @@ -860,7 +894,7 @@ pub(crate) fn plan_function_application( application: &FunctionApplication, output_name: Option<&str>, ) -> Result { - ensure_no_function_bindings_for_mutation(schema, "Function binding declaration")?; + ensure_supported_function_metadata(schema)?; if application.has_unknown_fields() { return Err(Error::NotSupported { message: "Function application contains fields from a newer contract".into(), @@ -910,8 +944,9 @@ pub(crate) fn plan_function_application( input.parameter )) })?; - let field = resolve_field_path(schema, path)?; - if field + let resolved = resolve_field_path(schema, path)?; + if resolved + .root .metadata() .get(COMPUTED_COLUMN_META_KEY) .map(String::as_str) @@ -921,6 +956,7 @@ pub(crate) fn plan_function_application( "Function input '{path}' is computed; computed-on-computed bindings are not supported" ))); } + let field = resolved.leaf; let parameter_field = ArrowField::new( input.parameter.clone(), field.data_type().clone(), @@ -2687,6 +2723,37 @@ mod tests { ]) } + fn valid_function_binding_schema( + title_nullable: bool, + body_nullable: bool, + binding: &FunctionBinding, + ) -> ArrowSchema { + let mut fields = function_binding_schema(title_nullable, body_nullable) + .fields() + .iter() + .map(|field| field.as_ref().clone()) + .collect::>(); + let inputs = binding + .inputs() + .iter() + .map(|input| input.field_path.clone()) + .collect::>(); + for output in binding.outputs() { + let index = fields + .iter() + .position(|field| field.name() == &output.output_name) + .unwrap(); + fields[index] = fields[index] + .clone() + .with_metadata(function_computed_column_metadata( + binding.binding_id(), + output.output_ordinal, + &inputs, + )); + } + ArrowSchema::new(fields) + } + #[test] fn test_non_nullable_function_inputs_can_bind_to_nullable_parameters() { let binding = FunctionBinding::from_json(include_str!( @@ -2694,7 +2761,11 @@ mod tests { )) .unwrap(); - ensure_binding_matches_schema(&function_binding_schema(false, false), &binding).unwrap(); + ensure_binding_matches_schema( + &valid_function_binding_schema(false, false, &binding), + &binding, + ) + .unwrap(); } #[test] @@ -2707,8 +2778,11 @@ mod tests { raw_binding["input_schema"]["fields"][0]["nullable"] = Value::Bool(false); let binding: FunctionBinding = serde_json::from_value(raw_binding).unwrap(); - let err = ensure_binding_matches_schema(&function_binding_schema(true, false), &binding) - .unwrap_err(); + let err = ensure_binding_matches_schema( + &valid_function_binding_schema(true, false, &binding), + &binding, + ) + .unwrap_err(); assert!( matches!(&err, Error::InvalidInput { message } if message.contains("input column 'title' is nullable") @@ -2719,6 +2793,73 @@ mod tests { ); } + #[test] + fn test_second_binding_rejects_outputs_without_reciprocal_metadata() { + let binding = FunctionBinding::from_json(include_str!( + "../../tests/fixtures/first_class_functions/v1/remote_function_binding.json" + )) + .unwrap(); + let schema = ArrowSchema::new_with_metadata( + function_binding_schema(true, true).fields().to_vec(), + HashMap::from([( + FUNCTION_BINDINGS_META_KEY.to_string(), + function_bindings_metadata(std::slice::from_ref(&binding)).unwrap(), + )]), + ); + + let err = plan_function_application( + &schema, + &named_struct_application( + r#"{"normalized_text":"secondary_text","token_count":"secondary_token_count"}"#, + ), + None, + ) + .unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { message } + if message.contains("declaration metadata") + && message.contains("fb_01K3TEXT")), + "{err:?}" + ); + } + + #[test] + fn test_persisted_nested_input_keeps_leaf_level_validation() { + let mut raw_binding: Value = serde_json::from_str(include_str!( + "../../tests/fixtures/first_class_functions/v1/remote_function_binding.json" + )) + .unwrap(); + raw_binding["inputs"][0]["field_path"] = Value::String("title.value".to_string()); + let binding: FunctionBinding = serde_json::from_value(raw_binding).unwrap(); + let title = ArrowField::new( + "title", + DataType::Struct(vec![ArrowField::new("value", DataType::Utf8, true)].into()), + true, + ) + .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()), + ])); + let mut fields = vec![title, 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}"), + }; + ArrowField::new(&output.output_name, data_type, true).with_metadata( + function_computed_column_metadata( + binding.binding_id(), + output.output_ordinal, + &["title.value".into(), "body".into()], + ), + ) + })); + + ensure_binding_matches_schema(&ArrowSchema::new(fields), &binding).unwrap(); + } + #[test] fn test_function_binding_metadata_survives_schema_round_trip() { let binding = FunctionBinding::from_json(include_str!( @@ -2767,9 +2908,36 @@ mod tests { output_ordinal: 1, } if binding_id == "fb_01K3TEXT" )); - let err = plan_function_application(&reopened, &named_struct_application("{}"), None) + let dependent_application = FunctionApplication::from_json( + r#"{ + "function":{"name":"dependent","version":"fv_dependent"}, + "inputs":[ + {"parameter":"text","kind":"column","value":{"path":"search_text"}} + ], + "output":{"kind":"scalar","arrow_type":"int64","nullable":false} + }"#, + ) + .unwrap(); + let err = plan_function_application(&reopened, &dependent_application, Some("dependent")) .unwrap_err(); - assert!(matches!(err, Error::NotSupported { .. })); + assert!( + matches!(&err, Error::InvalidInput { message } if message.contains("computed-on-computed")) + ); + let plan = plan_function_application( + &reopened, + &named_struct_application( + r#"{"normalized_text":"secondary_text","token_count":"secondary_token_count"}"#, + ), + None, + ) + .unwrap(); + assert_eq!( + plan.outputs + .iter() + .map(|output| output.output_name.as_str()) + .collect::>(), + ["secondary_text", "secondary_token_count"] + ); } #[test] @@ -2925,5 +3093,38 @@ mod tests { assert!( matches!(&err, Error::InvalidInput { message } if message.contains("computed-on-computed")) ); + + let nested_title = ArrowField::new( + "title", + DataType::Struct(vec![ArrowField::new("value", DataType::Utf8, true)].into()), + true, + ) + .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(), + "struct('value')".to_string(), + ), + ])); + let nested_schema = ArrowSchema::new(vec![nested_title, schema.field(1).as_ref().clone()]); + let nested_application = FunctionApplication::from_json( + r#"{ + "function":{"name":"text_features","version":"fv_exact"}, + "inputs":[ + {"parameter":"title","kind":"column","value":{"path":"title.value"}}, + {"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} + ]} + }"#, + ) + .unwrap(); + let err = plan_function_application(&nested_schema, &nested_application, None).unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { message } if message.contains("computed-on-computed")) + ); } }