diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 235049f97..9f8e1b649 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -12,6 +12,7 @@ __version__ = importlib.metadata.version("lancedb") from ._lancedb import connect as lancedb_connect from ._lancedb import FtsToken +from ._lancedb import Function from ._lancedb import tokenize as _tokenize from .common import URI, sanitize_uri from urllib.parse import urlparse @@ -507,6 +508,7 @@ __all__ = [ "FtsToken", "col", "Expr", + "Function", "func", "lit", "URI", diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index f87fd3d13..387c85e80 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -216,11 +216,21 @@ class BlobFile: def read_range(self, offset: int, length: int) -> bytes: ... def read_up_to(self, length: int) -> bytes: ... +class Function: + @property + def id(self) -> str: ... + @property + def parameters(self) -> tuple[tuple[str, pa.DataType], ...]: ... + @property + def output_type(self) -> pa.DataType: ... + @property + def output_nullable(self) -> bool: ... + class Job: @property def id(self) -> Optional[str]: ... async def status(self) -> str: ... - async def wait(self) -> None: ... + async def wait(self) -> Optional[Function]: ... async def cancel(self) -> None: ... class JobInfo: @@ -256,6 +266,8 @@ class JobDescription: def spec_json(self) -> Optional[str]: ... @property def failure(self) -> Optional[JobFailureInfo]: ... + @property + def result(self) -> Optional[Function]: ... class Table: def name(self) -> str: ... diff --git a/python/python/lancedb/job.py b/python/python/lancedb/job.py index d33b62cbf..c83e096db 100644 --- a/python/python/lancedb/job.py +++ b/python/python/lancedb/job.py @@ -10,6 +10,7 @@ from typing import Optional from lancedb.background_loop import LOOP from . import _lancedb +from ._lancedb import Function class AsyncJob: @@ -44,18 +45,22 @@ 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) -> Optional[Function]: """Wait until the operation reaches a terminal state. + Returns the success result when present (currently a + :class:`~lancedb.Function`), or `None` when the job finished without + one. + Raises `JobFailedError` if the operation failed, `JobCancelledError` if it was cancelled, and `TimeoutError` if `timeout` elapses first. """ if self._inner is None: - return + return None if timeout is None: - await self._inner.wait() + return await self._inner.wait() else: - await asyncio.wait_for(self._inner.wait(), timeout.total_seconds()) + return await asyncio.wait_for(self._inner.wait(), timeout.total_seconds()) async def cancel(self): """Request cancellation. Cancelling a finished operation is a no-op.""" @@ -88,15 +93,19 @@ class Job: return "finished" return LOOP.run(self._inner.status()) - def wait(self, timeout: Optional[timedelta] = None): + def wait(self, timeout: Optional[timedelta] = None) -> Optional[Function]: """Block until the operation reaches a terminal state. + Returns the success result when present (currently a + :class:`~lancedb.Function`), or `None` when the job finished without + one. + 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 None + return LOOP.run(self._inner.wait(timeout)) def cancel(self): """Request cancellation. Cancelling a finished operation is a no-op.""" diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index ce8d5bd6e..9463c387a 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -2306,3 +2306,228 @@ def test_remote_connection_jobs_surface(): assert job.status() == "failed" with pytest.raises(JobFailedError, match="worker died"): job.wait(timeout=timedelta(seconds=5)) + + +# Pinned Rust-canonical schema-only type IPC (base64). PyArrow's schema-only +# FileWriter bytes are not byte-identical to the Arrow Rust FileWriter used by +# the strict Function decoder, so these fixtures are derived from Rust serde. +_FIRST_CLASS_FUNCTION_JOB_RESULT_INT32_TYPE_IPC_B64 = ( + "QVJST1cxAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP" + "////94AAAAEAAAAAAACgAMAAoACQAEAAoAAAAQAAAAAAEEAAgACAAAAAQACAAAAAQAAAABAAAAFAAAABAAFAAQ" + "AA4ADwAEAAAACAAQAAAAGAAAACAAAAAAAAECHAAAAAgADAAEAAsACAAAACAAAAAAAAABAAAAAAAAAAAAAAAA/" + "////wAAAAAUAAAAAAAAAAwAFAASAAwACAAEAAwAAABsAAAAcAAAABAAAAAAAAQACAAIAAAABAAIAAAABAAAAA" + "EAAAAUAAAAEAAUABAADgAPAAQAAAAIABAAAAAYAAAAIAAAAAAAAQIcAAAACAAMAAQACwAIAAAAIAAAAAAAAAE" + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACQAAAAQVJST1cx" +) +_FIRST_CLASS_FUNCTION_JOB_RESULT_UTF8_TYPE_IPC_B64 = ( + "QVJST1cxAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP" + "////94AAAAEAAAAAAACgAMAAoACQAEAAoAAAAQAAAAAAEEAAgACAAAAAQACAAAAAQAAAABAAAAFAAAABAAFAAQ" + "AA4ADwAEAAAACAAQAAAAGAAAAAwAAAAAAAEFEAAAAAAAAAAEAAQABAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/" + "////wAAAAAQAAAADAAUABIADAAIAAQADAAAAGAAAABkAAAAEAAAAAAABAAIAAgAAAAEAAgAAAAEAAAAAQAAAB" + "QAAAAQABQAEAAOAA8ABAAAAAgAEAAAABgAAAAMAAAAAAABBRAAAAAAAAAABAAEAAQAAAAAAAAAAAAAAAAAAAA" + "AAAAAAAAAAIAAAABBUlJPVzE=" +) +_FIRST_CLASS_FUNCTION_JOB_RESULT_FUNCTION_ID = "fn.exact.python-job-result" +_FIRST_CLASS_FUNCTION_JOB_RESULT_ABSENT = object() +_FIRST_CLASS_FUNCTION_JOB_RESULT_NULL = object() + + +def _first_class_function_job_result_function_wire(): + int32_ipc = _FIRST_CLASS_FUNCTION_JOB_RESULT_INT32_TYPE_IPC_B64 + utf8_ipc = _FIRST_CLASS_FUNCTION_JOB_RESULT_UTF8_TYPE_IPC_B64 + return { + "kind": "function", + "format_version": 1, + "function": { + "format_version": 1, + "id": _FIRST_CLASS_FUNCTION_JOB_RESULT_FUNCTION_ID, + "signature": { + "parameters": [ + {"name": "x", "data_type_ipc": int32_ipc}, + {"name": "label", "data_type_ipc": utf8_ipc}, + ], + "output": { + "data_type_ipc": int32_ipc, + "nullable": True, + }, + }, + }, + } + + +def _first_class_function_job_result_none_wire(): + return {"kind": "none", "format_version": 1} + + +def _first_class_function_job_result_describe_body( + job_id, job_type, result=_FIRST_CLASS_FUNCTION_JOB_RESULT_ABSENT +): + body = { + "job_id": job_id, + "job_state": "DONE", + "job_type": job_type, + "creation_ms": 1, + "spec": {}, + } + if result is _FIRST_CLASS_FUNCTION_JOB_RESULT_NULL: + body["result"] = None + elif result is not _FIRST_CLASS_FUNCTION_JOB_RESULT_ABSENT: + body["result"] = result + return body + + +def _first_class_function_job_result_describe_handler(bodies_by_job_id): + def handler(request): + content_len = int(request.headers.get("Content-Length", 0)) + body = request.rfile.read(content_len) if content_len > 0 else b"" + payload = json.loads(body) if body else {} + if request.path != "/v1/jobs/describe": + request.send_response(404) + request.end_headers() + return + job_id = payload["job_id"] + if job_id not in bodies_by_job_id: + request.send_response(404) + request.end_headers() + return + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(json.dumps(bodies_by_job_id[job_id]).encode()) + + return handler + + +def _assert_exact_first_class_function_job_result(function): + assert isinstance(function, lancedb.Function) + assert function is not None + assert not isinstance(function, dict) + assert function.id == _FIRST_CLASS_FUNCTION_JOB_RESULT_FUNCTION_ID + assert function.parameters == (("x", pa.int32()), ("label", pa.utf8())) + assert function.output_type == pa.int32() + assert function.output_nullable is True + text = repr(function) + assert "Function" in text + assert _FIRST_CLASS_FUNCTION_JOB_RESULT_FUNCTION_ID in text + for token in ("definition", "source", "packages", "artifact", "digest", "secret"): + assert token not in text.lower() + + +def test_first_class_function_job_result_sync_wait_returns_exact_function(): + bodies = { + "job-register": _first_class_function_job_result_describe_body( + "job-register", + "register_function", + _first_class_function_job_result_function_wire(), + ) + } + with mock_lancedb_connection( + _first_class_function_job_result_describe_handler(bodies) + ) as db: + result = db.job("job-register").wait() + _assert_exact_first_class_function_job_result(result) + + timed_out = db.job("job-register").wait(timeout=timedelta(seconds=5)) + _assert_exact_first_class_function_job_result(timed_out) + + with pytest.raises(TypeError): + lancedb.Function() + with pytest.raises(AttributeError): + result.id = "mutated" + with pytest.raises(AttributeError): + result.parameters = () + with pytest.raises(AttributeError): + result.output_type = pa.int64() + with pytest.raises(AttributeError): + result.output_nullable = False + + +@pytest.mark.asyncio +async def test_first_class_function_job_result_async_wait_returns_exact_function(): + bodies = { + "job-register": _first_class_function_job_result_describe_body( + "job-register", + "register_function", + _first_class_function_job_result_function_wire(), + ) + } + async with mock_lancedb_connection_async( + _first_class_function_job_result_describe_handler(bodies) + ) as db: + result = await db.job("job-register").wait() + _assert_exact_first_class_function_job_result(result) + + timed_out = await db.job("job-register").wait(timeout=timedelta(seconds=5)) + _assert_exact_first_class_function_job_result(timed_out) + + +def test_first_class_function_job_result_no_result_wait_returns_none(): + bodies = { + "job-index-absent": _first_class_function_job_result_describe_body( + "job-index-absent", "create_index" + ), + "job-index-explicit": _first_class_function_job_result_describe_body( + "job-index-explicit", + "create_index", + _first_class_function_job_result_none_wire(), + ), + } + with mock_lancedb_connection( + _first_class_function_job_result_describe_handler(bodies) + ) as db: + assert db.job("job-index-absent").wait() is None + assert db.job("job-index-explicit").wait(timeout=timedelta(seconds=5)) is None + + +@pytest.mark.asyncio +async def test_first_class_function_job_result_async_no_result_wait_returns_none(): + bodies = { + "job-index-absent": _first_class_function_job_result_describe_body( + "job-index-absent", "create_index" + ), + "job-index-explicit": _first_class_function_job_result_describe_body( + "job-index-explicit", + "create_index", + _first_class_function_job_result_none_wire(), + ), + } + async with mock_lancedb_connection_async( + _first_class_function_job_result_describe_handler(bodies) + ) as db: + assert await db.job("job-index-absent").wait() is None + assert ( + await db.job("job-index-explicit").wait(timeout=timedelta(seconds=5)) + is None + ) + + +def test_first_class_function_job_result_get_job_result_projection(): + bodies = { + "job-register": _first_class_function_job_result_describe_body( + "job-register", + "register_function", + _first_class_function_job_result_function_wire(), + ), + "job-absent": _first_class_function_job_result_describe_body( + "job-absent", "create_index" + ), + "job-null": _first_class_function_job_result_describe_body( + "job-null", + "create_index", + _FIRST_CLASS_FUNCTION_JOB_RESULT_NULL, + ), + "job-explicit-none": _first_class_function_job_result_describe_body( + "job-explicit-none", + "create_index", + _first_class_function_job_result_none_wire(), + ), + } + with mock_lancedb_connection( + _first_class_function_job_result_describe_handler(bodies) + ) as db: + register_description = db.get_job("job-register") + _assert_exact_first_class_function_job_result(register_description.result) + + assert db.get_job("job-absent").result is None + assert db.get_job("job-null").result is None + assert db.get_job("job-explicit-none").result is None diff --git a/python/src/function.rs b/python/src/function.rs new file mode 100644 index 000000000..c0dcf72df --- /dev/null +++ b/python/src/function.rs @@ -0,0 +1,62 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use arrow::pyarrow::ToPyArrow; +use pyo3::{Bound, Py, PyAny, PyResult, Python, pyclass, pymethods, types::PyTuple}; + +/// Immutable first-class Function handle backed by the exact Rust value. +#[pyclass(frozen, skip_from_py_object)] +#[derive(Clone)] +pub struct Function { + inner: lancedb::function::Function, +} + +impl Function { + pub(crate) fn new(inner: lancedb::function::Function) -> Self { + Self { inner } + } + + /// Crate-private accessor for later call-authoring slices. + #[allow(dead_code)] + pub(crate) fn inner(&self) -> &lancedb::function::Function { + &self.inner + } +} + +#[pymethods] +impl Function { + #[getter] + fn id(&self) -> &str { + self.inner.id().as_str() + } + + #[getter] + fn parameters<'py>(&self, py: Python<'py>) -> PyResult> { + let parameters = self.inner.signature().parameters(); + let mut pairs = Vec::with_capacity(parameters.len()); + for parameter in parameters { + let data_type = parameter.data_type().to_pyarrow(py)?; + pairs.push((parameter.name(), data_type)); + } + PyTuple::new(py, pairs) + } + + #[getter] + fn output_type(&self, py: Python<'_>) -> PyResult> { + self.inner + .signature() + .output() + .data_type() + .to_pyarrow(py) + .map(|obj| obj.unbind()) + } + + #[getter] + fn output_nullable(&self) -> bool { + self.inner.signature().output().nullable() + } + + fn __repr__(&self) -> String { + format!("Function(id={:?})", self.inner.id().as_str()) + } +} diff --git a/python/src/job.rs b/python/src/job.rs index 56ee211f4..56ae7cda4 100644 --- a/python/src/job.rs +++ b/python/src/job.rs @@ -3,6 +3,7 @@ use std::sync::Arc; +use crate::function::Function; use crate::runtime::future_into_py; use pyo3::{Bound, PyAny, PyRef, PyResult, pyclass, pymethods}; @@ -21,6 +22,23 @@ impl Job { } } +/// Project a Rust [`lancedb::JobResult`] onto the Python success surface. +/// +/// Delegates variant interpretation to [`lancedb::JobResult::into_function`]: +/// no nested Function collapses to Python `None`; an exact Function becomes +/// the corresponding [`Function`] handle. +fn project_wait_result(result: lancedb::JobResult) -> Option { + result.into_function().map(Function::new) +} + +/// Project a describe `result` onto Python `Optional[Function]`. +/// +/// Rust `None`, `Some(JobResult::None)`, and JSON null all become Python +/// `None`. Only `Some(JobResult::Function)` becomes a [`Function`] handle. +fn project_description_result(result: Option) -> Option { + result.and_then(project_wait_result) +} + #[pymethods] impl Job { #[getter] @@ -39,8 +57,8 @@ impl Job { pub fn wait(self_: PyRef<'_, Self>) -> PyResult> { let inner = self_.inner.clone(); future_into_py(self_.py(), async move { - inner.wait().await.infer_error()?; - Ok(()) + let result = inner.wait().await.infer_error()?; + Ok(project_wait_result(result)) }) } @@ -115,6 +133,7 @@ pub struct JobDescription { creation_ms: i64, spec_json: Option, failure: Option, + result: Option, } #[pymethods] @@ -140,6 +159,7 @@ impl From for JobDescription { message: failure.message, retryable: failure.retryable, }), + result: project_description_result(description.result), } } } diff --git a/python/src/lib.rs b/python/src/lib.rs index 6b0c0cf97..1f02455b8 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -23,6 +23,7 @@ pub mod arrow; pub mod connection; pub mod error; pub mod expr; +pub mod function; pub mod header; pub mod index; pub mod job; @@ -45,6 +46,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::()?;