mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-24 23:18:39 +00:00
Merge latest origin/main into gatekeeper/fix-2057-1
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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: ...
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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<float32>[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",
|
||||
]
|
||||
@@ -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<Bound<'_, PyAny>> {
|
||||
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<Bound<'_, PyAny>> {
|
||||
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<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
|
||||
@@ -13,6 +13,23 @@ pub struct Job {
|
||||
inner: Arc<lancedb::Job>,
|
||||
}
|
||||
|
||||
/// 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<lancedb::Job<lancedb::function::FunctionVersion>>,
|
||||
}
|
||||
|
||||
impl FunctionJob {
|
||||
pub(crate) fn new(inner: lancedb::Job<lancedb::function::FunctionVersion>) -> 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<String> {
|
||||
self.inner.id().map(str::to_string)
|
||||
}
|
||||
|
||||
pub fn status(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(
|
||||
self_.py(),
|
||||
async move { inner.status().await.infer_error() },
|
||||
)
|
||||
}
|
||||
|
||||
pub fn wait(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
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<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.cancel().await.infer_error()?;
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl Job {
|
||||
#[getter]
|
||||
|
||||
@@ -47,6 +47,7 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<Session>()?;
|
||||
m.add_class::<Table>()?;
|
||||
m.add_class::<crate::job::Job>()?;
|
||||
m.add_class::<crate::job::FunctionJob>()?;
|
||||
m.add_class::<crate::job::JobInfo>()?;
|
||||
m.add_class::<crate::job::JobDescription>()?;
|
||||
m.add_class::<crate::job::JobFailureInfo>()?;
|
||||
|
||||
@@ -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<crate::job::Job<crate::function::FunctionVersion>> {
|
||||
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<str>,
|
||||
version: impl AsRef<str>,
|
||||
) -> Result<crate::function::FunctionVersion> {
|
||||
self.internal
|
||||
.get_function(name.as_ref(), version.as_ref())
|
||||
.await
|
||||
}
|
||||
|
||||
/// Rename a table in the database.
|
||||
///
|
||||
/// This is only supported in LanceDB Cloud.
|
||||
|
||||
@@ -241,6 +241,12 @@ fn job_op_not_supported<T>(what: &str) -> Result<T> {
|
||||
})
|
||||
}
|
||||
|
||||
fn function_catalog_not_supported<T>() -> Result<T> {
|
||||
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<Arc<dyn BaseTable>>;
|
||||
/// Register an immutable Function version through the remote catalog.
|
||||
async fn create_function_async(
|
||||
&self,
|
||||
_request: crate::function::FunctionRegistrationRequest,
|
||||
) -> Result<crate::job::Job<crate::function::FunctionVersion>> {
|
||||
function_catalog_not_supported()
|
||||
}
|
||||
/// Look up one exact immutable Function version.
|
||||
async fn get_function(
|
||||
&self,
|
||||
_name: &str,
|
||||
_version: &str,
|
||||
) -> Result<crate::function::FunctionVersion> {
|
||||
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.
|
||||
|
||||
@@ -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<String>,
|
||||
}
|
||||
|
||||
impl_json!(FunctionRegistrationRequest);
|
||||
|
||||
/// Exact FunctionVersion reference embedded in applications and bindings.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct FunctionVersionRef {
|
||||
|
||||
@@ -126,8 +126,7 @@ impl<T> Job<T>
|
||||
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<dyn JobHandle>) -> Self {
|
||||
Self {
|
||||
inner: JobInner::Handle {
|
||||
|
||||
@@ -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<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
})
|
||||
}
|
||||
|
||||
async fn create_function_async(
|
||||
&self,
|
||||
request: FunctionRegistrationRequest,
|
||||
) -> Result<Job<FunctionVersion>> {
|
||||
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<FunctionVersion> {
|
||||
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<crate::job::Job> {
|
||||
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));
|
||||
|
||||
@@ -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"
|
||||
));
|
||||
}
|
||||
}
|
||||
+1
@@ -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}}}
|
||||
Vendored
+46
@@ -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"
|
||||
]
|
||||
}
|
||||
Reference in New Issue
Block a user