mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-22 05:58:20 +00:00
feat: add scalar function authoring and catalog client (#3991)
## Problem The canonical Function wire values and typed remote Job contract do not yet provide a Python authoring surface or catalog client, so users cannot package a scalar callable, register it, or reopen the exact immutable Function version. ## Behavior This adds scalar-only `@udf` authoring with deterministic annotation or explicit Arrow schema validation, content-addressed Python artifacts, and an internal scalar-to-Arrow-batch adapter descriptor. Registration payloads model non-secret environment values and secret names only. Remote connections can submit `create_function_async` and receive a typed `Job<FunctionVersion>`, then reopen that exact version by name and version ID. Synchronous connections can call `create_function` to submit and wait for the immutable version in one operation. Local Function catalog operations return a stable `NotSupported` error. Shared Rust/Python golden payloads and mocked catalog responses freeze the request, typed terminal result, and exact lookup contract. ## Validation - Rust formatting, remote check, clippy, and focused LDB-1/LDB-2 tests - Python formatting, lint, and focused LDB-1/LDB-2 tests - Python API documentation build
This commit is contained in:
@@ -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>()?;
|
||||
|
||||
Reference in New Issue
Block a user