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.