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.
This commit is contained in:
Jack Ye
2026-10-25 09:50:27 +08:00
committed by GitHub
parent 0ea932dafd
commit 6b6778e684
8 changed files with 359 additions and 21 deletions
+1
View File
@@ -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
+120
View File
@@ -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)
+59 -21
View File
@@ -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(
+59
View File
@@ -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)
@@ -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)]),