From 426684cf1b07306724261b986c7011e8b64abf9e Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 21 Aug 2026 15:48:09 +0800 Subject: [PATCH] feat: add first-class function wire contracts (#3985) ## Problem Enterprise Function-backed computed columns need a stable SDK contract before Sophon catalog and execution endpoints can be added. The existing `Job` API can only represent unit terminal results, and there is no shared Rust/Python wire definition for immutable Function versions, applications, bindings, or refresh results. ## Behavior This introduces remote-only canonical Function values in Rust and Python, evolves `Job` to decode typed remote terminal results while keeping local spawned operations unit-typed, and fixes the cross-language contract with shared JSON golden fixtures. Unknown fields and discriminator values remain forward-decodable, while canonical output contains only fields known to the client. Function models contain secret names only. Sophon remains the sole owner of catalog persistence, environment bake, secret resolution, execution, and publication. This PR does not add authoring/catalog endpoints, local execution, refresh runners, or live Sophon E2E coverage. --- docs/src/python/python.md | 41 +- python/python/lancedb/__init__.py | 6 + python/python/lancedb/functions.py | 379 ++++++++++++++ .../tests/test_first_class_function_slice1.py | 225 ++++++++ rust/lancedb/src/function.rs | 489 ++++++++++++++++++ rust/lancedb/src/job.rs | 134 ++++- rust/lancedb/src/lib.rs | 2 + rust/lancedb/src/remote/db.rs | 55 +- rust/lancedb/src/remote/job.rs | 160 ++++-- rust/lancedb/src/remote/table.rs | 6 +- .../tests/first_class_function_slice1.rs | 181 +++++++ ...remote_function_application.canonical.json | 1 + .../v1/remote_function_application.json | 20 + .../v1/remote_function_application_float.json | 8 + .../v1/remote_function_binding.canonical.json | 1 + .../v1/remote_function_binding.json | 15 + .../v1/remote_function_job.json | 31 ++ .../v1/remote_function_version.canonical.json | 1 + .../v1/remote_refresh_job.json | 15 + .../v1/remote_refresh_result.canonical.json | 1 + ...t_without_published_version.canonical.json | 1 + ...resh_result_without_published_version.json | 6 + .../v1/remote_unit_job.json | 9 + 23 files changed, 1679 insertions(+), 108 deletions(-) create mode 100644 python/python/lancedb/functions.py create mode 100644 python/python/tests/test_first_class_function_slice1.py create mode 100644 rust/lancedb/src/function.rs create mode 100644 rust/lancedb/tests/first_class_function_slice1.rs create mode 100644 rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_application.canonical.json create mode 100644 rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_application.json create mode 100644 rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_application_float.json create mode 100644 rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_binding.canonical.json create mode 100644 rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_binding.json create mode 100644 rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_job.json create mode 100644 rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json create mode 100644 rust/lancedb/tests/fixtures/first_class_functions/v1/remote_refresh_job.json create mode 100644 rust/lancedb/tests/fixtures/first_class_functions/v1/remote_refresh_result.canonical.json create mode 100644 rust/lancedb/tests/fixtures/first_class_functions/v1/remote_refresh_result_without_published_version.canonical.json create mode 100644 rust/lancedb/tests/fixtures/first_class_functions/v1/remote_refresh_result_without_published_version.json create mode 100644 rust/lancedb/tests/fixtures/first_class_functions/v1/remote_unit_job.json diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 1d5975dee..a99c0236a 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -54,6 +54,42 @@ listing a storage directory. ::: lancedb.LsmWriteSpec +## Functions and Jobs + +::: lancedb.functions.FunctionArtifact + +::: lancedb.functions.FunctionParameter + +::: lancedb.functions.FunctionResultField + +::: lancedb.functions.FunctionOutput + +::: lancedb.functions.FunctionSignature + +::: lancedb.functions.PythonEnvironmentSpec + +::: lancedb.functions.FunctionVersion + +::: lancedb.functions.PythonRuntimeSpec + +::: lancedb.functions.FunctionVersionRef + +::: lancedb.functions.ApplicationInput + +::: lancedb.functions.FunctionApplication + +::: lancedb.functions.InputBinding + +::: lancedb.functions.OutputMapping + +::: lancedb.functions.FunctionBinding + +::: lancedb.functions.RefreshColumnResult + +::: lancedb.job.Job + +::: lancedb.job.AsyncJob + ## Expressions Type-safe expression builder for filters and projections. Use these instead @@ -153,8 +189,9 @@ The same option is available on `lancedb.tokenize(...)` and the deprecated ```python import lancedb -tokens = list(lancedb.tokenize("acme makes searchable data", - custom_stop_words=["acme"])) +tokens = list( + lancedb.tokenize("acme makes searchable data", custom_stop_words=["acme"]) +) ``` ::: lancedb.tokenize diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index e12ef4e86..a8a336a6d 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -22,6 +22,12 @@ from .remote.db import RemoteDBConnection from .expr import Expr, col, lit, func from .schema import blob, vector, BlobType from .job import AsyncJob, Job +from .functions import ( + FunctionApplication as FunctionApplication, + FunctionBinding as FunctionBinding, + FunctionVersion as FunctionVersion, + PythonRuntimeSpec as PythonRuntimeSpec, +) from .table import AsyncTable, Table from .types import BaseTokenizerType from ._lancedb import Session diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py new file mode 100644 index 000000000..4ffb65e2c --- /dev/null +++ b/python/python/lancedb/functions.py @@ -0,0 +1,379 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""Canonical values exchanged with LanceDB Enterprise Function services. + +These immutable models contain client/wire state only. Catalog persistence, +environment bake, secret resolution, and execution are owned by Sophon. +""" + +from __future__ import annotations + +import json +from collections.abc import Mapping +from typing import Any, Optional + +import pydantic +from pydantic import BaseModel, Field, conint + +_PYDANTIC_V2 = int(pydantic.VERSION.split(".", 1)[0]) >= 2 +if _PYDANTIC_V2: + from pydantic import field_validator, model_validator +else: + from pydantic import root_validator, validator + +_Int32 = conint(strict=True, ge=-(2**31), le=2**31 - 1) +_UInt32 = conint(strict=True, ge=0, le=2**32 - 1) +_UInt64 = conint(strict=True, ge=0, le=2**64 - 1) + + +class _FrozenDict(dict): + def _immutable(self, *args, **kwargs): + raise TypeError("remote canonical values are immutable") + + __setitem__ = _immutable + __delitem__ = _immutable + clear = _immutable + pop = _immutable + popitem = _immutable + setdefault = _immutable + update = _immutable + + def __ior__(self, other): + self._immutable() + + +def _freeze_value(value): + if isinstance(value, Mapping): + return _FrozenDict({key: _freeze_value(child) for key, child in value.items()}) + if isinstance(value, (list, tuple)): + return tuple(_freeze_value(child) for child in value) + return value + + +def _validate_literal(value): + if isinstance(value, float): + raise ValueError( + "floating-point Function literals are not part of the Slice 1 " + "canonical wire contract" + ) + if isinstance(value, int) and not isinstance(value, bool): + if not -(2**63) <= value <= 2**64 - 1: + raise ValueError( + "Function integer literal is outside the canonical JSON range" + ) + elif isinstance(value, Mapping): + for child in value.values(): + _validate_literal(child) + elif isinstance(value, (list, tuple)): + for child in value: + _validate_literal(child) + return value + + +def _known_wire_value(value): + if isinstance(value, _RemoteValue): + return value._known_dict() + if isinstance(value, Mapping): + return {key: _known_wire_value(child) for key, child in value.items()} + if isinstance(value, (list, tuple)): + return [_known_wire_value(child) for child in value] + return value + + +class _RemoteValue(BaseModel): + if _PYDANTIC_V2: + model_config = {"extra": "ignore", "frozen": True} + else: + + class Config: + allow_mutation = False + extra = "ignore" + + if _PYDANTIC_V2: + + @model_validator(mode="after") + def _freeze_mappings(self): + for name, value in self.__dict__.items(): + object.__setattr__(self, name, _freeze_value(value)) + return self + + else: + + @root_validator + def _freeze_mappings(cls, values): + return {name: _freeze_value(value) for name, value in values.items()} + + @classmethod + def from_json(cls, payload: str): + if _PYDANTIC_V2: + return cls.model_validate_json(payload) + return cls.parse_raw(payload) + + def _known_dict(self) -> dict[str, Any]: + fields = self.__class__.model_fields if _PYDANTIC_V2 else self.__fields__ + known = {} + for name, field in fields.items(): + value = getattr(self, name) + if value is None: + continue + required = field.is_required() if _PYDANTIC_V2 else field.required + if not required: + default_factory = field.default_factory + if default_factory is not None and value == default_factory(): + continue + if default_factory is None and value == field.default: + continue + known[name] = _known_wire_value(value) + return known + + def _copy(self, *, update: Mapping[str, Any]): + update = {name: _freeze_value(value) for name, value in update.items()} + if _PYDANTIC_V2: + return self.model_copy(update=update) + return self.copy(update=update) + + def to_canonical_json(self) -> str: + return json.dumps( + self._known_dict(), + ensure_ascii=False, + allow_nan=False, + sort_keys=True, + separators=(",", ":"), + ) + + +class FunctionArtifact(_RemoteValue): + """Content-addressed Python artifact identity.""" + + kind: str + digest: str + entrypoint: str + + +class FunctionParameter(_RemoteValue): + name: str + arrow_type: str + nullable: bool + + +class FunctionResultField(_RemoteValue): + name: str + arrow_type: str + nullable: bool + + +class FunctionOutput(_RemoteValue): + """Scalar or ordered named-struct output; unknown kinds remain decodable.""" + + kind: str + arrow_type: Optional[str] = None + nullable: Optional[bool] = None + fields: tuple[FunctionResultField, ...] = () + + +class FunctionSignature(_RemoteValue): + inputs: tuple[FunctionParameter, ...] + output: FunctionOutput + + +class PythonEnvironmentSpec(_RemoteValue): + """One Sophon-managed Python environment source.""" + + kind: str + packages: tuple[str, ...] = () + path: Optional[str] = None + modules: tuple[str, ...] = () + image: Optional[str] = None + + +class PythonRuntimeSpec(_RemoteValue): + """Remote runtime definition with non-secret environment values. + + V1 supports ``kind="python"``. Newer runtime kinds remain readable, while + their unknown payload fields are intentionally not retained by the client. + """ + + kind: str + python_version: Optional[str] = None + environment: Optional[PythonEnvironmentSpec] = None + env: Optional[Mapping[str, str]] = None + + if _PYDANTIC_V2: + + @model_validator(mode="after") + def _validate_runtime_kind(self): + if self.kind == "python": + if self.python_version is None: + raise ValueError("python runtime requires python_version") + if self.environment is None: + raise ValueError("python runtime requires environment") + else: + object.__setattr__(self, "python_version", None) + object.__setattr__(self, "environment", None) + object.__setattr__(self, "env", None) + return self + + else: + + @root_validator + def _validate_runtime_kind(cls, values): + if values.get("kind") == "python": + if values.get("python_version") is None: + raise ValueError("python runtime requires python_version") + if values.get("environment") is None: + raise ValueError("python runtime requires environment") + else: + values["python_version"] = None + values["environment"] = None + values["env"] = None + return values + + +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. + """ + + name: str + version: str + artifact: FunctionArtifact + signature: FunctionSignature + runtime: PythonRuntimeSpec + runtime_digest: str + environment_digest: str + required_secrets: tuple[str, ...] = () + created_at: str + + +class FunctionVersionRef(_RemoteValue): + name: str + version: str + + +class ApplicationInput(_RemoteValue): + """One parameter value. + + Slice 1 freezes integers, strings, booleans, nulls, arrays, and objects. + Floating-point literal encoding is deferred until Python authoring is + introduced with a language-neutral numeric representation. + """ + + parameter: str + kind: str + value: Any + + if _PYDANTIC_V2: + + @field_validator("value") + @classmethod + def _validate_value(cls, value): + return _validate_literal(value) + + else: + + @validator("value") + def _validate_value(cls, value): + return _validate_literal(value) + + +class FunctionApplication(_RemoteValue): + """Immutable pre-declaration application of an exact Function version.""" + + function: FunctionVersionRef + inputs: tuple[ApplicationInput, ...] + output: FunctionOutput + group_id: str + columns: Mapping[str, str] = Field(default_factory=dict) + + def rename(self, *, columns: Mapping[str, str]) -> FunctionApplication: + """Return a copy with result-field to table-column aliases.""" + if self.output.kind != "named_struct": + raise ValueError("rename(columns=...) requires a named-struct application") + result_fields = {field.name for field in self.output.fields} + unknown = set(columns) - result_fields + if unknown: + raise ValueError(f"unknown Function result fields: {sorted(unknown)!r}") + merged = dict(self.columns) + merged.update(columns) + destinations = tuple( + merged.get(field.name, field.name) for field in self.output.fields + ) + if len(set(destinations)) != len(destinations): + raise ValueError("FunctionApplication rename destinations must be unique") + return self._copy(update={"columns": merged}) + + +class InputBinding(_RemoteValue): + parameter: str + field_id: _Int32 + field_path: str + arrow_type: str + nullable: bool + + +class OutputMapping(_RemoteValue): + """One stable result-field mapping. + + Assignment state is outside the Slice 1 client contract. During the NULL + transition Lance exposes no public cell-flag identifier to persist here. + """ + + result_field: str + output_name: str + output_field_id: _Int32 + output_ordinal: _UInt32 + arrow_type: str + nullable: bool + + +class FunctionBinding(_RemoteValue): + """Immutable grouped binding persisted by the Enterprise table service.""" + + binding_id: str + revision: _UInt64 + function: FunctionVersionRef + group_id: str + inputs: tuple[InputBinding, ...] + outputs: tuple[OutputMapping, ...] + + +class RefreshColumnResult(_RemoteValue): + """Terminal result of a remote Function-column refresh Job.""" + + rows_assigned: _UInt64 + rows_failed: _UInt64 + rows_remaining: _UInt64 + source_version: _UInt64 + published_version: Optional[_UInt64] = None + + @property + def rows_filled(self) -> int: + """Deprecated compatibility alias for :attr:`rows_assigned`.""" + return self.rows_assigned + + @property + def version(self) -> Optional[int]: + """Deprecated compatibility alias for :attr:`published_version`.""" + return self.published_version + + +__all__ = [ + "ApplicationInput", + "FunctionApplication", + "FunctionArtifact", + "FunctionBinding", + "FunctionOutput", + "FunctionParameter", + "FunctionResultField", + "FunctionSignature", + "FunctionVersion", + "FunctionVersionRef", + "InputBinding", + "OutputMapping", + "PythonEnvironmentSpec", + "PythonRuntimeSpec", + "RefreshColumnResult", +] diff --git a/python/python/tests/test_first_class_function_slice1.py b/python/python/tests/test_first_class_function_slice1.py new file mode 100644 index 000000000..9f934507f --- /dev/null +++ b/python/python/tests/test_first_class_function_slice1.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import json +from pathlib import Path + +import pytest + +import lancedb.functions as functions +from lancedb.functions import ( + FunctionApplication, + FunctionBinding, + FunctionVersion, + PythonRuntimeSpec, + RefreshColumnResult, +) + + +FIXTURES = ( + Path(__file__).parents[3] + / "rust" + / "lancedb" + / "tests" + / "fixtures" + / "first_class_functions" + / "v1" +) + + +def fixture(name: str) -> str: + return (FIXTURES / name).read_text() + + +def job_result(name: str) -> dict: + return json.loads(fixture(name))["result"] + + +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_public_function_values_are_in_api_reference(): + docs = Path(__file__).parents[3] / "docs" / "src" / "python" / "python.md" + rendered = docs.read_text() + for name in functions.__all__: + assert f"::: lancedb.functions.{name}" in rendered + + +@pytest.mark.parametrize( + ("fixture_name", "canonical_name", "model", "nested_result"), + [ + ( + "remote_function_job.json", + "remote_function_version.canonical.json", + FunctionVersion, + True, + ), + ( + "remote_function_application.json", + "remote_function_application.canonical.json", + FunctionApplication, + False, + ), + ( + "remote_function_binding.json", + "remote_function_binding.canonical.json", + FunctionBinding, + False, + ), + ( + "remote_refresh_job.json", + "remote_refresh_result.canonical.json", + RefreshColumnResult, + True, + ), + ( + "remote_refresh_result_without_published_version.json", + "remote_refresh_result_without_published_version.canonical.json", + RefreshColumnResult, + False, + ), + ], +) +def test_python_and_rust_share_remote_canonical_goldens( + fixture_name, canonical_name, model, nested_result +): + value = json.loads(fixture(fixture_name)) + if nested_result: + value = value["result"] + decoded = model.from_json(json.dumps(value)) + assert decoded.to_canonical_json() == fixture(canonical_name).strip() + + +def test_function_version_identity_is_immutable_and_exact(): + value = job_result("remote_function_job.json") + version = FunctionVersion.from_json(json.dumps(value)) + assert version.name == "embed" + assert version.version == "fv_01K3EXACT" + assert version.required_secrets == ("HF_TOKEN",) + + with pytest.raises((TypeError, ValueError)): + version.version = "fv_changed" + with pytest.raises(TypeError, match="immutable"): + version.runtime.env["TOKENIZERS_PARALLELISM"] = "true" + + changed = dict(value) + changed["version"] = "fv_changed" + assert FunctionVersion(**changed) != version + + +def test_unknown_fields_and_discriminators_are_forward_decodable(): + value = job_result("remote_function_job.json") + value["future_version_metadata"] = {"retention_class": "catalog"} + value["runtime"] = {"kind": "wasm", "module_digest": "sha256:wasm"} + value["signature"]["output"]["kind"] = "future_output_shape" + + version = FunctionVersion.from_json(json.dumps(value)) + assert version.runtime.kind == "wasm" + assert version.runtime.python_version is None + assert version.runtime.environment is None + assert json.loads(version.to_canonical_json())["runtime"] == {"kind": "wasm"} + assert version.signature.output.kind == "future_output_shape" + + +def test_function_application_uses_rename_columns_only(): + application = FunctionApplication.from_json( + fixture("remote_function_application.json") + ) + renamed = application.rename(columns={"normalized_text": "body_normalized"}) + + assert application.columns["normalized_text"] == "search_text" + assert renamed.columns["normalized_text"] == "body_normalized" + assert renamed.function == application.function + assert renamed.group_id == application.group_id + assert not hasattr(application, "rename_outputs") + with pytest.raises(TypeError, match="immutable"): + renamed.columns["normalized_text"] = "changed" + with pytest.raises(TypeError, match="immutable"): + application.inputs[0].value["path"] = "changed" + + with pytest.raises(ValueError, match="unknown Function result fields"): + application.rename(columns={"missing": "search_text"}) + with pytest.raises(ValueError, match="destinations must be unique"): + application.rename(columns={"normalized_text": "same", "token_count": "same"}) + + bare_value = json.loads(fixture("remote_function_application.json")) + bare_value.pop("columns") + bare = FunctionApplication(**bare_value) + with pytest.raises(ValueError, match="destinations must be unique"): + bare.rename(columns={"normalized_text": "token_count"}) + + +def test_binding_and_refresh_result_keep_stable_remote_fields(): + binding = FunctionBinding.from_json(fixture("remote_function_binding.json")) + assert binding.revision == 3 + assert binding.function.version == "fv_01K3TEXT" + assert [output.output_ordinal for output in binding.outputs] == [0, 1] + + result = RefreshColumnResult.from_json( + json.dumps(job_result("remote_refresh_job.json")) + ) + assert result.rows_filled == result.rows_assigned + assert result.version == result.published_version + + result = RefreshColumnResult.from_json( + fixture("remote_refresh_result_without_published_version.json") + ) + assert result.published_version is None + assert RefreshColumnResult.from_json(result.to_canonical_json()) == result + + +def test_function_literal_numeric_domain_matches_rust(): + with pytest.raises(ValueError, match="floating-point Function literals"): + FunctionApplication.from_json(fixture("remote_function_application_float.json")) + + value = json.loads(fixture("remote_function_application_float.json")) + value["inputs"][0]["value"] = 2**64 + with pytest.raises(ValueError, match="outside the canonical JSON range"): + FunctionApplication.from_json(json.dumps(value)) + + +def test_empty_default_maps_have_stable_canonical_bytes(): + runtime = PythonRuntimeSpec( + kind="python", python_version="3.12", environment={"kind": "pip"} + ) + assert runtime.to_canonical_json() == ( + '{"environment":{"kind":"pip"},"kind":"python","python_version":"3.12"}' + ) + + value = json.loads(fixture("remote_function_application.json")) + value.pop("columns") + application = FunctionApplication.from_json(json.dumps(value)) + assert "columns" not in json.loads(application.to_canonical_json()) + + +@pytest.mark.parametrize("field", ["rows_assigned", "source_version"]) +def test_refresh_result_rejects_non_u64_values(field): + value = job_result("remote_refresh_job.json") + value[field] = -1 + with pytest.raises(ValueError): + RefreshColumnResult.from_json(json.dumps(value)) + + value[field] = "1" + with pytest.raises(ValueError): + RefreshColumnResult.from_json(json.dumps(value)) + + +def test_canonical_client_values_contain_secret_names_only(): + version = FunctionVersion.from_json( + json.dumps(job_result("remote_function_job.json")) + ) + canonical = json.loads(version.to_canonical_json()) + assert canonical["required_secrets"] == ["HF_TOKEN"] + assert_no_secret_values(canonical) diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs new file mode 100644 index 000000000..087a00b90 --- /dev/null +++ b/rust/lancedb/src/function.rs @@ -0,0 +1,489 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Canonical values exchanged with the Enterprise Function service. +//! +//! This module contains client/wire values only. Catalog persistence, +//! environment bake, secret resolution, and execution are owned by Sophon. + +use std::collections::BTreeMap; + +use serde::de::{self, DeserializeOwned}; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; +use serde_json::Value; + +use crate::{Error, Result}; + +fn invalid_json(error: impl std::fmt::Display) -> Error { + Error::InvalidInput { + message: format!("invalid remote Function JSON: {error}"), + } +} + +fn write_canonical_json(value: &Value, output: &mut String) -> serde_json::Result<()> { + match value { + Value::Object(map) => { + output.push('{'); + let mut entries = map.iter().collect::>(); + entries.sort_unstable_by_key(|(key, _)| *key); + for (index, (key, value)) in entries.into_iter().enumerate() { + if index != 0 { + output.push(','); + } + output.push_str(&serde_json::to_string(key)?); + output.push(':'); + write_canonical_json(value, output)?; + } + output.push('}'); + } + Value::Array(values) => { + output.push('['); + for (index, value) in values.iter().enumerate() { + if index != 0 { + output.push(','); + } + write_canonical_json(value, output)?; + } + output.push(']'); + } + other => output.push_str(&serde_json::to_string(other)?), + } + Ok(()) +} + +fn canonical_json(value: &T) -> Result { + let value = serde_json::to_value(value).map_err(invalid_json)?; + let mut output = String::new(); + write_canonical_json(&value, &mut output).map_err(invalid_json)?; + Ok(output) +} + +fn from_json(json: &str) -> Result { + serde_json::from_str(json).map_err(invalid_json) +} + +fn validate_literal(value: &Value) -> Result<()> { + match value { + Value::Number(number) if number.is_f64() => Err(Error::InvalidInput { + message: "floating-point Function literals are not part of the Slice 1 canonical wire contract" + .to_string(), + }), + Value::Array(values) => values.iter().try_for_each(validate_literal), + Value::Object(values) => values.values().try_for_each(validate_literal), + _ => Ok(()), + } +} + +macro_rules! impl_json { + ($type:ty) => { + impl $type { + /// Decode a remote value. Unknown fields and discriminator values + /// are accepted so newer servers remain readable. + pub fn from_json(json: &str) -> Result { + from_json(json) + } + + /// Encode the known client contract with bytewise-sorted JSON keys. + pub fn to_canonical_json(&self) -> Result { + canonical_json(self) + } + } + }; +} + +/// Packaged Python artifact identity. Source bytes are never part of this value. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct FunctionArtifact { + pub kind: String, + pub digest: String, + pub entrypoint: String, +} + +/// One ordered Arrow input parameter. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct FunctionParameter { + pub name: String, + pub arrow_type: String, + pub nullable: bool, +} + +/// One field of an ordered named-struct result. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct FunctionResultField { + pub name: String, + pub arrow_type: String, + pub nullable: bool, +} + +/// Scalar or named-struct Function output. +/// +/// `kind` remains a string so unknown future result shapes can be decoded. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct FunctionOutput { + pub kind: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub arrow_type: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub nullable: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub fields: Vec, +} + +/// Ordered language-neutral Function signature. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct FunctionSignature { + pub inputs: Vec, + pub output: FunctionOutput, +} + +/// One Python environment source. +/// +/// The selected source is interpreted by Sophon. `kind` is open for forward +/// compatible decoding. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct PythonEnvironmentSpec { + pub kind: String, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub packages: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub path: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub modules: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub image: Option, +} + +/// Reproducible Python runtime definition understood by Sophon. +/// +/// `env` contains non-secret values. Secret values have no client model; +/// [`FunctionVersion::required_secrets`] contains names only. +#[derive(Debug, Clone, PartialEq, Eq)] +#[non_exhaustive] +pub enum PythonRuntimeSpec { + /// The V1 Sophon-managed Python runtime. + Python { + python_version: String, + environment: PythonEnvironmentSpec, + env: BTreeMap, + }, + /// A runtime kind introduced by a newer server. + /// + /// Unknown payload fields are intentionally not retained because the + /// client does not proxy catalog values. + Unrecognized { kind: String }, +} + +impl PythonRuntimeSpec { + /// The wire discriminator reported by Sophon. + pub fn kind(&self) -> &str { + match self { + Self::Python { .. } => "python", + Self::Unrecognized { kind } => kind, + } + } + + /// The Python version for the V1 runtime, or `None` for an unknown kind. + pub fn python_version(&self) -> Option<&str> { + match self { + Self::Python { python_version, .. } => Some(python_version), + Self::Unrecognized { .. } => None, + } + } + + /// The Python environment for the V1 runtime, or `None` for an unknown kind. + pub fn environment(&self) -> Option<&PythonEnvironmentSpec> { + match self { + Self::Python { environment, .. } => Some(environment), + Self::Unrecognized { .. } => None, + } + } + + /// Non-secret environment variables, or `None` for an unknown kind. + pub fn env(&self) -> Option<&BTreeMap> { + match self { + Self::Python { env, .. } => Some(env), + Self::Unrecognized { .. } => None, + } + } +} + +#[derive(Deserialize)] +struct PythonRuntimeWire { + kind: String, + #[serde(default)] + python_version: Option, + #[serde(default)] + environment: Option, + #[serde(default)] + env: BTreeMap, +} + +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 }) + } + } +} + +impl Serialize for PythonRuntimeSpec { + fn serialize(&self, serializer: S) -> std::result::Result { + #[derive(Serialize)] + struct PythonRuntimeRef<'a> { + kind: &'static str, + python_version: &'a str, + environment: &'a PythonEnvironmentSpec, + #[serde(skip_serializing_if = "BTreeMap::is_empty")] + env: &'a BTreeMap, + } + + #[derive(Serialize)] + struct UnrecognizedRuntimeRef<'a> { + kind: &'a str, + } + + match self { + Self::Python { + python_version, + environment, + env, + } => PythonRuntimeRef { + kind: "python", + python_version, + environment, + env, + } + .serialize(serializer), + Self::Unrecognized { kind } => UnrecognizedRuntimeRef { kind }.serialize(serializer), + } + } +} + +/// 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. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct FunctionVersion { + name: String, + version: String, + artifact: FunctionArtifact, + signature: FunctionSignature, + runtime: PythonRuntimeSpec, + runtime_digest: String, + environment_digest: String, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + required_secrets: Vec, + created_at: String, +} + +impl FunctionVersion { + pub fn name(&self) -> &str { + &self.name + } + + pub fn version(&self) -> &str { + &self.version + } + + pub fn artifact(&self) -> &FunctionArtifact { + &self.artifact + } + + pub fn signature(&self) -> &FunctionSignature { + &self.signature + } + + pub fn runtime(&self) -> &PythonRuntimeSpec { + &self.runtime + } + + pub fn runtime_digest(&self) -> &str { + &self.runtime_digest + } + + pub fn environment_digest(&self) -> &str { + &self.environment_digest + } + + /// Required secret names. Resolved values exist only inside Sophon. + pub fn required_secrets(&self) -> &[String] { + &self.required_secrets + } + + pub fn created_at(&self) -> &str { + &self.created_at + } +} + +impl_json!(FunctionVersion); + +/// Exact FunctionVersion reference embedded in applications and bindings. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct FunctionVersionRef { + pub name: String, + pub version: String, +} + +/// Parameter binding in a FunctionApplication. +/// +/// `kind` remains open until Python authoring is added in Slice 2. Slice 1 +/// freezes JSON integers, strings, booleans, nulls, arrays, and objects as +/// canonical literal values. Floating-point literals are rejected until a +/// language-neutral numeric representation is defined. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ApplicationInput { + pub parameter: String, + pub kind: String, + pub value: Value, +} + +/// Pre-declaration application of an exact FunctionVersion. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct FunctionApplication { + function: FunctionVersionRef, + inputs: Vec, + output: FunctionOutput, + group_id: String, + #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] + columns: BTreeMap, +} + +impl FunctionApplication { + pub fn function(&self) -> &FunctionVersionRef { + &self.function + } + + pub fn inputs(&self) -> &[ApplicationInput] { + &self.inputs + } + + pub fn output(&self) -> &FunctionOutput { + &self.output + } + + pub fn group_id(&self) -> &str { + &self.group_id + } + + pub fn columns(&self) -> &BTreeMap { + &self.columns + } + + /// Decode a remote application after validating the Slice 1 literal domain. + pub fn from_json(json: &str) -> Result { + let application: Self = from_json(json)?; + application + .inputs + .iter() + .try_for_each(|input| validate_literal(&input.value))?; + Ok(application) + } + + /// Encode the application with bytewise-sorted JSON keys. + pub fn to_canonical_json(&self) -> Result { + self.inputs + .iter() + .try_for_each(|input| validate_literal(&input.value))?; + canonical_json(self) + } +} + +/// Stable table input bound to a registered parameter. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct InputBinding { + pub parameter: String, + pub field_id: i32, + pub field_path: String, + pub arrow_type: String, + pub nullable: bool, +} + +/// Ordered result-field to table-field mapping for a grouped binding. +/// +/// Assignment state is not part of the Slice 1 client contract. During the +/// NULL transition there is no public Lance cell-flag identifier to persist. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct OutputMapping { + pub result_field: String, + pub output_name: String, + pub output_field_id: i32, + pub output_ordinal: u32, + pub arrow_type: String, + pub nullable: bool, +} + +/// Immutable grouped binding persisted by the Enterprise table service. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct FunctionBinding { + binding_id: String, + revision: u64, + function: FunctionVersionRef, + group_id: String, + inputs: Vec, + outputs: Vec, +} + +impl FunctionBinding { + pub fn binding_id(&self) -> &str { + &self.binding_id + } + + pub fn revision(&self) -> u64 { + self.revision + } + + pub fn function(&self) -> &FunctionVersionRef { + &self.function + } + + pub fn group_id(&self) -> &str { + &self.group_id + } + + pub fn inputs(&self) -> &[InputBinding] { + &self.inputs + } + + pub fn outputs(&self) -> &[OutputMapping] { + &self.outputs + } +} + +impl_json!(FunctionBinding); + +/// Stable terminal result of a remote Function-column refresh Job. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct RefreshColumnResult { + pub rows_assigned: u64, + pub rows_failed: u64, + pub rows_remaining: u64, + pub source_version: u64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub published_version: Option, +} + +impl RefreshColumnResult { + /// Deprecated compatibility alias for `rows_assigned`. + pub fn rows_filled(&self) -> u64 { + self.rows_assigned + } + + /// Deprecated compatibility alias for `published_version`. + pub fn version(&self) -> Option { + self.published_version + } +} + +impl_json!(RefreshColumnResult); diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index d77dd6974..0f880e398 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -6,6 +6,8 @@ use std::sync::Arc; use async_trait::async_trait; +use serde::de::DeserializeOwned; +use serde_json::Value; use tokio::sync::watch; use tokio::task::{AbortHandle, JoinHandle}; @@ -19,43 +21,127 @@ pub(crate) trait JobHandle: Send + Sync { None } async fn status(&self) -> Result; - async fn wait(&self) -> Result<()>; + async fn wait(&self) -> Result; async fn cancel(&self) -> Result<()>; } +/// A backend-neutral successful terminal result. +/// +/// Local operations do not carry a value. Remote operations may carry JSON +/// that the public [`Job`] decodes according to its result type. +pub(crate) struct TerminalResult { + #[allow(dead_code)] // Typed remote submit endpoints consume this after Slice 1. + value: Option, + #[allow(dead_code)] // Preserved so typed decode errors retain request correlation. + request_id: Option, +} + +impl TerminalResult { + pub(crate) fn local() -> Self { + Self { + value: None, + request_id: None, + } + } + + pub(crate) fn remote(value: Option, request_id: String) -> Self { + Self { + value, + request_id: Some(request_id), + } + } + + #[allow(dead_code)] // Exercised by the remote typed-result fixtures in Slice 1. + fn decode(self) -> Result { + let request_id = self.request_id.unwrap_or_default(); + let value = self.value.ok_or_else(|| Error::Http { + source: "successful typed job response did not contain a result".into(), + request_id: request_id.clone(), + status_code: None, + })?; + serde_json::from_value(value).map_err(|error| Error::Http { + source: format!("failed to parse typed job result: {error}").into(), + request_id, + status_code: None, + }) + } +} + +type ResultDecoder = fn(TerminalResult) -> Result; + +enum JobInner { + Handle { + handle: Box, + decode: ResultDecoder, + }, + Completed(T), +} + /// A handle to an operation that may still be running. /// /// The operation may already be complete when the handle is created. -pub struct Job { - handle: Option>, +pub struct Job +where + T: Clone + Send + Sync + 'static, +{ + inner: JobInner, } -impl std::fmt::Debug for Job { +impl std::fmt::Debug for Job +where + T: Clone + Send + Sync + 'static, +{ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("Job") .field("id", &self.id()) - .field("done", &self.handle.is_none()) + .field("done", &matches!(self.inner, JobInner::Completed(_))) .finish() } } -impl Job { +impl Job<()> { /// A job whose operation finished before the handle was created. pub(crate) fn new_done() -> Self { - Self { handle: None } + Self { + inner: JobInner::Completed(()), + } } pub(crate) fn new(handle: Box) -> Self { Self { - handle: Some(handle), + inner: JobInner::Handle { + handle, + decode: |_| Ok(()), + }, } } - /// A job running as a task in this process. + /// A unit-result job running as a task in this process. pub(crate) fn spawned(task: JoinHandle>) -> Self { Self::new(Box::new(SpawnedJob::new(task))) } +} +impl Job +where + T: Clone + DeserializeOwned + Send + Sync + 'static, +{ + /// Construct a typed remote Job before result-specific submit APIs are added. + #[allow(dead_code)] + pub(crate) fn new_typed(handle: Box) -> Self { + Self { + inner: JobInner::Handle { + handle, + decode: TerminalResult::decode::, + }, + } + } +} + +impl Job +where + T: Clone + Send + Sync + 'static, +{ /// Identifies the operation on the server that is running it. /// /// Returned for correlating with server logs or the jobs API. Operations @@ -63,7 +149,10 @@ impl Job { /// value is opaque: parsing it or storing it to resume the job later is /// not supported. pub fn id(&self) -> Option<&str> { - self.handle.as_ref().and_then(|handle| handle.id()) + match &self.inner { + JobInner::Handle { handle, .. } => handle.id(), + JobInner::Completed(_) => None, + } } /// The operation's current lifecycle state: "running", "finished", @@ -73,9 +162,9 @@ impl Job { /// terminal failure state, or retry. States a newer server reports that /// this client version does not know pass through as-is. pub async fn status(&self) -> Result { - match &self.handle { - None => Ok("finished".to_string()), - Some(handle) => handle.status().await, + match &self.inner { + JobInner::Handle { handle, .. } => handle.status().await, + JobInner::Completed(_) => Ok("finished".to_string()), } } @@ -83,10 +172,10 @@ impl Job { /// /// Returns [`crate::Error::JobFailed`] if the operation failed and /// [`crate::Error::JobCancelled`] if it was cancelled. - pub async fn wait(&self) -> Result<()> { - match &self.handle { - None => Ok(()), - Some(handle) => handle.wait().await, + pub async fn wait(&self) -> Result { + match &self.inner { + JobInner::Handle { handle, decode } => decode(handle.wait().await?), + JobInner::Completed(result) => Ok(result.clone()), } } @@ -94,9 +183,9 @@ impl Job { /// /// Cancelling an operation that already finished is a no-op. pub async fn cancel(&self) -> Result<()> { - match &self.handle { - None => Ok(()), - Some(handle) => handle.cancel().await, + match &self.inner { + JobInner::Handle { handle, .. } => handle.cancel().await, + JobInner::Completed(_) => Ok(()), } } } @@ -162,7 +251,7 @@ impl JobHandle for SpawnedJob { Ok(label.to_string()) } - async fn wait(&self) -> Result<()> { + async fn wait(&self) -> Result { let mut outcome = self.outcome.clone(); let settled = outcome .wait_for(|outcome| outcome.is_some()) @@ -172,7 +261,8 @@ impl JobHandle for SpawnedJob { })? .clone() .expect("wait_for returns once an outcome is set"); - settled.into_result() + settled.into_result()?; + Ok(TerminalResult::local()) } async fn cancel(&self) -> Result<()> { diff --git a/rust/lancedb/src/lib.rs b/rust/lancedb/src/lib.rs index 70d023ccc..291dcaf65 100644 --- a/rust/lancedb/src/lib.rs +++ b/rust/lancedb/src/lib.rs @@ -181,6 +181,7 @@ pub mod dataloader; pub mod embeddings; pub mod error; pub mod expr; +pub mod function; pub mod index; pub mod io; pub mod ipc; @@ -205,6 +206,7 @@ use serde::{Deserialize, Serialize}; pub use blob::{BlobRangeRequest, blob, is_blob}; pub use connection::{ConnectNamespaceBuilder, Connection}; pub use error::{Error, JobFailure, Result}; +pub use function::FunctionVersion; pub use job::Job; use lance_index::vector::ApproxMode as LanceApproxMode; use lance_linalg::distance::DistanceType as LanceDistanceType; diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 45a0bd925..03a13cb4e 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -25,7 +25,7 @@ use crate::database::{ }; use crate::error::Result; use crate::job::Job; -use crate::remote::job::RemoteJob; +use crate::remote::job::{DescribeJobResponse, RemoteJob, job_state_to_client}; use crate::remote::util::stream_as_body; use crate::table::BaseTable; @@ -472,48 +472,6 @@ struct RemoteListJobsResponse { page_token: Option, } -/// The server's account of why a job failed. Absent from older servers, -/// which report only the terminal state. -#[derive(serde::Deserialize)] -struct RemoteReportedFailure { - #[serde(default)] - phase: Option, - #[serde(default)] - message: Option, - #[serde(default)] - retryable: Option, -} - -#[derive(serde::Deserialize)] -struct RemoteDescribeJobResponse { - job_id: String, - #[serde(default)] - job_type: String, - job_state: String, - #[serde(default)] - creation_ms: i64, - #[serde(default)] - spec: serde_json::Value, - #[serde(default)] - failure: Option, -} - -/// Server job states -> the client vocabulary ("running" / "finished" / -/// "failed" / "cancelled"). Covers both the describe enum (IN_PROGRESS / -/// DONE / FAILED / CANCELLED) and the registry's lowercase list-row states -/// (in_progress / succeeded / failed / canceled / timed_out). States this -/// client version does not know (e.g. created, queued) pass through as-is. -fn job_state_to_client(state: &str) -> String { - match state { - "IN_PROGRESS" | "in_progress" => "running", - "DONE" | "done" | "succeeded" => "finished", - "FAILED" | "failed" | "TIMED_OUT" | "timed_out" => "failed", - "CANCELLED" | "cancelled" | "canceled" => "cancelled", - other => other, - } - .to_string() -} - /// Bound on `list_jobs` page walking; a warning is logged when the listing /// is truncated at this many pages. const MAX_LIST_JOBS_PAGES: usize = 100; @@ -586,19 +544,14 @@ impl Database for RemoteDatabase { }) => return Ok(None), Err(err) => return Err(err), }; - let body: RemoteDescribeJobResponse = rsp.json().await.err_to_http(request_id)?; + let body: DescribeJobResponse = rsp.json().await.err_to_http(request_id)?; Ok(Some(JobDescription { job_id: body.job_id, job_type: body.job_type, state: job_state_to_client(&body.job_state), creation_ms: body.creation_ms, spec: body.spec, - failure: body.failure.map(|reported| crate::error::JobFailure { - phase: reported.phase, - message: reported.message, - retryable: reported.retryable, - source: None, - }), + failure: body.failure.map(|reported| reported.into_job_failure()), })) } @@ -2507,7 +2460,7 @@ mod tests { http::Response::builder() .status(200) .body(format!( - r#"{{"job_id": "job-1", "job_type": "create_index", "job_state": "{}", "creation_ms": 1}}"#, + r#"{{"job_id": "job-1", "job_type": "create_function", "job_state": "{}", "creation_ms": 1, "result": {{"name": "embed", "version": "fv_1"}}}}"#, state )) .unwrap() diff --git a/rust/lancedb/src/remote/job.rs b/rust/lancedb/src/remote/job.rs index 2fc99da59..0d41dbb35 100644 --- a/rust/lancedb/src/remote/job.rs +++ b/rust/lancedb/src/remote/job.rs @@ -8,10 +8,10 @@ use std::time::Duration; use async_trait::async_trait; use tokio::time::sleep; -use serde::{Deserialize, Deserializer}; +use serde::Deserialize; use crate::error::{Error, JobFailure, Result}; -use crate::job::JobHandle; +use crate::job::{JobHandle, TerminalResult}; use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient}; /// Delay before the second job-state poll; doubles up to [`MAX_POLL_INTERVAL`]. @@ -29,12 +29,6 @@ enum JobState { Other(String), } -impl<'de> Deserialize<'de> for JobState { - fn deserialize>(deserializer: D) -> std::result::Result { - Ok(Self::from(String::deserialize(deserializer)?.as_str())) - } -} - impl JobState { /// The client vocabulary label for this state. fn client_label(&self) -> String { @@ -51,22 +45,26 @@ impl JobState { impl From<&str> for JobState { fn from(state: &str) -> Self { match state { - "IN_PROGRESS" => Self::InProgress, - "CANCELLED" => Self::Cancelled, + "IN_PROGRESS" | "in_progress" => Self::InProgress, + "CANCELLED" | "cancelled" | "canceled" => Self::Cancelled, // The server reports a timed-out job as FAILED on describe; // accept the raw registry state too in case a future server // stops folding it. - "FAILED" | "TIMED_OUT" => Self::Failed, - "DONE" => Self::Done, + "FAILED" | "failed" | "TIMED_OUT" | "timed_out" => Self::Failed, + "DONE" | "done" | "succeeded" => Self::Done, other => Self::Other(other.to_string()), } } } +pub(super) fn job_state_to_client(state: &str) -> String { + JobState::from(state).client_label() +} + /// The server's account of why a job failed. Absent from older servers, which /// report only the terminal state. #[derive(Deserialize)] -struct ReportedFailure { +pub(super) struct ReportedFailure { #[serde(default)] phase: Option, #[serde(default)] @@ -75,11 +73,43 @@ struct ReportedFailure { retryable: Option, } +/// Forward-compatible `/v1/jobs/describe` wire envelope. #[derive(Deserialize)] -struct DescribeJobResponse { - job_state: JobState, +pub(super) struct DescribeJobResponse { #[serde(default)] - failure: Option, + pub(super) job_id: String, + #[serde(default)] + pub(super) job_type: String, + pub(super) job_state: String, + #[serde(default)] + pub(super) creation_ms: i64, + #[serde(default)] + pub(super) spec: serde_json::Value, + #[serde(default)] + result: Option, + #[serde(default)] + pub(super) failure: Option, +} + +impl ReportedFailure { + pub(super) fn into_job_failure(self) -> JobFailure { + JobFailure { + phase: self.phase, + message: self.message, + retryable: self.retryable, + source: None, + } + } +} + +impl DescribeJobResponse { + fn state(&self) -> JobState { + JobState::from(self.job_state.as_str()) + } + + fn into_terminal_result(self, request_id: String) -> TerminalResult { + TerminalResult::remote(self.result, request_id) + } } pub struct RemoteJob { @@ -93,7 +123,7 @@ impl RemoteJob { } /// One `/v1/jobs/describe` round trip. - async fn describe(&self) -> Result { + async fn describe(&self) -> Result<(String, DescribeJobResponse)> { let request = self .client .post("/v1/jobs/describe") @@ -104,10 +134,10 @@ impl RemoteJob { let description: DescribeJobResponse = serde_json::from_str(&body).map_err(|e| Error::Http { source: format!("failed to parse job description: {}", e).into(), - request_id, + request_id: request_id.clone(), status_code: None, })?; - Ok(description) + Ok((request_id, description)) } } @@ -118,26 +148,21 @@ impl JobHandle for RemoteJob { } async fn status(&self) -> Result { - Ok(self.describe().await?.job_state.client_label()) + Ok(self.describe().await?.1.state().client_label()) } - async fn wait(&self) -> Result<()> { + async fn wait(&self) -> Result { let mut interval = INITIAL_POLL_INTERVAL; loop { - let description = self.describe().await?; - match description.job_state { - JobState::Done => return Ok(()), + let (request_id, description) = self.describe().await?; + match description.state() { + JobState::Done => return Ok(description.into_terminal_result(request_id)), JobState::Failed => { return Err(Error::JobFailed { job_id: Some(self.job_id.clone()), failure: description .failure - .map(|reported| JobFailure { - phase: reported.phase, - message: reported.message, - retryable: reported.retryable, - source: None, - }) + .map(ReportedFailure::into_job_failure) .unwrap_or_default(), }); } @@ -168,3 +193,78 @@ impl JobHandle for RemoteJob { .map(|_| ()) } } + +#[cfg(test)] +mod tests { + use async_trait::async_trait; + + use crate::Result; + use crate::function::{FunctionVersion, RefreshColumnResult}; + use crate::job::{Job, JobHandle, TerminalResult}; + + use super::DescribeJobResponse; + + const FUNCTION_JOB: &str = + include_str!("../../tests/fixtures/first_class_functions/v1/remote_function_job.json"); + const REFRESH_JOB: &str = + include_str!("../../tests/fixtures/first_class_functions/v1/remote_refresh_job.json"); + const UNIT_JOB: &str = + include_str!("../../tests/fixtures/first_class_functions/v1/remote_unit_job.json"); + const MISSING_RESULT_JOB: &str = r#"{"job_state":"DONE"}"#; + + struct FixtureRemoteJob(&'static str); + + #[async_trait] + impl JobHandle for FixtureRemoteJob { + async fn status(&self) -> Result { + Ok("finished".to_string()) + } + + async fn wait(&self) -> Result { + let description: DescribeJobResponse = + serde_json::from_str(self.0).expect("remote job fixture"); + Ok(description.into_terminal_result("fixture-request".to_string())) + } + + async fn cancel(&self) -> Result<()> { + Ok(()) + } + } + + #[tokio::test] + async fn typed_remote_job_fixtures_decode_terminal_results() { + let function = Job::::new_typed(Box::new(FixtureRemoteJob(FUNCTION_JOB))); + let result = function.wait().await.expect("typed FunctionVersion result"); + assert_eq!(result.version(), "fv_01K3EXACT"); + + let refresh = + Job::::new_typed(Box::new(FixtureRemoteJob(REFRESH_JOB))); + let result = refresh.wait().await.expect("typed RefreshColumnResult"); + assert_eq!(result.rows_assigned, 999_998_800); + assert_eq!(result.rows_filled(), result.rows_assigned); + + let unit = Job::new(Box::new(FixtureRemoteJob(UNIT_JOB))); + unit.wait() + .await + .expect("unit result ignores additive remote payloads"); + } + + #[tokio::test] + async fn typed_remote_job_requires_a_terminal_result() { + let typed = + Job::::new_typed(Box::new(FixtureRemoteJob(MISSING_RESULT_JOB))); + let error = typed.wait().await.unwrap_err(); + assert!( + error + .to_string() + .contains("successful typed job response did not contain a result") + ); + } + + #[test] + fn remote_wire_unknown_fields_are_forward_decodable() { + let response: DescribeJobResponse = + serde_json::from_str(FUNCTION_JOB).expect("function job fixture"); + assert_eq!(response.job_state, "DONE"); + } +} diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index a0a4cebc2..328f2a708 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -162,13 +162,13 @@ impl crate::job::JobHandle for FreshnessJob { crate::job::JobHandle::status(&self.inner).await } - async fn wait(&self) -> Result<()> { - crate::job::JobHandle::wait(&self.inner).await?; + async fn wait(&self) -> Result { + let result = crate::job::JobHandle::wait(&self.inner).await?; let version = self.version.read().await; if version.is_none() { self.freshness.lock().unwrap().checkout_baseline = Some(SystemTime::now()); } - Ok(()) + Ok(result) } async fn cancel(&self) -> Result<()> { diff --git a/rust/lancedb/tests/first_class_function_slice1.rs b/rust/lancedb/tests/first_class_function_slice1.rs new file mode 100644 index 000000000..dab05fe48 --- /dev/null +++ b/rust/lancedb/tests/first_class_function_slice1.rs @@ -0,0 +1,181 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use std::fs; +use std::path::PathBuf; + +use lancedb::function::{ + FunctionApplication, FunctionBinding, FunctionVersion, RefreshColumnResult, +}; +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 job_result(name: &str) -> Value { + serde_json::from_str::(&fixture(name)).expect("remote Job fixture")["result"].clone() +} + +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" + ), + "client canonical value must not model resolved secret material" + ); + assert_no_secret_values(value); + } + } + Value::Array(values) => values.iter().for_each(assert_no_secret_values), + _ => {} + } +} + +#[test] +fn function_version_job_result_matches_shared_canonical_golden() { + let result = job_result("remote_function_job.json"); + let version = FunctionVersion::from_json(&result.to_string()).expect("FunctionVersion result"); + + assert_eq!(version.name(), "embed"); + assert_eq!(version.version(), "fv_01K3EXACT"); + assert_eq!(version.runtime_digest(), "sha256:runtime"); + assert_eq!(version.required_secrets(), &["HF_TOKEN"]); + assert_eq!( + version.to_canonical_json().expect("canonical JSON"), + fixture("remote_function_version.canonical.json").trim() + ); +} + +#[test] +fn version_identity_is_immutable_and_exact() { + let original = job_result("remote_function_job.json"); + let version = + FunctionVersion::from_json(&original.to_string()).expect("FunctionVersion result"); + let reopened = version.clone(); + assert_eq!(reopened, version); + assert_eq!(reopened.name(), version.name()); + assert_eq!(reopened.version(), version.version()); + + let mut changed = original; + changed["version"] = Value::String("fv_01K3DIFFERENT".to_string()); + let changed = FunctionVersion::from_json(&changed.to_string()).expect("changed version"); + assert_ne!(changed, version); +} + +#[test] +fn application_and_binding_match_shared_remote_goldens() { + let application = FunctionApplication::from_json(&fixture("remote_function_application.json")) + .expect("application fixture"); + assert_eq!(application.function().version, "fv_01K3TEXT"); + assert_eq!(application.output().kind, "named_struct"); + assert_eq!(application.inputs().len(), 2); + assert_eq!( + application.to_canonical_json().expect("canonical JSON"), + fixture("remote_function_application.canonical.json").trim() + ); + + let binding = FunctionBinding::from_json(&fixture("remote_function_binding.json")) + .expect("binding fixture"); + assert_eq!(binding.revision(), 3); + assert_eq!(binding.function().version, "fv_01K3TEXT"); + assert_eq!(binding.outputs()[0].output_ordinal, 0); + assert_eq!(binding.outputs()[1].output_ordinal, 1); + assert_eq!( + binding.to_canonical_json().expect("canonical JSON"), + fixture("remote_function_binding.canonical.json").trim() + ); +} + +#[test] +fn refresh_job_result_matches_shared_canonical_golden() { + let result = job_result("remote_refresh_job.json"); + let result = RefreshColumnResult::from_json(&result.to_string()).expect("refresh result"); + assert_eq!(result.rows_assigned, 999_998_800); + assert_eq!(result.rows_filled(), result.rows_assigned); + assert_eq!(result.version(), result.published_version); + assert_eq!( + result.to_canonical_json().expect("canonical JSON"), + fixture("remote_refresh_result.canonical.json").trim() + ); + + let result = RefreshColumnResult::from_json(&fixture( + "remote_refresh_result_without_published_version.json", + )) + .expect("optional version"); + assert_eq!(result.published_version, None); + assert_eq!( + result + .to_canonical_json() + .expect("canonical result without version"), + fixture("remote_refresh_result_without_published_version.canonical.json").trim() + ); + assert_eq!( + RefreshColumnResult::from_json( + &result + .to_canonical_json() + .expect("canonical result without version") + ) + .expect("round-trip result without version"), + result + ); +} + +#[test] +fn unknown_fields_and_discriminators_are_forward_decodable() { + let mut result = job_result("remote_function_job.json"); + result["future_version_metadata"] = serde_json::json!({"retention_class": "catalog"}); + result["runtime"] = serde_json::json!({ + "kind": "wasm", + "module_digest": "sha256:wasm" + }); + result["signature"]["output"]["kind"] = Value::String("future_output_shape".to_string()); + + let version = FunctionVersion::from_json(&result.to_string()).expect("future remote value"); + assert_eq!(version.runtime().kind(), "wasm"); + assert_eq!(version.runtime().python_version(), None); + assert_eq!(version.signature().output.kind, "future_output_shape"); + assert_eq!( + serde_json::from_str::( + &version.to_canonical_json().expect("canonical future value") + ) + .expect("canonical JSON")["runtime"], + serde_json::json!({"kind": "wasm"}) + ); +} + +#[test] +fn floating_point_application_literals_are_rejected_consistently() { + let error = FunctionApplication::from_json(&fixture("remote_function_application_float.json")) + .unwrap_err(); + assert!( + error + .to_string() + .contains("floating-point Function literals") + ); +} + +#[test] +fn canonical_client_values_contain_secret_names_only() { + let result = job_result("remote_function_job.json"); + let version = FunctionVersion::from_json(&result.to_string()).expect("FunctionVersion result"); + let canonical: Value = serde_json::from_str( + &version + .to_canonical_json() + .expect("canonical FunctionVersion"), + ) + .expect("canonical JSON"); + + assert_eq!( + canonical["required_secrets"], + serde_json::json!(["HF_TOKEN"]) + ); + assert_no_secret_values(&canonical); +} diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_application.canonical.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_application.canonical.json new file mode 100644 index 000000000..05da9fe39 --- /dev/null +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_application.canonical.json @@ -0,0 +1 @@ +{"columns":{"normalized_text":"search_text","token_count":"search_token_count"},"function":{"name":"text_features","version":"fv_01K3TEXT"},"group_id":"fg_01K3TEXT","inputs":[{"kind":"column","parameter":"title","value":{"path":"title"}},{"kind":"column","parameter":"body","value":{"path":"body"}}],"output":{"fields":[{"arrow_type":"utf8","name":"normalized_text","nullable":false},{"arrow_type":"int64","name":"token_count","nullable":false}],"kind":"named_struct"}} diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_application.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_application.json new file mode 100644 index 000000000..44aeff460 --- /dev/null +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_application.json @@ -0,0 +1,20 @@ +{ + "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", + "token_count": "search_token_count" + }, + "future_application": {"declaration_mode": "managed"} +} diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_application_float.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_application_float.json new file mode 100644 index 000000000..47724eee0 --- /dev/null +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_application_float.json @@ -0,0 +1,8 @@ +{ + "function": {"name": "score", "version": "fv_01K3FLOAT"}, + "inputs": [ + {"parameter": "threshold", "kind": "literal", "value": 1e-7} + ], + "output": {"kind": "scalar", "arrow_type": "bool", "nullable": false}, + "group_id": "fg_01K3FLOAT" +} 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 new file mode 100644 index 000000000..c548c4a58 --- /dev/null +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_binding.canonical.json @@ -0,0 +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} 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 new file mode 100644 index 000000000..5d8193eea --- /dev/null +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_binding.json @@ -0,0 +1,15 @@ +{ + "binding_id": "fb_01K3TEXT", + "revision": 3, + "function": {"name": "text_features", "version": "fv_01K3TEXT"}, + "group_id": "fg_01K3TEXT", + "inputs": [ + {"parameter": "title", "field_id": 11, "field_path": "title", "arrow_type": "utf8", "nullable": true}, + {"parameter": "body", "field_id": 12, "field_path": "body", "arrow_type": "utf8", "nullable": true} + ], + "outputs": [ + {"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} + ], + "future_binding": {"metadata_revision": 1} +} diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_job.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_job.json new file mode 100644 index 000000000..6ba4eb226 --- /dev/null +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_job.json @@ -0,0 +1,31 @@ +{ + "job_id": "job_function_01K3", + "job_type": "create_function", + "job_state": "DONE", + "creation_ms": 1787270400000, + "spec": {"name": "embed"}, + "result": { + "name": "embed", + "version": "fv_01K3EXACT", + "artifact": { + "kind": "python_callable", + "digest": "sha256:code", + "entrypoint": "embed" + }, + "signature": { + "inputs": [{"name": "text", "arrow_type": "utf8", "nullable": true}], + "output": {"kind": "scalar", "arrow_type": "list", "nullable": false} + }, + "runtime": { + "kind": "python", + "python_version": "3.12", + "environment": {"kind": "pip", "packages": ["sentence-transformers>=3"]}, + "env": {"TOKENIZERS_PARALLELISM": "false"} + }, + "runtime_digest": "sha256:runtime", + "environment_digest": "sha256:environment", + "required_secrets": ["HF_TOKEN"], + "created_at": "2026-08-21T00:00:00Z" + }, + "future_job": {"trace_id": "trace-1"} +} diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json new file mode 100644 index 000000000..7ab632a98 --- /dev/null +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json @@ -0,0 +1 @@ +{"artifact":{"digest":"sha256:code","entrypoint":"embed","kind":"python_callable"},"created_at":"2026-08-21T00:00:00Z","environment_digest":"sha256:environment","name":"embed","required_secrets":["HF_TOKEN"],"runtime":{"env":{"TOKENIZERS_PARALLELISM":"false"},"environment":{"kind":"pip","packages":["sentence-transformers>=3"]},"kind":"python","python_version":"3.12"},"runtime_digest":"sha256:runtime","signature":{"inputs":[{"arrow_type":"utf8","name":"text","nullable":true}],"output":{"arrow_type":"list","kind":"scalar","nullable":false}},"version":"fv_01K3EXACT"} diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_refresh_job.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_refresh_job.json new file mode 100644 index 000000000..bb2490bc8 --- /dev/null +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_refresh_job.json @@ -0,0 +1,15 @@ +{ + "job_id": "job_refresh_01K3", + "job_type": "refresh_function_columns", + "job_state": "DONE", + "creation_ms": 1787270400001, + "spec": {"table": "documents", "binding_revision": 3}, + "result": { + "rows_assigned": 999998800, + "rows_failed": 0, + "rows_remaining": 0, + "source_version": 812, + "published_version": 919, + "future_result": {"committed_fragment_groups": 100} + } +} diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_refresh_result.canonical.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_refresh_result.canonical.json new file mode 100644 index 000000000..cda9287a3 --- /dev/null +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_refresh_result.canonical.json @@ -0,0 +1 @@ +{"published_version":919,"rows_assigned":999998800,"rows_failed":0,"rows_remaining":0,"source_version":812} diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_refresh_result_without_published_version.canonical.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_refresh_result_without_published_version.canonical.json new file mode 100644 index 000000000..69c9c96c1 --- /dev/null +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_refresh_result_without_published_version.canonical.json @@ -0,0 +1 @@ +{"rows_assigned":120,"rows_failed":0,"rows_remaining":0,"source_version":812} diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_refresh_result_without_published_version.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_refresh_result_without_published_version.json new file mode 100644 index 000000000..27231411e --- /dev/null +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_refresh_result_without_published_version.json @@ -0,0 +1,6 @@ +{ + "rows_assigned": 120, + "rows_failed": 0, + "rows_remaining": 0, + "source_version": 812 +} diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_unit_job.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_unit_job.json new file mode 100644 index 000000000..7f2b4c6de --- /dev/null +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_unit_job.json @@ -0,0 +1,9 @@ +{ + "job_id": "job_index_01K3", + "job_type": "create_index", + "job_state": "DONE", + "creation_ms": 1787270400002, + "spec": {"column": "vector"}, + "result": {"future_information": "ignored by Job<()>"}, + "future_job": {"trace_id": "trace-2"} +}