mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
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.
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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, "<registered-udf>", "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(
|
||||
|
||||
Reference in New Issue
Block a user