From 3994b8fcecbac1b053863df59a36972d39af7add Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 21 Aug 2026 21:36:00 +0800 Subject: [PATCH] feat(python): expose remote function refresh jobs --- python/python/lancedb/_lancedb.pyi | 8 ++++- python/python/lancedb/expr.py | 5 +-- python/python/lancedb/functions.py | 37 ++++++++++++++++++++ python/python/lancedb/job.py | 24 +++++++++++++ python/python/lancedb/table.py | 4 ++- python/src/job.rs | 56 ++++++++++++++++++++++++++++++ python/src/lib.rs | 1 + python/src/table.rs | 2 +- rust/lancedb/src/job.rs | 9 +++++ 9 files changed, 141 insertions(+), 5 deletions(-) diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 59537f45a..49d3bb671 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -235,6 +235,12 @@ class FunctionJob: async def wait(self) -> str: ... async def cancel(self) -> None: ... +class RefreshJob: + id: 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: ... @@ -354,7 +360,7 @@ class Table: self, application_json: str, output_name: Optional[str] ) -> AddColumnsResult: ... async def refresh_column(self, column: str) -> RefreshColumnResult: ... - async def refresh_column_async(self, column: str) -> Job: ... + async def refresh_column_async(self, column: str) -> RefreshJob: ... async def add_columns_with_schema(self, schema: pa.Schema) -> AddColumnsResult: ... async def alter_columns( self, columns: list[dict[str, Any]] diff --git a/python/python/lancedb/expr.py b/python/python/lancedb/expr.py index e8b2d63a4..d16ba95d7 100644 --- a/python/python/lancedb/expr.py +++ b/python/python/lancedb/expr.py @@ -85,8 +85,9 @@ class Expr: # for dict keys / set membership. __hash__ = None # type: ignore[assignment] - def __init__(self, inner: PyExpr) -> None: + def __init__(self, inner: PyExpr, *, column_path: str | None = None) -> None: self._inner = inner + self._column_path = column_path # ── comparisons ────────────────────────────────────────────────────────── @@ -273,7 +274,7 @@ def col(name: str) -> Expr: >>> col("age") > lit(18) Expr((age > 18)) """ - return Expr(expr_col(name)) + return Expr(expr_col(name), column_path=name) def lit(value: Union[bool, int, float, str, bytes, date, datetime, Decimal]) -> Expr: diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 1613e03b4..5979ba5fe 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -20,6 +20,7 @@ import re import sys import textwrap import types +import uuid from collections.abc import Mapping from datetime import date, datetime from typing import ( @@ -265,6 +266,42 @@ class FunctionVersion(_RemoteValue): required_secrets: tuple[str, ...] = () created_at: str + def __call__(self, **inputs: Any) -> FunctionApplication: + """Bind this exact version to table columns as one grouped application.""" + from lancedb.expr import Expr + + parameters = tuple(parameter.name for parameter in self.signature.inputs) + missing = [parameter for parameter in parameters if parameter not in inputs] + unknown = sorted(set(inputs) - set(parameters)) + if missing or unknown: + details = [] + if missing: + details.append(f"missing inputs: {missing!r}") + if unknown: + details.append(f"unknown inputs: {unknown!r}") + raise TypeError("invalid Function inputs (" + "; ".join(details) + ")") + + bindings = [] + for parameter in parameters: + value = inputs[parameter] + if not isinstance(value, Expr) or value._column_path is None: + raise TypeError( + f"Function input {parameter!r} must be a direct col(...) reference" + ) + bindings.append( + ApplicationInput( + parameter=parameter, + kind="column", + value={"path": value._column_path}, + ) + ) + return FunctionApplication( + function=FunctionVersionRef(name=self.name, version=self.version), + inputs=tuple(bindings), + output=self.signature.output, + group_id=f"fg_{uuid.uuid4().hex}", + ) + class FunctionRegistrationRequest(_RemoteValue): """Stable remote registration envelope produced by :func:`udf`. diff --git a/python/python/lancedb/job.py b/python/python/lancedb/job.py index 7bd600a74..3283dbe1c 100644 --- a/python/python/lancedb/job.py +++ b/python/python/lancedb/job.py @@ -130,3 +130,27 @@ class _FunctionJobAdapter: def _function_job(inner: "_lancedb.FunctionJob") -> AsyncJob[FunctionVersion]: return AsyncJob(_FunctionJobAdapter(inner)) + + +class _RefreshJobAdapter: + def __init__(self, inner: "_lancedb.RefreshJob"): + 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): + from .functions import RefreshColumnResult + + return RefreshColumnResult.from_json(await self._inner.wait()) + + async def cancel(self): + await self._inner.cancel() + + +def _refresh_job(inner: "_lancedb.RefreshJob") -> AsyncJob: + return AsyncJob(_RefreshJobAdapter(inner)) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 913ab5289..133e8d087 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -6148,7 +6148,9 @@ class AsyncTable: >>> asyncio.run(refresh_in_background()) 'finished' """ - return AsyncJob(await self._inner.refresh_column_async(column)) + from .job import _refresh_job + + return _refresh_job(await self._inner.refresh_column_async(column)) async def alter_columns( self, *alterations: Iterable[dict[str, Any]] diff --git a/python/src/job.rs b/python/src/job.rs index a08b958a6..c4d634e71 100644 --- a/python/src/job.rs +++ b/python/src/job.rs @@ -22,6 +22,12 @@ pub struct FunctionJob { inner: Arc>, } +/// Python bridge for a typed remote Function-column refresh job. +#[pyclass] +pub struct RefreshJob { + inner: Arc, +} + impl FunctionJob { pub(crate) fn new(inner: lancedb::Job) -> Self { Self { @@ -38,6 +44,56 @@ impl Job { } } +impl RefreshJob { + pub(crate) fn new(inner: lancedb::Job) -> Self { + Self { + inner: Arc::new(inner), + } + } +} + +#[pymethods] +impl RefreshJob { + #[getter] + pub fn id(&self) -> Option { + self.inner.id().map(str::to_string) + } + + pub fn status(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner.clone(); + future_into_py( + self_.py(), + async move { inner.status().await.infer_error() }, + ) + } + + pub fn wait(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + let value = inner.wait_raw_json().await.infer_error()?.ok_or_else(|| { + pyo3::exceptions::PyValueError::new_err( + "successful refresh job did not contain a result", + ) + })?; + let result: lancedb::function::RefreshColumnResult = serde_json::from_value(value) + .map_err(|error| { + pyo3::exceptions::PyValueError::new_err(format!( + "invalid refresh job result: {error}" + )) + })?; + result.to_canonical_json().infer_error() + }) + } + + pub fn cancel(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + inner.cancel().await.infer_error()?; + Ok(()) + }) + } +} + #[pymethods] impl FunctionJob { #[getter] diff --git a/python/src/lib.rs b/python/src/lib.rs index 756b3557f..99cb8c54e 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -48,6 +48,7 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/python/src/table.rs b/python/src/table.rs index cb4752cce..be931672d 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1584,7 +1584,7 @@ impl Table { let inner = self_.inner_ref()?.clone(); future_into_py(self_.py(), async move { let job = inner.refresh_column_async(column).await.infer_error()?; - Ok(crate::job::Job::new(job)) + Ok(crate::job::RefreshJob::new(job)) }) } diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index 1a76c7683..f7aeb6eda 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -120,6 +120,15 @@ impl Job<()> { pub(crate) fn spawned(task: JoinHandle>) -> Self { Self::new(Box::new(SpawnedJob::new(task))) } + + /// Wait for an untyped job while preserving a remote terminal payload. + #[doc(hidden)] + pub async fn wait_raw_json(&self) -> Result> { + match &self.inner { + JobInner::Handle { handle, .. } => Ok(handle.wait().await?.value), + JobInner::Completed(()) => Ok(None), + } + } } impl Job