From 6b6778e68495466e78cc786ebc912601c57c8b5d Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Sat, 24 Oct 2026 18:50:27 -0700 Subject: [PATCH] feat: infer lazy blob reader Function inputs (#4663) Infer Blob V2 input access from the Python Function signature. `video: BlobReader` requests a lazy, seekable reader; `video: bytes` requests the complete payload. Users do not need a separate `blob_inputs` registration option. ## Python experience For a video table joined to a table of frame timestamps, the Function can pass its reader directly to a media decoder: ```python import io import av from lancedb import BlobReader, udf @udf(pip=["av", "Pillow"]) def extract_frame(video: BlobReader, timestamp_s: float) -> bytes: with av.open(video) as container: stream = container.streams.video[0] container.seek(int(timestamp_s / stream.time_base), stream=stream) for frame in container.decode(stream): if frame.time is not None and frame.time >= timestamp_s: output = io.BytesIO() frame.to_image().save(output, format="JPEG") return output.getvalue() raise ValueError("No frame at or after the requested timestamp") ``` After registering `extract_frame`, the query keeps the usual Function call: ```sql SELECT f.video_id, f.timestamp_s, extract_frame(v.video, f.timestamp_s) AS image FROM videos AS v JOIN frames AS f ON v.id = f.video_id ``` The returned JPEG can feed a subsequent image-model step. The reader provides `read`, `readinto`, `seek`, `tell`, `size()`, and `read_range(offset, length)`; range reads leave the sequential cursor unchanged. Each row invocation owns its reader, which closes after the call. Decoder seeks can request the parts of the MP4 needed to locate and decode the requested frame; the amount read depends on the file layout and codec. Functions that need the whole value can keep a bytes signature: ```python @udf def video_length(video: bytes) -> int: return len(video) ``` ## Execution contract - Signature inference maps `BlobReader` to a Blob V2 field and records `blob_access="reader"`. `bytes` keeps its binary Arrow type and records `blob_access="bytes"`; it can also bind to a Blob V2 source. Explicit schemas support the same choice, including a Blob V2 field annotated as `bytes`. - Reader execution carries a compact blob reference into the Function worker. Payload reads happen when Python calls the reader, with requests split into at most 1 MiB ranges. The full video does not need to enter the Arrow input batch just to invoke the Function. - Bytes execution materializes the full source before the Function runs and remains subject to eager-input memory admission. Calling `reader.read()` without a size also materializes the remaining contents in Python; lazy access depends on how the Function consumes it. - Input access metadata survives Python/Rust Function serialization and registration. Packaged Functions import the worker-provided reader implementation. This PR defines the SDK contract; execution requires corresponding server support. - This changes blob inputs. The example still returns ordinary image bytes; it does not add streamed outputs or run an image model. Adds coverage for signature inference, explicit schemas, nullable readers, signature round trips, seeking, bounded range reads, buffering, and invalid responses. --- docs/src/python/python.md | 8 ++ docs/web/management/function.mdx | 34 +++++ python/python/lancedb/__init__.py | 1 + python/python/lancedb/_blob_reader.py | 120 ++++++++++++++++++ python/python/lancedb/functions.py | 80 +++++++++--- python/python/tests/test_blob_reader.py | 59 +++++++++ .../tests/test_first_class_function_slice2.py | 55 ++++++++ rust/lancedb/src/function.rs | 23 ++++ 8 files changed, 359 insertions(+), 21 deletions(-) create mode 100644 python/python/lancedb/_blob_reader.py create mode 100644 python/python/tests/test_blob_reader.py diff --git a/docs/src/python/python.md b/docs/src/python/python.md index c465d0ac4..5d1393924 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -304,6 +304,14 @@ listing a storage directory. ## Functions and Jobs +Annotate a blob input with `lancedb.BlobReader` for lazy, seekable access, or +with `bytes` to materialize its contents before the Function runs. Reader +inputs infer the Blob V2 type without an explicit input schema. Readers are +scoped to each invocation and support `read`, `readinto`, `seek`, `tell`, +`size`, and cursor-independent `read_range`. + +::: lancedb.BlobReader + ::: lancedb.functions.FunctionArtifact ::: lancedb.functions.FunctionParameter diff --git a/docs/web/management/function.mdx b/docs/web/management/function.mdx index 8e5b53e4a..6df99b5fd 100644 --- a/docs/web/management/function.mdx +++ b/docs/web/management/function.mdx @@ -205,6 +205,40 @@ Function output blobs follow the same compaction and cleanup lifecycle as ingested Managed blobs. Retained versions continue to own their payloads; cleanup reclaims an object only after its retained references are gone. +### Draft: incremental blob inputs (upcoming) + +This feature requires an SDK that exports `BlobReader` and a compatible +Enterprise runtime with reader input support enabled. Coordinate the runtime +upgrade before registering or executing reader Functions. + +Annotate a row Function parameter with `BlobReader` to read a Blob v2 input +incrementally. The signature selects the access mode without another +registration option. For example, read the first MP4 box type without loading +the whole video: + +```python +from lancedb import BlobReader, udf + + +@udf +def first_box_type(video: BlobReader) -> bytes: + video.seek(4) + return video.read(4) +``` + +The worker passes a seekable reader backed by the stored blob. `read(size)` +reads from the current cursor; `read_range(offset, length)` reads a byte range +without moving it. `size()` returns the payload size without reading it. +Each row call receives its own reader, which closes when the call returns or +raises. Do not retain it for another invocation. + +Use `video: bytes` when the Function needs the complete payload. Bytes inputs +materialize before invocation and remain subject to eager input limits. +Reader inputs keep the unread payload out of the input batch, but calling +`video.read()` without a size loads all remaining bytes in Python. Reader +inputs are supported in row mode. Return values remain ordinary values such +as `bytes`; this feature does not add streamed outputs. + ## Declare the environment {#environment} The function runs remotely, so its dependencies are part of its definition: diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 6b902c466..a3aa55209 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -24,6 +24,7 @@ from .remote import ClientConfig from .remote.db import RemoteDBConnection from .expr import Expr, col, lit, func from .schema import blob, vector +from ._blob_reader import BlobReader as BlobReader from .job import AsyncJob, Job from .sql import AsyncQuery as AsyncSqlQuery from .sql import Query as SqlQuery diff --git a/python/python/lancedb/_blob_reader.py b/python/python/lancedb/_blob_reader.py new file mode 100644 index 000000000..2fd0f00ca --- /dev/null +++ b/python/python/lancedb/_blob_reader.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""Seekable binary input for Functions that consume Blob V2 values lazily.""" + +import io +import operator +import urllib.request +from collections.abc import Callable +from http import HTTPStatus + +_READ_CHUNK_BYTES = 1024 * 1024 +_READ_TIMEOUT_SECONDS = 300 + + +class BlobReader(io.RawIOBase): + """Read a blob incrementally inside a Function. + + Annotating a Function input with ``BlobReader`` requests a lazy, seekable + reader. Annotating it with ``bytes`` requests materialized bytes instead. + ``read_range`` does not change the sequential cursor. Readers are scoped + to the invocation and must not be retained for a later call. + + For local calls, construct a reader from a range-reading callable and the + blob's size. The callable accepts an offset and length and returns bytes. + """ + + def __init__(self, read_range: Callable[[int, int], bytes], size: int): + super().__init__() + self._read_range = read_range + self._size = operator.index(size) + if self._size < 0: + raise ValueError("blob size must be nonnegative") + self._position = 0 + + @classmethod + def _from_uri(cls, uri: str, size: int): + opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) + + def read_range(offset, length): + request = urllib.request.Request( + uri, headers={"Range": f"bytes={offset}-{offset + length - 1}"} + ) + with opener.open(request, timeout=_READ_TIMEOUT_SECONDS) as response: + if response.status != HTTPStatus.PARTIAL_CONTENT: + raise OSError( + "blob range service did not return a partial response" + ) + data = response.read(length + 1) + if len(data) != length: + raise OSError( + "blob range service returned an invalid payload length" + ) + return data + + return cls(read_range, size) + + def readable(self) -> bool: + return not self.closed + + def seekable(self) -> bool: + return not self.closed + + def size(self) -> int: + """Return the blob size without reading its payload.""" + self._checkClosed() + return self._size + + def tell(self) -> int: + self._checkClosed() + return self._position + + def seek(self, offset: int, whence: int = io.SEEK_SET) -> int: + self._checkClosed() + offset = operator.index(offset) + if whence == io.SEEK_SET: + position = offset + elif whence == io.SEEK_CUR: + position = self._position + offset + elif whence == io.SEEK_END: + position = self._size + offset + else: + raise ValueError(f"invalid whence: {whence}") + if position < 0: + raise ValueError("negative seek position") + self._position = position + return position + + def read_range(self, offset: int, length: int) -> bytes: + """Read a byte range without changing the cursor; truncate at EOF.""" + self._checkClosed() + offset, length = operator.index(offset), operator.index(length) + if offset < 0 or length < 0: + raise ValueError("offset and length must be nonnegative") + remaining = min(length, max(0, self._size - offset)) + chunks = [] + while remaining: + count = min(remaining, _READ_CHUNK_BYTES) + chunk = self._read_range(offset, count) + if len(chunk) != count: + raise OSError("blob range reader returned an invalid payload length") + chunks.append(chunk) + offset += count + remaining -= count + return b"".join(chunks) + + def read(self, size: int = -1) -> bytes: + self._checkClosed() + size = -1 if size is None else operator.index(size) + if size < 0: + size = max(0, self._size - self._position) + data = self.read_range(self._position, size) + self._position += len(data) + return data + + def readinto(self, buffer) -> int: + view = memoryview(buffer).cast("B") + data = self.read(len(view)) + view[: len(data)] = data + return len(data) diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 7f5776b32..42c4b27e1 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -33,6 +33,7 @@ from typing import ( Annotated, Any, Callable, + Literal, Optional, Union, get_args, @@ -52,7 +53,9 @@ from pydantic import ( model_validator, ) +from .schema import blob as _blob_field from .schema import is_blob_v2_field as _is_blob_v2_field +from ._blob_reader import BlobReader from .secrets import EnvVarSecret _Int32 = conint(strict=True, ge=-(2**31), le=2**31 - 1) @@ -213,6 +216,7 @@ class FunctionParameter(_RemoteValue): name: str arrow_type: str nullable: bool + blob_access: Optional[Literal["bytes", "reader"]] = None class FunctionResultField(_OpenRemoteValue): @@ -1096,6 +1100,43 @@ def _infer_signature( if (input_schema is None) != (output_schema is None): raise ValueError("input_schema and output_schema must be provided together") + try: + annotations = get_type_hints(function, include_extras=True) + except Exception as error: + raise TypeError(f"failed to resolve Function annotations: {error}") from error + + def input_parameter(field): + annotation = annotations.get(field.name) + if get_origin(annotation) in (Union, types.UnionType): + choices = [ + value for value in get_args(annotation) if value is not type(None) + ] + if len(choices) == 1: + annotation = choices[0] + arrow_type = _canonical_arrow_field(field) + if annotation is BlobReader: + if arrow_type != _FUNCTION_BLOB_V2_TYPE: + raise TypeError( + f"BlobReader input '{field.name}' requires a Blob V2 field" + ) + access = "reader" + elif annotation is bytes: + if arrow_type != _FUNCTION_BLOB_V2_TYPE and not ( + pa.types.is_binary(field.type) or pa.types.is_large_binary(field.type) + ): + raise TypeError( + f"bytes input '{field.name}' requires binary or Blob V2" + ) + access = "bytes" + else: + access = "bytes" if arrow_type == _FUNCTION_BLOB_V2_TYPE else None + return FunctionParameter( + name=field.name, + arrow_type=arrow_type, + nullable=field.nullable, + blob_access=access, + ) + if input_schema is not None: if not isinstance(input_schema, pa.Schema): raise TypeError("input_schema must be a PyArrow Schema") @@ -1110,20 +1151,9 @@ def _infer_signature( "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_field(field), - nullable=field.nullable, - ) - for field in input_schema - ) + inputs = tuple(input_parameter(field) 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 ] @@ -1132,15 +1162,21 @@ def _infer_signature( 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_field( - pa.field(parameter.name, data_type, nullable=nullable) - ), - nullable=nullable, + annotation = annotations[parameter.name] + base = annotation + nullable = False + if get_origin(base) in (Union, types.UnionType): + choices = [value for value in get_args(base) if value is not type(None)] + if len(choices) == 1: + base, nullable = choices[0], True + if base is BlobReader: + inputs.append( + input_parameter(_blob_field(parameter.name, nullable=nullable)) ) + continue + data_type, nullable = _annotation_type(annotation) + inputs.append( + input_parameter(pa.field(parameter.name, data_type, nullable=nullable)) ) output_type, output_nullable = _annotation_type(annotations["return"]) if output_nullable: @@ -1239,7 +1275,7 @@ _PYTHON_CALLABLE_ARTIFACT = "python_callable" _PYTHON_BUNDLE_ARTIFACT = "python_bundle" # Modules the Function image places beside user code. _RESERVED_SOURCE_MODULES = frozenset( - {"user_function", "lance_source", "lance_callable"} + {"user_function", "lance_source", "lance_callable", "lance_blob"} ) # The service accepts bundle paths made of these components only. _SOURCE_PATH_COMPONENT = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") @@ -1499,6 +1535,8 @@ class _SourcePackager: """One module-level line that rebinds `name` to `value` in the artifact: an import for modules and importable classes/functions, a literal otherwise.""" + if value is BlobReader: + return f"from lance_blob import BlobReader as {name}" if isinstance(value, types.ModuleType): if value.__name__.split(".")[0] in _NAMESPACE_MODULES: raise ValueError( diff --git a/python/python/tests/test_blob_reader.py b/python/python/tests/test_blob_reader.py new file mode 100644 index 000000000..dd12409af --- /dev/null +++ b/python/python/tests/test_blob_reader.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import io + +import pytest + +from lancedb import BlobReader + + +def test_reader_is_lazy_and_seekable(): + requests = [] + size = 64 * 1024 * 1024 + + def read_range(offset, length): + requests.append((offset, length)) + return bytes((index % 251 for index in range(offset, offset + length))) + + with BlobReader(read_range, size) as reader: + assert requests == [] + assert reader.size() == size + reader.seek(-4, io.SEEK_END) + assert reader.read(10) == bytes(index % 251 for index in range(size - 4, size)) + assert reader.read(1) == b"" + assert requests == [(size - 4, 4)] + assert reader.read_range(1, 3) == b"\x01\x02\x03" + assert reader.tell() == size + reader.seek(2) + target = bytearray(3) + assert reader.readinto(target) == 3 + assert target == b"\x02\x03\x04" + reader.seek(-1, io.SEEK_CUR) + assert reader.tell() == 4 + with pytest.raises(ValueError, match="closed"): + reader.read(1) + + +def test_reader_supports_buffered_io_and_bounds_range_requests(): + requests = [] + payload = b"x" * (3 * 1024 * 1024 + 7) + + def read_range(offset, length): + requests.append((offset, length)) + return payload[offset : offset + length] + + with io.BufferedReader(BlobReader(read_range, len(payload))) as reader: + assert reader.read() == payload + assert requests + assert all(length <= 1024 * 1024 for _, length in requests) + + +def test_reader_rejects_invalid_ranges_and_short_reads(): + reader = BlobReader(lambda offset, length: b"", 10) + with pytest.raises(ValueError): + reader.seek(-1) + with pytest.raises(ValueError): + reader.read_range(-1, 1) + with pytest.raises(OSError, match="payload length"): + reader.read(1) diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 9ee3e388a..874034564 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -23,6 +23,7 @@ import pyarrow as pa import pytest import lancedb +from lancedb import BlobReader from lancedb.functions import ( PythonRuntimeSpec, SecretBinding, @@ -966,6 +967,60 @@ def test_blob_fields_use_the_scalar_function_semantic_type(): assert signature.output.arrow_type == "blob_v2" +def test_blob_input_access_is_inferred_from_annotations(): + @udf + def read_blob(video: BlobReader, timestamp_s: float) -> bytes: + video.seek(int(timestamp_s)) + return video.read(4) + + @udf + def read_bytes(video: bytes) -> int: + return len(video) + + reader = read_blob.registration_request.signature + assert reader.inputs[0].arrow_type == "blob_v2" + assert reader.inputs[0].blob_access == "reader" + assert reader.inputs[1].blob_access is None + assert reader.output.arrow_type == "binary" + eager = read_bytes.registration_request.signature.inputs[0] + assert eager.arrow_type == "binary" + assert eager.blob_access == "bytes" + encoded = json.loads(read_blob.registration_request.to_canonical_json()) + assert encoded["signature"]["inputs"][0]["blob_access"] == "reader" + + +def test_nullable_blob_reader_and_explicit_blob_bytes(): + @udf + def nullable(video: Optional[BlobReader]) -> int: + return 0 if video is None else video.size() + + @udf( + input_schema=pa.schema([lancedb.blob("video", nullable=False)]), + output_schema=pa.field("size", pa.int64(), nullable=False), + ) + def eager(video: bytes) -> int: + return len(video) + + parameter = nullable.registration_request.signature.inputs[0] + assert parameter.nullable is True + assert parameter.blob_access == "reader" + parameter = eager.registration_request.signature.inputs[0] + assert parameter.arrow_type == "blob_v2" + assert parameter.blob_access == "bytes" + + +def test_blob_reader_rejects_an_incompatible_explicit_schema(): + def read(video: BlobReader) -> int: + return video.size() + + with pytest.raises(TypeError, match="requires a Blob V2 field"): + udf( + read, + input_schema=pa.schema([pa.field("video", pa.int64(), nullable=False)]), + output_schema=pa.field("size", pa.int64(), nullable=False), + ) + + def test_whole_named_struct_function_can_include_a_blob_result_field(): @udf( input_schema=pa.schema([lancedb.blob("image", nullable=False)]), diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index 98d915118..36e6c30c8 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -156,11 +156,20 @@ pub struct FunctionArtifact { } /// One ordered Arrow input parameter. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BlobAccess { + Bytes, + Reader, +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct FunctionParameter { pub name: String, pub arrow_type: String, pub nullable: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub blob_access: Option, } /// One field of an ordered named-struct result. @@ -956,6 +965,20 @@ mod conda_environment_tests { mod tests { use super::*; + #[test] + fn blob_reader_access_survives_signature_round_trip() { + let parameter: FunctionParameter = serde_json::from_value(serde_json::json!({ + "name": "video", "arrow_type": "blob_v2", "nullable": false, + "blob_access": "reader" + })) + .unwrap(); + assert_eq!(parameter.blob_access, Some(BlobAccess::Reader)); + assert_eq!( + serde_json::to_value(parameter).unwrap()["blob_access"], + "reader" + ); + } + /// Canonical form is what the FunctionVersion hash is taken over, so key /// order must come from the keys and not from however serde happened to /// emit them. Nesting is included because the sort is recursive.