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