From 10c99734795f3c639978d6f1de49d123bf6e7aa9 Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Fri, 30 Oct 2026 05:50:29 -0700 Subject: [PATCH] docs: explain vectorized blob UDFs Document vectorized `pyarrow.Array` UDFs for Blob v2 inputs and outputs. Clarify annotation-selected array execution and the adapter-v1 client/executor boundary. Add public regression coverage for registering, hydrating, executing, publishing, and reading back Blob arrays, including large, empty, and null values. --- python/python/lancedb/functions.py | 23 ++- .../tests/test_first_class_function_slice2.py | 166 ++++++++++++++++++ rust/lancedb/src/function.rs | 9 +- 3 files changed, 193 insertions(+), 5 deletions(-) diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 2125542af..64804dc13 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -1839,7 +1839,11 @@ class UdfDefinition: Instances are created with :func:`udf`. Calling an instance executes the original Python function, or constructs the original class, which keeps local unit testing ordinary. Remote execution adapts the callable to the - internal Arrow batch ABI described by the registration artifact. + internal Arrow batch ABI described by the registration artifact. By + default it invokes the callable once per row. A callable whose parameters + are all annotated as ``pyarrow.Array`` receives whole Arrow arrays and + returns a ``pyarrow.Array`` instead. This authoring choice is independent + of the internal artifact adapter in the registration request. """ def __init__( @@ -2144,7 +2148,10 @@ def udf( Parameters ---------- function : function or class, optional - The synchronous callable, or the class whose instances are called. + The synchronous callable, or the class whose instances are called. For + a function, annotate every parameter as ``pyarrow.Array`` to opt into + vectorized execution; a matching return annotation documents that the + callable returns an Arrow array. name : str, optional The remote Function name. Defaults to the callable name. Must start with an ASCII letter or underscore and contain only ASCII @@ -2224,6 +2231,18 @@ def udf( 3.0 >>> [field.name for field in scale.registration_request.signature.initialization] ['factor', 'offset'] + >>> import lancedb + >>> @udf( + ... input_schema=pa.schema([lancedb.blob("image", nullable=False)]), + ... output_schema=lancedb.blob("copy", nullable=False), + ... ) + ... def copy_blobs(image: pa.Array) -> pa.Array: + ... return image + + Blob v2 input is exposed as ``pyarrow.LargeBinaryArray`` in vectorized + mode (and as ``bytes`` in row mode). Return a ``LargeBinaryArray`` with the + declared nullability. Use :func:`lancedb.blob` in the explicit schemas; + plain ``pyarrow.large_binary()`` does not carry Blob v2 semantics. """ def decorate(target: Any) -> UdfDefinition: diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 562a90eb4..e1e3784d6 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -7,6 +7,7 @@ import base64 import contextlib import functools import importlib.util +import inspect import types from datetime import date import http.server @@ -15,6 +16,7 @@ import os from pathlib import Path import subprocess import sys +import typing import urllib.parse from typing import Optional @@ -474,6 +476,33 @@ def _run_packaged(definition, *args): return namespace[definition.registration_request.artifact.entrypoint](*args) +def _execute_registered_array_udf_v1(request: dict, batch: pa.RecordBatch) -> pa.Array: + """Exercise adapter-v1 annotation dispatch over a registered artifact.""" + assert request["artifact"]["adapter"] == { + "kind": "scalar_to_arrow_batch", + "version": 1, + } + source = base64.b64decode(request["artifact"]["content"]["data"]) + namespace = { + "pa": pa, + "pyarrow": pa, + "Array": pa.Array, + "RecordBatch": pa.RecordBatch, + "typing": typing, + } + exec(compile(source, "", "exec"), namespace) + function = namespace[request["artifact"]["entrypoint"]] + parameters = tuple(inspect.signature(function).parameters.values()) + annotations = typing.get_type_hints(function) + assert parameters and all( + annotations.get(parameter.name) is pa.Array for parameter in parameters + ) + assert annotations.get("return") is pa.Array + result = function(*(batch.column(parameter.name) for parameter in parameters)) + assert isinstance(result, pa.Array) + return result + + def test_udf_conda_environment(): @udf(conda=["scipy", "numpy"], conda_channels=["conda-forge", "defaults"]) def halve(value: float) -> float: @@ -1229,6 +1258,57 @@ def test_blob_signature_rejects_collection_below_a_struct(): return len(value["images"]) +def test_blob_fields_support_vectorized_pyarrow_arrays(): + @udf( + input_schema=pa.schema([lancedb.blob("image", nullable=True)]), + output_schema=lancedb.blob("result", nullable=False), + ) + def copy_blobs(image: pa.Array) -> pa.Array: + return pa.array( + [value.as_py() if value.is_valid else b"" for value in image], + type=pa.large_binary(), + ) + + values = pa.array([b"large blob", b"", None], type=pa.large_binary()) + request = json.loads(copy_blobs.registration_request.to_canonical_json()) + batch = pa.RecordBatch.from_arrays([values], names=["image"]) + result = _execute_registered_array_udf_v1(request, batch) + assert isinstance(result, pa.LargeBinaryArray) + assert result.to_pylist() == [b"large blob", b"", b""] + + signature = copy_blobs.registration_request.signature + assert signature.inputs[0].arrow_type == "blob_v2" + assert signature.inputs[0].nullable is True + assert signature.output.arrow_type == "blob_v2" + assert signature.output.nullable is False + assert copy_blobs.registration_request.artifact.adapter.kind == ( + "scalar_to_arrow_batch" + ) + assert copy_blobs.registration_request.artifact.adapter.version == 1 + + source = base64.b64decode( + copy_blobs.registration_request.artifact.content.data + ).decode("utf-8") + assert "def copy_blobs(image: pa.Array) -> pa.Array:" in source + + +def test_documented_blob_array_identity_executes_packaged_artifact(): + @udf( + input_schema=pa.schema([lancedb.blob("image", nullable=False)]), + output_schema=lancedb.blob("copy", nullable=False), + ) + def copy_blobs(image: pa.Array) -> pa.Array: + return image + + values = pa.array([b"large blob", b""], type=pa.large_binary()) + request = json.loads(copy_blobs.registration_request.to_canonical_json()) + result = _execute_registered_array_udf_v1( + request, + pa.RecordBatch.from_arrays([values], names=["image"]), + ) + assert result.equals(values) + + def test_named_struct_function_can_include_a_blob_result_field(): @udf( input_schema=pa.schema([lancedb.blob("image", nullable=False)]), @@ -1883,6 +1963,92 @@ def test_building_a_binding_contacts_no_server(): assert state["requests"] == [] +def test_registered_vectorized_blob_udf_hydrates_executes_and_publishes(tmp_path): + @udf( + input_schema=pa.schema([lancedb.blob("image", nullable=True)]), + output_schema=lancedb.blob("copy", nullable=False), + ) + def copy_blobs(image: pa.Array) -> pa.Array: + assert isinstance(image, pa.LargeBinaryArray) + return pa.array( + [ + b"missing" if value.as_py() is None else b"copy:" + value.as_py() + for value in image + ], + type=pa.large_binary(), + ) + + with _mock_remote_function_catalog() as (host, state): + remote = lancedb.connect( + "db://dev", + api_key="fake", + host_override=host, + client_config={"retry_config": {"retries": 0}}, + ) + remote.create_function(copy_blobs) + registration = state["requests"][0][1] + + definition = json.loads(copy_blobs.registration_request.to_canonical_json()) + assert registration["artifact"] == definition["artifact"] + assert registration["signature"] == definition["signature"] + + local = lancedb.connect(tmp_path) + source = local.create_table( + "blob_input", + schema=pa.schema( + [pa.field("id", pa.int64()), lancedb.blob("image", nullable=True)] + ), + ) + source.add( + pa.Table.from_arrays( + [ + pa.array([0, 1, 2], type=pa.int64()), + pa.array([b"large blob", b"", None], type=pa.large_binary()), + ], + names=["id", "image"], + ) + ) + input_hits = ( + source.search() + .with_row_id(True) + .limit(3) + .to_arrow() + .sort_by([("id", "ascending")]) + ) + hydrated = source.fetch_blobs("image", input_hits) + assert isinstance(hydrated, pa.LargeBinaryArray) + result = _execute_registered_array_udf_v1( + registration, + pa.RecordBatch.from_arrays([hydrated], names=["image"]), + ) + assert result.to_pylist() == [b"copy:large blob", b"copy:", b"missing"] + + published = local.create_table( + "blob_output", + schema=pa.schema( + [pa.field("id", pa.int64()), lancedb.blob("copy", nullable=False)] + ), + ) + published.add( + pa.Table.from_arrays( + [pa.array(range(len(result)), type=pa.int64()), result], + names=["id", "copy"], + ) + ) + output_hits = ( + published.search() + .with_row_id(True) + .limit(3) + .to_arrow() + .sort_by([("id", "ascending")]) + ) + assert published.fetch_blobs("copy", output_hits).to_pylist() == [ + b"copy:large blob", + b"copy:", + b"missing", + ] + + def test_blocking_remote_registration_returns_function_version(): with _mock_remote_function_catalog() as (host, state): db = lancedb.connect( diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index 128a7efdf..be06be2fa 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -559,9 +559,12 @@ pub struct FunctionArtifactContent { /// 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. +/// This is part of the executor request envelope, not the user-facing UDF +/// authoring mode. For version 1, the executor selects vectorized execution +/// when every callable parameter is annotated as `pyarrow.Array`; it passes one +/// Arrow array per declared input and requires an Arrow array matching the +/// declared output. Other callables are invoked row by row. Both modes use the +/// same internal adapter kind. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct PythonAdapterSpec { pub kind: String,