mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-26 13:36:13 +00:00
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:
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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(
|
||||
|
||||
@@ -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)]),
|
||||
|
||||
Reference in New Issue
Block a user