diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 9a7f2c415..6cdd6e855 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -17,7 +17,7 @@ from .db import AsyncConnection, DBConnection, LanceDBConnection from .remote import ClientConfig from .remote.db import RemoteDBConnection from .expr import Expr, col, lit, func -from .schema import vector +from .schema import blob, vector, BlobType from .table import AsyncTable, Table from ._lancedb import Session from .namespace import ( @@ -467,6 +467,8 @@ __all__ = [ "lit", "URI", "sanitize_uri", + "blob", + "BlobType", "vector", "DBConnection", "LanceDBConnection", diff --git a/python/python/lancedb/_blob.py b/python/python/lancedb/_blob.py new file mode 100644 index 000000000..338c973be --- /dev/null +++ b/python/python/lancedb/_blob.py @@ -0,0 +1,420 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""Blob fetch API and v2 projection helpers.""" + +from __future__ import annotations + +import io +from collections.abc import Awaitable, Callable, Iterable +from typing import TYPE_CHECKING, Optional, Union + +import pyarrow as pa + +from .expr import Expr +from .schema import blob_v2_column_paths +from .types import BlobMode, QueryProjection, QueryProjectionSpec +from .util import get_uri_scheme + +if TYPE_CHECKING: + from _typeshed import WriteableBuffer + + from .remote.table import RemoteTable + from .table import AsyncTable, Table + +BLOB_MODE_TO_HANDLING = { + "lazy": "blobs_descriptions", + "bytes": "all_binary", + "descriptions": "blobs_descriptions", +} + +ROW_ID_FIELD_NAME = "_lance_row_id" + +FetchBlobsSync = Callable[[str, pa.Table], pa.Array | pa.ChunkedArray] +FetchBlobsAsync = Callable[[str, pa.Table], Awaitable[pa.Array | pa.ChunkedArray]] + + +class BlobFile(io.RawIOBase): + """Seekable lazy handle from :meth:`~lancedb.table.Table.fetch_blob_files`. + + Bytes load on ``read`` or ``read_range``, not when the handle is opened. + Use :meth:`aread` from async code. + """ + + def __init__(self, inner) -> None: + self._inner = inner + + async def aread(self) -> bytes: + return await self._inner.read() + + def close(self) -> None: + self._inner.close() + + @property + def closed(self) -> bool: + return self._inner.is_closed() + + def readable(self) -> bool: + return True + + def seekable(self) -> bool: + return True + + def seek(self, offset: int, whence: int = io.SEEK_SET) -> int: + if whence == io.SEEK_SET: + self._inner.seek(offset) + elif whence == io.SEEK_CUR: + self._inner.seek(self._inner.tell() + offset) + elif whence == io.SEEK_END: + self._inner.seek(self._inner.size() + offset) + else: + raise ValueError(f"invalid whence: {whence}") + return self._inner.tell() + + def tell(self) -> int: + return self._inner.tell() + + def size(self) -> int: + return self._inner.size() + + def readall(self) -> bytes: + return self._inner.read_bytes() + + def read(self, size: int = -1) -> bytes: + if size == -1: + return self._inner.read_bytes() + return super().read(size) + + def read_range(self, offset: int, length: int) -> bytes: + return self._inner.read_range(offset, length) + + def readinto(self, b: WriteableBuffer) -> int: + view = memoryview(b).cast("B") + chunk = self._inner.read_up_to(len(view)) + view[: len(chunk)] = chunk + return len(chunk) + + def __repr__(self) -> str: + return f"" + + +def validate_blob_mode(blob_mode: BlobMode) -> None: + if blob_mode not in BLOB_MODE_TO_HANDLING: + modes = ", ".join(repr(mode) for mode in BLOB_MODE_TO_HANDLING) + raise ValueError(f"blob_mode must be one of {modes}, got {blob_mode!r}") + + +def supports_blob_auto_row_id(table: Table | AsyncTable | RemoteTable) -> bool: + """Blob auto row-id applies to native tables, not LanceDB Cloud.""" + from .remote.table import RemoteTable + + if isinstance(table, RemoteTable): + return False + + inner = getattr(table, "_inner", None) + if inner is not None: + uri = inner.database().uri + if isinstance(uri, str) and get_uri_scheme(uri) == "db": + return False + + return True + + +def projection_includes_blob_column( + projection: QueryProjection, + blob_columns: Iterable[str], +) -> bool: + columns = set(blob_columns) + if not columns: + return False + if projection is None: + return True + for output, source in _iter_projection_pairs(projection): + if output in columns or source in columns: + return True + return False + + +def blob_v2_projection_sources( + schema: pa.Schema, + projection: QueryProjection, +) -> dict[str, str]: + blob_columns = blob_v2_column_paths(schema) + if not blob_columns: + return {} + columns = set(blob_columns) + if projection is None: + return {column: column for column in blob_columns} + return { + output: source + for output, source in _iter_projection_pairs(projection) + if source in columns + } + + +def v2_projection_needs_row_id( + schema: pa.Schema, + projection: QueryProjection, + *, + with_row_id: bool, +) -> bool: + if with_row_id: + return False + return projection_includes_blob_column(projection, blob_v2_column_paths(schema)) + + +def blob_auto_row_id_for_scan( + table: Table | AsyncTable | RemoteTable, + schema: pa.Schema, + projection: QueryProjection, + *, + with_row_id: bool | None, +) -> bool: + if with_row_id is not None: + return False + if not supports_blob_auto_row_id(table): + return False + return v2_projection_needs_row_id(schema, projection, with_row_id=False) + + +def finalize_blob_query_table( + tbl: pa.Table, + *, + user_requested_row_id: bool, + blob_auto_row_id: bool, + blob_paths: Iterable[str] = (), +) -> pa.Table: + if user_requested_row_id or not blob_auto_row_id: + return tbl + return stash_auto_row_ids(tbl, blob_paths) + + +async def replace_v2_blob_columns_with_bytes( + tbl: pa.Table, + blob_sources: dict[str, str], + fetch_blobs: FetchBlobsAsync, +) -> pa.Table: + for output_name, source_name in blob_sources.items(): + if output_name not in tbl.column_names: + continue + blobs = await fetch_blobs(source_name, tbl) + tbl = _set_blob_column(tbl, output_name, blobs) + return tbl + + +def replace_v2_blob_columns_with_bytes_sync( + tbl: pa.Table, + blob_sources: dict[str, str], + fetch_blobs: FetchBlobsSync, +) -> pa.Table: + for output_name, source_name in blob_sources.items(): + if output_name not in tbl.column_names: + continue + blobs = fetch_blobs(source_name, tbl) + tbl = _set_blob_column(tbl, output_name, blobs) + return tbl + + +def stash_auto_row_ids(tbl: pa.Table, blob_paths: Iterable[str]) -> pa.Table: + if "_rowid" not in tbl.column_names: + raise ValueError("query result has no '_rowid' column to hide") + + present_paths = [p for p in blob_paths if p.split(".")[0] in tbl.column_names] + if not present_paths: + raise ValueError("query result has no blob v2 column to carry a row id") + + row_ids = tbl["_rowid"] + if isinstance(row_ids, pa.ChunkedArray): + row_ids = row_ids.combine_chunks() + row_ids = row_ids.cast(pa.uint64()) + + for path in present_paths: + tbl = _embed_row_id_in_column(tbl, path, row_ids) + return tbl.drop_columns(["_rowid"]) + + +def read_row_ids_from_hits(hits: pa.Table, blob_column: str) -> list[int]: + if "_rowid" in hits.column_names: + return hits["_rowid"].to_pylist() + + try: + leaf = _leaf_struct_column(hits, blob_column) + if ROW_ID_FIELD_NAME in leaf.type.names: + return leaf.field(ROW_ID_FIELD_NAME).to_pylist() + except KeyError: + pass + + # blob_column is the source name; aliased projections use the output name in hits. + row_ids = _find_row_id_in_any_column(hits) + if row_ids is not None: + return row_ids + + raise ValueError( + f"query result has no '_rowid' column and no '{ROW_ID_FIELD_NAME}' " + f"field on blob column '{blob_column}'. Pass fresh blob query " + "results, call .with_row_id(True), or pass a list of row ids." + ) + + +def _find_row_id_in_any_column(tbl: pa.Table) -> Optional[list[int]]: + for name in tbl.column_names: + column = tbl.column(name) + if isinstance(column, pa.ChunkedArray): + column = column.combine_chunks() + row_ids = _find_row_id_in_struct(column) + if row_ids is not None: + return row_ids + return None + + +def _find_row_id_in_struct(array: pa.Array) -> Optional[list[int]]: + if not pa.types.is_struct(array.type): + return None + if ROW_ID_FIELD_NAME in array.type.names: + return array.field(ROW_ID_FIELD_NAME).to_pylist() + for i in range(array.type.num_fields): + row_ids = _find_row_id_in_struct(array.field(i)) + if row_ids is not None: + return row_ids + return None + + +def _iter_projection_pairs( + projection: QueryProjectionSpec, +) -> Iterable[tuple[str, str]]: + if isinstance(projection, dict): + for name, expr in projection.items(): + if isinstance(expr, str): + yield name, expr + elif isinstance(expr, Expr): + yield name, expr.to_sql() + return + for column in projection: + if isinstance(column, str): + yield column, column + elif isinstance(column, tuple) and len(column) == 2: + name, expr = column + if isinstance(expr, str): + yield name, expr + elif isinstance(expr, Expr): + yield name, expr.to_sql() + + +def _set_blob_column(tbl: pa.Table, output_name: str, blobs: pa.Array) -> pa.Table: + index = tbl.schema.get_field_index(output_name) + return tbl.set_column(index, pa.field(output_name, blobs.type), [blobs]) + + +def _embed_row_id_in_column(tbl: pa.Table, path: str, row_ids: pa.Array) -> pa.Table: + def add_row_id(children: list, child_fields: list) -> None: + children.append(row_ids) + child_fields.append(pa.field(ROW_ID_FIELD_NAME, pa.uint64(), nullable=False)) + + return _transform_struct_column(tbl, path, add_row_id) + + +def strip_auto_row_ids(tbl: pa.Table, blob_paths: Iterable[str]) -> pa.Table: + """Remove any `_lance_row_id` field embedded in blob descriptor structs. + + For read-only descriptor views (`blob_mode="descriptions"`) that never + fetch bytes, so have no use for the row id. + """ + + def drop_row_id(children: list, child_fields: list) -> None: + for i, field in enumerate(child_fields): + if field.name == ROW_ID_FIELD_NAME: + del children[i], child_fields[i] + return + + for path in blob_paths: + if path.split(".")[0] not in tbl.column_names: + continue + tbl = _transform_struct_column(tbl, path, drop_row_id) + return tbl + + +def _transform_struct_column( + tbl: pa.Table, path: str, leaf_transform: Callable[[list, list], None] +) -> pa.Table: + top_name, *rest = path.split(".") + top_index = tbl.schema.get_field_index(top_name) + top_field = tbl.schema.field(top_index) + top_array = tbl.column(top_name) + if isinstance(top_array, pa.ChunkedArray): + top_array = top_array.combine_chunks() + + new_array, new_field = _rebuild_struct(top_array, top_field, rest, leaf_transform) + return tbl.set_column(top_index, new_field, new_array) + + +def _rebuild_struct( + struct_array: pa.StructArray, + struct_field: pa.Field, + remaining_path: list[str], + leaf_transform: Callable[[list, list], None], +) -> tuple[pa.StructArray, pa.Field]: + null_mask = struct_array.is_null() + if not remaining_path: + children = [struct_array.field(i) for i in range(struct_array.type.num_fields)] + child_fields = list(struct_array.type) + leaf_transform(children, child_fields) + new_array = pa.StructArray.from_arrays( + children, fields=child_fields, mask=null_mask + ) + else: + child_name = remaining_path[0] + child_index = struct_array.type.get_field_index(child_name) + child_array = struct_array.field(child_index) + child_field = struct_array.type.field(child_index) + new_child_array, new_child_field = _rebuild_struct( + child_array, child_field, remaining_path[1:], leaf_transform + ) + + children = [] + child_fields = [] + for i in range(struct_array.type.num_fields): + field = struct_array.type.field(i) + if field.name == child_name: + children.append(new_child_array) + child_fields.append(new_child_field) + else: + children.append(struct_array.field(i)) + child_fields.append(field) + new_array = pa.StructArray.from_arrays( + children, fields=child_fields, mask=null_mask + ) + + new_field = pa.field( + struct_field.name, + new_array.type, + nullable=struct_field.nullable, + metadata=struct_field.metadata, + ) + return new_array, new_field + + +def _leaf_struct_column(tbl: pa.Table, path: str) -> pa.StructArray: + parts = path.split(".") + column = tbl.column(parts[0]) + if isinstance(column, pa.ChunkedArray): + column = column.combine_chunks() + for part in parts[1:]: + column = column.field(part) + return column + + +def _normalize_blob_row_ids( + row_ids: Union[list[int], pa.Table], blob_column: str +) -> list[int]: + if isinstance(row_ids, pa.Table): + return read_row_ids_from_hits(row_ids, blob_column) + if isinstance(row_ids, (pa.Array, pa.ChunkedArray)): + raise ValueError( + "pass a query table with _rowid, not a column array " + "(use fetch_blobs('image', hits), not fetch_blobs('image', hits['image']))" + ) + return list(row_ids) + + +def _wrap_blob_files(handles: Iterable[object]) -> list[Optional[BlobFile]]: + return [BlobFile(handle) if handle is not None else None for handle in handles] diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 18757ebfb..390235d8c 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -181,6 +181,17 @@ class Connection(object): self, ) -> Dict[str, Any]: ... +class BlobFile: + async def read(self) -> bytes: ... + def read_bytes(self) -> bytes: ... + def close(self) -> None: ... + def is_closed(self) -> bool: ... + def seek(self, position: int) -> None: ... + def tell(self) -> int: ... + def size(self) -> int: ... + def read_range(self, offset: int, length: int) -> bytes: ... + def read_up_to(self, length: int) -> bytes: ... + class Table: def name(self) -> str: ... def __repr__(self) -> str: ... @@ -258,6 +269,13 @@ class Table: def query(self) -> Query: ... def take_offsets(self, offsets: list[int]) -> TakeQuery: ... def take_row_ids(self, row_ids: list[int]) -> TakeQuery: ... + async def blob_columns(self) -> list[str]: ... + async def fetch_blobs( + self, column: str, row_ids: list[int] + ) -> pa.LargeBinaryArray: ... + async def fetch_blob_files( + self, column: str, row_ids: list[int] + ) -> list[Optional[BlobFile]]: ... def vector_search(self) -> VectorQuery: ... class Tags: diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 748f88af5..97be089c6 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -15,10 +15,12 @@ from typing import ( List, Literal, Optional, + Protocol, Tuple, Type, TypeVar, Union, + runtime_checkable, ) import deprecation @@ -39,15 +41,21 @@ from .expr import Expr from .rerankers.base import Reranker from .rerankers.rrf import RRFReranker from .rerankers.util import check_reranker_result +from .schema import is_blob_like_field, schema_has_blob_field from .util import flatten_columns - -BlobMode = Literal["lazy", "bytes", "descriptions"] - -_BLOB_MODE_TO_HANDLING = { - "lazy": "blobs_descriptions", - "bytes": "all_binary", - "descriptions": "blobs_descriptions", -} +from ._blob import ( + BLOB_MODE_TO_HANDLING, + FetchBlobsAsync, + FetchBlobsSync, + blob_auto_row_id_for_scan, + blob_v2_projection_sources, + finalize_blob_query_table, + replace_v2_blob_columns_with_bytes, + replace_v2_blob_columns_with_bytes_sync, + supports_blob_auto_row_id, + validate_blob_mode, +) +from .types import BlobMode, QueryProjection if TYPE_CHECKING: import sys @@ -73,25 +81,22 @@ if TYPE_CHECKING: T = TypeVar("T", bound="LanceModel") -def _validate_blob_mode(blob_mode: BlobMode) -> None: - if blob_mode not in _BLOB_MODE_TO_HANDLING: - modes = ", ".join(repr(mode) for mode in _BLOB_MODE_TO_HANDLING) - raise ValueError(f"blob_mode must be one of {modes}, got {blob_mode!r}") +@runtime_checkable +class _LanceScanner(Protocol): + projected_schema: pa.Schema | None + schema: pa.Schema | None + def to_pandas(self, blob_mode: BlobMode | None = ..., **kwargs) -> pd.DataFrame: ... -def _field_is_blob(field: pa.Field) -> bool: - metadata = field.metadata or {} - return metadata.get(b"lance-encoding:blob") == b"true" or ( - metadata.get("lance-encoding:blob") == "true" - ) + def to_pyarrow(self): ... + def to_table(self) -> pa.Table: ... -def _schema_has_blob_field(schema: pa.Schema) -> bool: - return any(_field_is_blob(field) for field in schema) + def to_reader(self): ... def _blob_mode_requires_native_pandas(blob_mode: BlobMode, schema: pa.Schema) -> bool: - return blob_mode in _BLOB_MODE_TO_HANDLING and _schema_has_blob_field(schema) + return blob_mode in BLOB_MODE_TO_HANDLING and schema_has_blob_field(schema) def _unsupported_blob_pandas_error(reason: str) -> RuntimeError: @@ -140,13 +145,7 @@ def _combine_where( return f"({existing_sql}) AND ({new_sql})" -def _projection_to_scanner_kwargs( - columns: Optional[ - Union[ - List[str], List[Tuple[str, Union[str, Expr]]], Dict[str, Union[str, Expr]] - ] - ], -) -> Dict[str, Any]: +def _projection_to_scanner_kwargs(columns: QueryProjection) -> Dict[str, Any]: if columns is None: return {} if isinstance(columns, list): @@ -171,7 +170,11 @@ def _projection_to_scanner_kwargs( def _scanner_kwargs_for_query( - query: Query, blob_mode: BlobMode, dataset: Optional[Any] = None + query: Query, + blob_mode: BlobMode, + dataset: Optional[Any] = None, + *, + with_row_id: Optional[bool] = None, ) -> Dict[str, Any]: fragments = _scanner_fragments_for_query(query, dataset) kwargs = { @@ -179,10 +182,10 @@ def _scanner_kwargs_for_query( "filter": _filter_to_sql(query.filter), "limit": query.limit, "offset": query.offset, - "with_row_id": query.with_row_id, + "with_row_id": with_row_id if with_row_id is not None else query.with_row_id, "with_row_address": query.with_row_address, "fast_search": query.fast_search, - "blob_handling": _BLOB_MODE_TO_HANDLING[blob_mode], + "blob_handling": BLOB_MODE_TO_HANDLING[blob_mode], "fragments": fragments, } return {key: value for key, value in kwargs.items() if value is not None} @@ -215,11 +218,11 @@ def _scanner_fragments_for_query(query: Query, dataset: Optional[Any]) -> Option def _ensure_lazy_blob_frame( df: "pd.DataFrame", schema: pa.Schema, blob_mode: BlobMode ) -> "pd.DataFrame": - if blob_mode != "lazy" or not _schema_has_blob_field(schema) or len(df) == 0: + if blob_mode != "lazy" or not schema_has_blob_field(schema) or len(df) == 0: return df for field in schema: - if not _field_is_blob(field) or field.name not in df.columns: + if not is_blob_like_field(field) or field.name not in df.columns: continue value = df[field.name].iloc[0] if value is not None and not hasattr(value, "readall"): @@ -229,7 +232,7 @@ def _ensure_lazy_blob_frame( return df -def _scanner_to_table(scanner: Any) -> pa.Table: +def _scanner_to_table(scanner: _LanceScanner) -> pa.Table: if hasattr(scanner, "to_pyarrow"): reader = scanner.to_pyarrow() return reader.read_all() @@ -239,7 +242,9 @@ def _scanner_to_table(scanner: Any) -> pa.Table: return reader.read_all() -def _scanner_to_pandas(scanner: Any, blob_mode: BlobMode, **kwargs) -> "pd.DataFrame": +def _scanner_to_pandas( + scanner: _LanceScanner, blob_mode: BlobMode, **kwargs +) -> pd.DataFrame: schema = getattr(scanner, "projected_schema", None) if schema is None: schema = getattr(scanner, "schema", None) @@ -260,13 +265,71 @@ def _scanner_to_pandas(scanner: Any, blob_mode: BlobMode, **kwargs) -> "pd.DataF return df tbl = _scanner_to_table(scanner) - if blob_mode == "lazy" and _schema_has_blob_field(tbl.schema): + if blob_mode == "lazy" and schema_has_blob_field(tbl.schema): raise _unsupported_blob_pandas_error( "the Lance scanner does not expose to_pandas" ) return tbl.to_pandas(**kwargs) +def _finish_plain_scan_pandas( + scanner: _LanceScanner, + *, + blob_mode: BlobMode, + blob_sources: dict[str, str], + fetch_blobs: FetchBlobsSync, + strip_auto_row_id: bool, + flatten: Optional[Union[int, bool]], + **kwargs, +) -> pd.DataFrame: + if blob_sources: + tbl = _scanner_to_table(scanner) + tbl = replace_v2_blob_columns_with_bytes_sync(tbl, blob_sources, fetch_blobs) + if strip_auto_row_id and "_rowid" in tbl.column_names: + tbl = tbl.drop_columns(["_rowid"]) + if flatten is not None: + tbl = flatten_columns(tbl, flatten) + return tbl.to_pandas(**kwargs) + if flatten is not None: + tbl = flatten_columns(_scanner_to_table(scanner), flatten) + if strip_auto_row_id and "_rowid" in tbl.column_names: + tbl = tbl.drop_columns(["_rowid"]) + return tbl.to_pandas(**kwargs) + df = _scanner_to_pandas(scanner, blob_mode, **kwargs) + if strip_auto_row_id and "_rowid" in df.columns: + return df.drop(columns=["_rowid"]) + return df + + +async def _finish_plain_scan_pandas_async( + scanner: _LanceScanner, + *, + blob_mode: BlobMode, + blob_sources: dict[str, str], + fetch_blobs: FetchBlobsAsync, + strip_auto_row_id: bool, + flatten: Optional[Union[int, bool]], + **kwargs, +) -> pd.DataFrame: + if blob_sources: + tbl = _scanner_to_table(scanner) + tbl = await replace_v2_blob_columns_with_bytes(tbl, blob_sources, fetch_blobs) + if strip_auto_row_id and "_rowid" in tbl.column_names: + tbl = tbl.drop_columns(["_rowid"]) + if flatten is not None: + tbl = flatten_columns(tbl, flatten) + return tbl.to_pandas(**kwargs) + if flatten is not None: + tbl = flatten_columns(_scanner_to_table(scanner), flatten) + if strip_auto_row_id and "_rowid" in tbl.column_names: + tbl = tbl.drop_columns(["_rowid"]) + return tbl.to_pandas(**kwargs) + df = _scanner_to_pandas(scanner, blob_mode, **kwargs) + if strip_auto_row_id and "_rowid" in df.columns: + return df.drop(columns=["_rowid"]) + return df + + # Pydantic validation function for vector queries def ensure_vector_query( val: Any, @@ -674,7 +737,7 @@ class Query(pydantic.BaseModel): distance_type: Optional[str] = None # which columns to return in the results (dict values may be str or Expr) - columns: Optional[Union[List[str], Dict[str, Union[str, Expr]]]] = None + columns: QueryProjection = None # minimum number of IVF partitions to search # @@ -958,7 +1021,7 @@ class LanceQueryBuilder(ABC): Forwarded to pyarrow.Table.to_pandas after query execution and optional flattening. """ - _validate_blob_mode(blob_mode) + validate_blob_mode(blob_mode) output_schema = getattr(self, "output_schema", None) if output_schema is not None: schema = output_schema() @@ -1017,6 +1080,11 @@ class LanceQueryBuilder(ABC): Execute the query and return the results as a pyarrow [RecordBatchReader](https://arrow.apache.org/docs/python/generated/pyarrow.RecordBatchReader.html) + For v2 blob projections, ``to_batches`` keeps the auto ``_rowid`` + column visible so batch consumers can call ``fetch_blobs``. Use + ``to_arrow``, ``to_list``, or ``to_pandas`` if you want LanceDB to hide + auto row ids in the final collected result. + Parameters ---------- batch_size: int @@ -1195,6 +1263,42 @@ class LanceQueryBuilder(ABC): self._with_row_id = with_row_id return self + def _user_requested_row_id(self) -> bool: + return self._with_row_id is True + + def _blob_auto_row_id_enabled(self) -> bool: + if not supports_blob_auto_row_id(self._table): + return False + return blob_auto_row_id_for_scan( + self._table, + self._table.schema, + self._columns, + with_row_id=self._with_row_id, + ) + + def _scan_needs_row_id(self) -> bool: + return self._user_requested_row_id() or self._blob_auto_row_id_enabled() + + def _query_for_scan(self) -> Query: + query = self.to_query_object() + if self._scan_needs_row_id(): + query.with_row_id = True + return query + + def _finalize_blob_query_table(self, tbl: pa.Table) -> pa.Table: + blob_auto_row_id = self._blob_auto_row_id_enabled() + blob_paths = ( + blob_v2_projection_sources(self._table.schema, self._columns).keys() + if blob_auto_row_id + else () + ) + return finalize_blob_query_table( + tbl, + user_requested_row_id=self._user_requested_row_id(), + blob_auto_row_id=blob_auto_row_id, + blob_paths=blob_paths, + ) + def with_row_address(self, with_row_address: bool = True) -> Self: """Set whether to return row addresses. @@ -1371,13 +1475,29 @@ class LanceQueryBuilder(ABC): return None dataset = self._table.to_lance() - scanner = dataset.scanner( - **_scanner_kwargs_for_query(query, blob_mode, dataset) + blob_auto_row_id = self._blob_auto_row_id_enabled() + blob_sources = ( + blob_v2_projection_sources(self._table.schema, query.columns) + if blob_mode == "bytes" + else {} + ) + scanner = dataset.scanner( + **_scanner_kwargs_for_query( + query, + "descriptions" if blob_sources else blob_mode, + dataset, + with_row_id=query.with_row_id or blob_auto_row_id or bool(blob_sources), + ) + ) + return _finish_plain_scan_pandas( + scanner, + blob_mode=blob_mode, + blob_sources=blob_sources, + fetch_blobs=self._table.fetch_blobs, + strip_auto_row_id=blob_auto_row_id, + flatten=flatten, + **kwargs, ) - if flatten is not None: - tbl = flatten_columns(_scanner_to_table(scanner), flatten) - return tbl.to_pandas(**kwargs) - return _scanner_to_pandas(scanner, blob_mode, **kwargs) @abstractmethod def to_query_object(self) -> Query: @@ -1625,7 +1745,9 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): The maximum time to wait for the query to complete. If None, wait indefinitely. """ - return self.to_batches(timeout=timeout).read_all() + return self._finalize_blob_query_table( + self.to_batches(timeout=timeout).read_all() + ) def to_query_object(self) -> Query: """ @@ -1685,7 +1807,7 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): vector = self._query if isinstance(self._query, list) else self._query.tolist() if isinstance(vector[0], np.ndarray): vector = [v.tolist() for v in vector] - query = self.to_query_object() + query = self._query_for_scan() result_set = self._table._execute_query( query, batch_size=batch_size, timeout=timeout ) @@ -1891,13 +2013,13 @@ class LanceFtsQueryBuilder(LanceQueryBuilder): query, PhraseQuery ): raise TypeError("Please use PhraseQuery for phrase queries.") - query = self.to_query_object() + query = self._query_for_scan() results = self._table._execute_query(query, timeout=timeout) results = results.read_all() if self._reranker is not None: results = self._reranker.rerank_fts(self._query, results) check_reranker_result(results) - return results + return self._finalize_blob_query_table(results) def to_batches( self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None @@ -1925,7 +2047,9 @@ class LanceFtsQueryBuilder(LanceQueryBuilder): class LanceEmptyQueryBuilder(LanceQueryBuilder): def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table: - return self.to_batches(timeout=timeout).read_all() + return self._finalize_blob_query_table( + self.to_batches(timeout=timeout).read_all() + ) def to_query_object(self) -> Query: return Query( @@ -1947,7 +2071,7 @@ class LanceEmptyQueryBuilder(LanceQueryBuilder): def to_batches( self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None ) -> pa.RecordBatchReader: - query = self.to_query_object() + query = self._query_for_scan() return self._table._execute_query(query, batch_size=batch_size, timeout=timeout) def rerank(self, reranker: Reranker) -> LanceEmptyQueryBuilder: @@ -2051,15 +2175,25 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): fts_results = fts_future.result() vector_results = vector_future.result() - return self._combine_hybrid_results( + results = self._combine_hybrid_results( fts_results=fts_results, vector_results=vector_results, norm=self._norm, fts_query=self._fts_query._query, reranker=self._reranker, limit=self._limit, - with_row_ids=self._with_row_id, + with_row_ids=True, ) + return self._finish_hybrid_results(results) + + def _finish_hybrid_results(self, results: pa.Table) -> pa.Table: + if self._user_requested_row_id(): + return results + if self._blob_auto_row_id_enabled(): + return self._finalize_blob_query_table(results) + if "_rowid" in results.column_names: + return results.drop(["_rowid"]) + return results @staticmethod def _combine_hybrid_results( @@ -2530,6 +2664,9 @@ class AsyncQueryBase(object): self._with_row_address = None self._fragments = None self._fragment_ids = None + self._with_row_id = None + self._blob_auto_row_id = False + self._blob_paths: tuple[str, ...] = () def to_query_object(self) -> Query: """ @@ -2539,11 +2676,46 @@ class AsyncQueryBase(object): python and more easily serializable. """ query = Query.from_inner(self._inner.to_query_request()) + query.with_row_id = self._user_requested_row_id() query.with_row_address = self._with_row_address query.fragments = self._fragments query.fragment_ids = self._fragment_ids return query + def _user_requested_row_id(self) -> bool: + return self._with_row_id is True + + def _blob_auto_row_id_enabled(self) -> bool: + return self._blob_auto_row_id + + def _finalize_blob_query_table(self, tbl: pa.Table) -> pa.Table: + return finalize_blob_query_table( + tbl, + user_requested_row_id=self._user_requested_row_id(), + blob_auto_row_id=self._blob_auto_row_id_enabled(), + blob_paths=self._blob_paths, + ) + + async def _maybe_add_blob_row_id(self) -> None: + if self._table is None or not supports_blob_auto_row_id(self._table): + self._blob_auto_row_id = False + self._blob_paths = () + return + + req = self._inner.to_query_request() + schema = await self._table.schema() + self._blob_auto_row_id = blob_auto_row_id_for_scan( + self._table, + schema, + req.select, + with_row_id=self._with_row_id, + ) + if not self._blob_auto_row_id: + self._blob_paths = () + return + self._blob_paths = tuple(blob_v2_projection_sources(schema, req.select).keys()) + self._inner.with_row_id() + def select(self, columns: Union[List[str], dict[str, str]]) -> Self: """ Return only the specified columns. @@ -2596,6 +2768,7 @@ class AsyncQueryBase(object): """ Include the _rowid column in the results. """ + self._with_row_id = True self._inner.with_row_id() return self @@ -2642,6 +2815,7 @@ class AsyncQueryBase(object): If not specified, no timeout is applied. If the query does not complete within the specified time, an error will be raised. """ + await self._maybe_add_blob_row_id() return AsyncRecordBatchReader( await self._inner.execute( max_batch_length=max_batch_length, timeout=timeout @@ -2672,8 +2846,8 @@ class AsyncQueryBase(object): complete within the specified time, an error will be raised. """ batch_iter = await self.to_batches(timeout=timeout) - return pa.Table.from_batches( - await batch_iter.read_all(), schema=batch_iter.schema + return self._finalize_blob_query_table( + pa.Table.from_batches(await batch_iter.read_all(), schema=batch_iter.schema) ) async def to_list(self, timeout: Optional[timedelta] = None) -> List[dict]: @@ -2740,7 +2914,7 @@ class AsyncQueryBase(object): Forwarded to pyarrow.Table.to_pandas after query execution and optional flattening. """ - _validate_blob_mode(blob_mode) + validate_blob_mode(blob_mode) if hasattr(self._inner, "output_schema"): schema = await self.output_schema() if _blob_mode_requires_native_pandas(blob_mode, schema): @@ -2781,14 +2955,36 @@ class AsyncQueryBase(object): if not _query_is_plain_scan(query): return None + schema = await self._table.schema() + blob_auto_row_id = blob_auto_row_id_for_scan( + self._table, + schema, + query.columns, + with_row_id=self._with_row_id, + ) + blob_sources = ( + blob_v2_projection_sources(schema, query.columns) + if blob_mode == "bytes" + else {} + ) dataset = await self._table._to_lance() scanner = dataset.scanner( - **_scanner_kwargs_for_query(query, blob_mode, dataset) + **_scanner_kwargs_for_query( + query, + "descriptions" if blob_sources else blob_mode, + dataset, + with_row_id=query.with_row_id or blob_auto_row_id or bool(blob_sources), + ) + ) + return await _finish_plain_scan_pandas_async( + scanner, + blob_mode=blob_mode, + blob_sources=blob_sources, + fetch_blobs=self._table.fetch_blobs, + strip_auto_row_id=blob_auto_row_id, + flatten=flatten, + **kwargs, ) - if flatten is not None: - tbl = flatten_columns(_scanner_to_table(scanner), flatten) - return tbl.to_pandas(**kwargs) - return _scanner_to_pandas(scanner, blob_mode, **kwargs) async def to_polars( self, @@ -3573,9 +3769,24 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase): fts_query = AsyncFTSQuery(self._inner.to_fts_query(), self._table) vec_query = AsyncVectorQuery(self._inner.to_vector_query(), self._table) - # save the row ID choice that was made on the query builder and force it - # to actually fetch the row ids because we need this for reranking - with_row_ids = self._inner.get_with_row_id() + req = fts_query._inner.to_query_request() + blob_auto_row_id = False + blob_paths: tuple[str, ...] = () + if self._table is not None and supports_blob_auto_row_id(self._table): + schema = await self._table.schema() + blob_auto_row_id = blob_auto_row_id_for_scan( + self._table, + schema, + req.select, + with_row_id=self._with_row_id, + ) + if blob_auto_row_id: + blob_paths = tuple( + blob_v2_projection_sources(schema, req.select).keys() + ) + self._blob_auto_row_id = blob_auto_row_id + self._blob_paths = blob_paths + fts_query.with_row_id() vec_query.with_row_id() @@ -3591,8 +3802,14 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase): fts_query=fts_query.get_query(), reranker=self._reranker, limit=self._inner.get_limit(), - with_row_ids=with_row_ids, + with_row_ids=True, ) + if ( + not self._user_requested_row_id() + and not blob_auto_row_id + and "_rowid" in result.column_names + ): + result = result.drop(["_rowid"]) return AsyncRecordBatchReader(result, max_batch_length=max_batch_length) diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index ecf3c417f..772e0055e 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -994,6 +994,19 @@ class RemoteTable(Table): "migrate_v2_manifest_paths() is not supported on the LanceDB Cloud" ) + def blob_columns(self) -> list[str]: + raise NotImplementedError( + "blob_columns() is not yet supported on the LanceDB Cloud" + ) + + def fetch_blobs(self, column: str, row_ids) -> pa.LargeBinaryArray: + raise NotImplementedError("fetch_blobs() is not supported on LanceDB Cloud") + + def fetch_blob_files(self, column: str, row_ids): + raise NotImplementedError( + "fetch_blob_files() is not supported on LanceDB Cloud" + ) + def head(self, n=5) -> pa.Table: """ Return the first `n` rows of the table. diff --git a/python/python/lancedb/schema.py b/python/python/lancedb/schema.py index e8e73f2e8..33adbbae3 100644 --- a/python/python/lancedb/schema.py +++ b/python/python/lancedb/schema.py @@ -2,10 +2,134 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors -"""Schema related utilities.""" +"""Schema helpers for Lance blob columns.""" import pyarrow as pa +_BLOB_EXTENSION_NAME = "lance.blob.v2" +_BLOB_V1_KEY = "lance-encoding:blob" +_ARROW_EXT_NAME_KEY = "ARROW:extension:name" + + +class BlobType(pa.ExtensionType): + """PyArrow extension type for a Lance blob v2 column. + + Queries return descriptors; call :meth:`~lancedb.table.Table.fetch_blob_files` + for lazy reads or :meth:`~lancedb.table.Table.fetch_blobs` for eager bytes. + """ + + def __init__(self) -> None: + storage_type = pa.struct( + [ + pa.field("data", pa.large_binary(), nullable=True), + pa.field("uri", pa.utf8(), nullable=True), + pa.field("position", pa.uint64(), nullable=True), + pa.field("size", pa.uint64(), nullable=True), + ] + ) + super().__init__(storage_type, _BLOB_EXTENSION_NAME) + + def __arrow_ext_serialize__(self) -> bytes: + return b"" + + @classmethod + def __arrow_ext_deserialize__( + cls, storage_type: pa.DataType, serialized: bytes + ) -> "BlobType": + return cls() + + def __reduce__(self): + # Ensure pickle round-trips on older pyarrow (apache/arrow#35599). + return type(self).__arrow_ext_deserialize__, ( + self.storage_type, + self.__arrow_ext_serialize__(), + ) + + +try: + pa.register_extension_type(BlobType()) # type: ignore[arg-type] +except pa.ArrowKeyError: + pass + + +def _metadata_value(metadata: dict, key: str): + return metadata.get(key.encode()) or metadata.get(key) + + +def _metadata_marks_blob_v2(metadata: dict) -> bool: + if not metadata: + return False + + extension_name = _metadata_value(metadata, _ARROW_EXT_NAME_KEY) + return extension_name in (_BLOB_EXTENSION_NAME, _BLOB_EXTENSION_NAME.encode()) + + +def _metadata_marks_legacy_blob(metadata: dict) -> bool: + if not metadata: + return False + + return _metadata_value(metadata, _BLOB_V1_KEY) in ("true", b"true") + + +def is_blob_v2_field(field: pa.Field) -> bool: + """Return True if `field` declares a blob v2 extension column.""" + field_type = field.type + if ( + isinstance(field_type, pa.ExtensionType) + and field_type.extension_name == _BLOB_EXTENSION_NAME + ): + return True + return _metadata_marks_blob_v2(field.metadata or {}) + + +def is_blob_like_field(field: pa.Field) -> bool: + """Blob detection for ``to_pandas(blob_mode=...)`` and scanner paths only. + + Matches v2 extension fields on table schema, legacy ``lance-encoding:blob`` + storage columns, and v2 query descriptor fields (the engine tags those with + the same metadata). Not used for fetch or auto ``_rowid``. + """ + return is_blob_v2_field(field) or _metadata_marks_legacy_blob(field.metadata or {}) + + +def _collect_blob_paths(schema: pa.Schema, is_blob) -> list[str]: + paths: list[str] = [] + + def walk(fields, prefix: str) -> None: + for field in fields: + path = f"{prefix}.{field.name}" if prefix else field.name + if is_blob(field): + paths.append(path) + elif pa.types.is_struct(field.type): + walk(field.type, path) + elif ( + pa.types.is_list(field.type) + or pa.types.is_large_list(field.type) + or pa.types.is_fixed_size_list(field.type) + ): + walk([field.type.value_field], path) + + walk(schema, "") + return paths + + +def blob_column_paths(schema: pa.Schema) -> list[str]: + """Dotted paths of blob-like columns (v2 extension or legacy metadata).""" + return _collect_blob_paths(schema, is_blob_like_field) + + +def blob_v2_column_paths(schema: pa.Schema) -> list[str]: + return _collect_blob_paths(schema, is_blob_v2_field) + + +def schema_has_blob_field(schema: pa.Schema) -> bool: + return bool(blob_column_paths(schema)) + + +def blob(name: str, nullable: bool = True) -> pa.Field: + """Create a Lance blob v2 column field.""" + return pa.field(name, BlobType(), nullable=nullable) + def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataType: """A help function to create a vector type. diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index bc11a4ce7..98f14dbfe 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -29,6 +29,14 @@ from urllib.parse import urlparse from lancedb.scannable import _register_optional_converters, to_scannable from . import __version__ +from ._blob import ( + BlobFile, + _normalize_blob_row_ids, + _wrap_blob_files, + strip_auto_row_ids, + validate_blob_mode, +) +from .types import BlobMode from lancedb.arrow import peek_reader from lancedb.background_loop import LOOP, embedding_executor from .dependencies import ( @@ -88,10 +96,7 @@ from .util import ( value_to_sql, ) from .index import lang_mapping - -BlobMode = Literal["lazy", "bytes", "descriptions"] - -_VALID_BLOB_MODES = ("lazy", "bytes", "descriptions") +from .schema import blob_v2_column_paths, schema_has_blob_field def _should_push_down_query_table( @@ -100,23 +105,6 @@ def _should_push_down_query_table( return namespace_client is not None and "QueryTable" in pushdown_operations -def _validate_blob_mode(blob_mode: BlobMode) -> None: - if blob_mode not in _VALID_BLOB_MODES: - modes = ", ".join(repr(mode) for mode in _VALID_BLOB_MODES) - raise ValueError(f"blob_mode must be one of {modes}, got {blob_mode!r}") - - -def _field_is_blob(field: pa.Field) -> bool: - metadata = field.metadata or {} - return metadata.get(b"lance-encoding:blob") == b"true" or ( - metadata.get("lance-encoding:blob") == "true" - ) - - -def _schema_has_blob_field(schema: pa.Schema) -> bool: - return any(_field_is_blob(field) for field in schema) - - _MODEL_BACKED_TOKENIZER_PREFIXES = ("jieba", "lindera") _MODEL_BACKED_TOKENIZER_ERRORS = ( "unknown base tokenizer", @@ -1523,6 +1511,31 @@ class Table(ABC): A query object that can be executed to get the rows. """ + @abstractmethod + def blob_columns(self) -> list[str]: + """Names of the blob v2 columns declared on this table.""" + + @abstractmethod + def fetch_blobs( + self, column: str, row_ids: Union[list[int], pa.Table] + ) -> pa.LargeBinaryArray: + """Materialize full blob bytes for ``column`` at the given rows. + + Convenience for small payloads. For large values use + :meth:`fetch_blob_files`. + """ + + @abstractmethod + def fetch_blob_files( + self, column: str, row_ids: Union[list[int], pa.Table] + ) -> "list[Optional[BlobFile]]": + """Open lazy, seekable :class:`~lancedb._blob.BlobFile` handles. + + Prefer this over :meth:`fetch_blobs` for large payloads. ``row_ids`` is + a ``list[int]`` or query ``pyarrow.Table`` with ``_rowid`` (or stashed + row-id metadata). Null rows are ``None``. Local tables only. + """ + @abstractmethod def _execute_query( self, @@ -2204,6 +2217,19 @@ class LanceTable(Table): def take_row_ids(self, row_ids: list[int]) -> LanceTakeQueryBuilder: return LanceTakeQueryBuilder(self._table.take_row_ids(row_ids)) + def blob_columns(self) -> list[str]: + return LOOP.run(self._table.blob_columns()) + + def fetch_blobs( + self, column: str, row_ids: Union[list[int], pa.Table] + ) -> pa.LargeBinaryArray: + return LOOP.run(self._table.fetch_blobs(column, row_ids)) + + def fetch_blob_files( + self, column: str, row_ids: Union[list[int], pa.Table] + ) -> "list[Optional[BlobFile]]": + return LOOP.run(self._table.fetch_blob_files(column, row_ids)) + @property def tags(self) -> Tags: """Tag management for the table. @@ -2399,9 +2425,14 @@ class LanceTable(Table): ------- pd.DataFrame """ - _validate_blob_mode(blob_mode) - if blob_mode == "descriptions" or not _schema_has_blob_field(self.schema): - return self.to_arrow().to_pandas(**kwargs) + validate_blob_mode(blob_mode) + if blob_mode == "descriptions" or not schema_has_blob_field(self.schema): + arrow_tbl = self.to_arrow() + if blob_mode == "descriptions": + arrow_tbl = strip_auto_row_ids( + arrow_tbl, blob_v2_column_paths(self.schema) + ) + return arrow_tbl.to_pandas(**kwargs) if ( blob_mode == "lazy" @@ -2410,6 +2441,9 @@ class LanceTable(Table): ): return self.to_arrow().to_pandas(**kwargs) + if blob_mode == "bytes" and blob_v2_column_paths(self.schema): + return self.search().to_pandas(blob_mode=blob_mode, **kwargs) + return self.to_lance().to_pandas(blob_mode=blob_mode, **kwargs) def to_arrow(self) -> pa.Table: @@ -4539,14 +4573,18 @@ class AsyncTable: ------- pd.DataFrame """ - _validate_blob_mode(blob_mode) - if blob_mode == "descriptions" or not _schema_has_blob_field( - await self.schema() - ): - return (await self.to_arrow()).to_pandas(**kwargs) + validate_blob_mode(blob_mode) + schema = await self.schema() + if blob_mode == "descriptions" or not schema_has_blob_field(schema): + arrow_tbl = await self.to_arrow() + if blob_mode == "descriptions": + arrow_tbl = strip_auto_row_ids(arrow_tbl, blob_v2_column_paths(schema)) + return arrow_tbl.to_pandas(**kwargs) if blob_mode == "lazy" and get_uri_scheme(await self.uri()) == "memory": return (await self.to_arrow()).to_pandas(**kwargs) + if blob_mode == "bytes" and blob_v2_column_paths(schema): + return await self.query().to_pandas(blob_mode=blob_mode, **kwargs) return (await self._to_lance()).to_pandas(blob_mode=blob_mode, **kwargs) async def to_arrow(self) -> pa.Table: @@ -5647,6 +5685,24 @@ class AsyncTable: """ return AsyncTakeQuery(self._inner.take_row_ids(row_ids), self) + async def blob_columns(self) -> list[str]: + return await self._inner.blob_columns() + + async def fetch_blobs( + self, column: str, row_ids: Union[list[int], pa.Table] + ) -> pa.LargeBinaryArray: + return await self._inner.fetch_blobs( + column, _normalize_blob_row_ids(row_ids, column) + ) + + async def fetch_blob_files( + self, column: str, row_ids: Union[list[int], pa.Table] + ) -> "list[Optional[BlobFile]]": + handles = await self._inner.fetch_blob_files( + column, _normalize_blob_row_ids(row_ids, column) + ) + return _wrap_blob_files(handles) + @property def tags(self) -> AsyncTags: """Tag management for the dataset. diff --git a/python/python/lancedb/types.py b/python/python/lancedb/types.py index 8bab57e0e..d2edd318a 100644 --- a/python/python/lancedb/types.py +++ b/python/python/lancedb/types.py @@ -1,11 +1,24 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The LanceDB Authors -from typing import Literal +from __future__ import annotations + +from typing import Dict, List, Literal, Optional, Tuple, Union + +from .expr import Expr # Query type literals QueryType = Literal["vector", "fts", "hybrid", "auto"] +BlobMode = Literal["lazy", "bytes", "descriptions"] + +QueryProjectionSpec = Union[ + List[str], + List[Tuple[str, Union[str, Expr]]], + Dict[str, Union[str, Expr]], +] +QueryProjection = Optional[QueryProjectionSpec] + # Distance type literals DistanceType = Literal["l2", "cosine", "dot"] DistanceTypeWithHamming = Literal["l2", "cosine", "dot", "hamming"] diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py new file mode 100644 index 000000000..05af14298 --- /dev/null +++ b/python/python/tests/test_blob.py @@ -0,0 +1,562 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import io + +import pyarrow as pa +import pyarrow.compute as pc +import pytest + +import lancedb +from lancedb._blob import read_row_ids_from_hits, stash_auto_row_ids +from lancedb.index import FTS +from lancedb.schema import blob_column_paths, blob_v2_column_paths + + +def _blob_table(name, rows): + db = lancedb.connect("memory:///") + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = db.create_table(name, schema=schema) + table.add(rows) + return table + + +def _blob_array(name, values): + blob_type = lancedb.blob(name).type + storage_type = blob_type.storage_type + storage = pa.StructArray.from_arrays( + [ + pa.array(values, type=pa.large_binary()), + pa.array([None] * len(values), type=pa.string()), + pa.array([None] * len(values), type=pa.uint64()), + pa.array([None] * len(values), type=pa.uint64()), + ], + fields=list(storage_type), + ) + return pa.ExtensionArray.from_storage(blob_type, storage) + + +def _row_ids_by_id(table): + hits = table.search().with_row_id(True).limit(1000).to_arrow() + assert "_rowid" in hits.column_names + return dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist())) + + +def test_blob_factory_declares_v2_field(): + field = lancedb.blob("image") + assert isinstance(field.type, pa.ExtensionType) + assert field.type.extension_name == "lance.blob.v2" + + +def test_blob_v2_column_paths_include_list_children(): + schema = pa.schema( + [ + pa.field("id", pa.int64()), + pa.field("info", pa.struct([lancedb.blob("blob")])), + pa.field("images", pa.list_(lancedb.blob("image"))), + pa.field("large_images", pa.large_list(lancedb.blob("large_image"))), + pa.field( + "fixed_images", + pa.list_(lancedb.blob("fixed_image"), list_size=2), + ), + ] + ) + + assert blob_v2_column_paths(schema) == [ + "info.blob", + "images.image", + "large_images.large_image", + "fixed_images.fixed_image", + ] + + +def _legacy_v1_table(name): + db = lancedb.connect("memory:///") + schema = pa.schema( + [ + pa.field("id", pa.int64()), + pa.field( + "legacy", pa.large_binary(), metadata={"lance-encoding:blob": "true"} + ), + ] + ) + table = db.create_table(name, schema=schema) + table.add([{"id": 1, "legacy": b"old"}]) + return table + + +def test_blob_v2_column_paths_exclude_legacy_metadata(): + schema = pa.schema( + [ + pa.field("id", pa.int64()), + lancedb.blob("image"), + pa.field( + "legacy", pa.large_binary(), metadata={"lance-encoding:blob": "true"} + ), + ] + ) + assert blob_v2_column_paths(schema) == ["image"] + assert blob_column_paths(schema) == ["image", "legacy"] + + +def test_blob_v2_paths_match_blob_columns(): + table = _blob_table("paths_match", [{"id": 1, "image": b"x"}]) + assert blob_v2_column_paths(table.schema) == table.blob_columns() + + db = lancedb.connect("memory:///") + info = pa.StructArray.from_arrays( + [ + pa.array(["first"], type=pa.string()), + _blob_array("blob", [b"nested"]), + ], + names=["name", "blob"], + ) + data = pa.Table.from_arrays( + [pa.array([1], type=pa.int64()), info], + names=["id", "info"], + ) + nested = db.create_table("nested_paths", data=data) + assert blob_v2_column_paths(nested.schema) == nested.blob_columns() + + +def test_auto_row_id_stash_round_trip(): + table = _blob_table( + "stash_round_trip", + [{"id": 1, "image": b"alpha"}, {"id": 2, "image": b"beta"}], + ) + hits = table.search().with_row_id(True).limit(10).to_arrow() + row_ids = hits["_rowid"].to_pylist() + + stashed = stash_auto_row_ids(hits, ["image"]) + + assert "_rowid" not in stashed.column_names + assert stashed.schema.field("image").metadata == hits.schema.field("image").metadata + assert read_row_ids_from_hits(stashed, "image") == row_ids + + +def test_blob_query_omits_auto_row_id(): + table = _blob_table("rowid", [{"id": 1, "image": b"x"}]) + hits = table.search().limit(10).to_arrow() + assert "_rowid" not in hits.column_names + + +def test_blob_query_explicit_row_id_opt_in(): + table = _blob_table("explicit_rowid", [{"id": 1, "image": b"x"}]) + hits = table.search().with_row_id(True).limit(10).to_arrow() + assert "_rowid" in hits.column_names + + +def test_table_to_pandas_descriptions_mode_omits_row_id(): + table = _blob_table("descriptions_no_leak", [{"id": 1, "image": b"x"}]) + df = table.to_pandas(blob_mode="descriptions") + descriptor = df["image"].iloc[0] + assert "_lance_row_id" not in descriptor + assert set(descriptor.keys()) == {"kind", "position", "size", "blob_id", "blob_uri"} + + +@pytest.mark.asyncio +async def test_async_table_to_pandas_descriptions_mode_omits_row_id(): + db = await lancedb.connect_async("memory:///") + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = await db.create_table("descriptions_no_leak_async", schema=schema) + await table.add([{"id": 1, "image": b"x"}]) + df = await table.to_pandas(blob_mode="descriptions") + descriptor = df["image"].iloc[0] + assert "_lance_row_id" not in descriptor + assert set(descriptor.keys()) == {"kind", "position", "size", "blob_id", "blob_uri"} + + +def test_fetch_blobs_round_trip(): + table = _blob_table( + "round_trip", + [{"id": 1, "image": b"alpha"}, {"id": 2, "image": b"beta"}], + ) + by_id = _row_ids_by_id(table) + blobs = table.fetch_blobs("image", [by_id[1], by_id[2]]) + assert [blobs[0].as_py(), blobs[1].as_py()] == [b"alpha", b"beta"] + + +def test_fetch_blobs_accepts_query_result(): + table = _blob_table("from_result", [{"id": 1, "image": b"gamma"}]) + hits = table.search().limit(10).to_arrow() + assert "_rowid" not in hits.column_names + blobs = table.fetch_blobs("image", hits) + assert {blobs[i].as_py() for i in range(len(blobs))} == {b"gamma"} + + +def test_fetch_blobs_null_alignment(): + table = _blob_table( + "nulls", + [{"id": 1, "image": b"present"}, {"id": 2, "image": None}], + ) + by_id = _row_ids_by_id(table) + request = [by_id[1], by_id[2], by_id[1]] + blobs = table.fetch_blobs("image", request) + assert len(blobs) == len(request) + assert blobs[0].as_py() == b"present" + assert blobs[1].as_py() is None + assert blobs[2].as_py() == b"present" + + +def test_fetch_blobs_nested_path(): + db = lancedb.connect("memory:///") + info = pa.StructArray.from_arrays( + [ + pa.array(["first", "second"], type=pa.string()), + _blob_array("blob", [b"nested-alpha", b"nested-beta"]), + ], + names=["name", "blob"], + ) + data = pa.Table.from_arrays( + [pa.array([1, 2], type=pa.int64()), info], + names=["id", "info"], + ) + table = db.create_table("nested", data=data) + + by_id = _row_ids_by_id(table) + blobs = table.fetch_blobs("info.blob", [by_id[1], by_id[2]]) + assert [blobs[0].as_py(), blobs[1].as_py()] == [b"nested-alpha", b"nested-beta"] + + +def test_fetch_blob_files_lazy_read(): + payload = b"lazy-read" * 100 + table = _blob_table("lazy", [{"id": 1, "image": payload}]) + by_id = _row_ids_by_id(table) + handles = table.fetch_blob_files("image", [by_id[1]]) + assert len(handles) == 1 + assert handles[0].read() == payload + + +def test_fetch_blob_files_null_alignment(): + table = _blob_table( + "lazy_nulls", + [{"id": 1, "image": b"here"}, {"id": 2, "image": None}], + ) + by_id = _row_ids_by_id(table) + handles = table.fetch_blob_files("image", [by_id[2], by_id[1]]) + assert len(handles) == 2 + assert handles[0] is None + assert handles[1].read() == b"here" + + +def test_fetch_blobs_rejects_non_blob_column(): + table = _blob_table("reject", [{"id": 1, "image": b"x"}]) + with pytest.raises(ValueError, match="not a blob column"): + table.fetch_blobs("id", [0]) + + +def test_legacy_v1_query_omits_auto_row_id(): + table = _legacy_v1_table("legacy_v1") + hits = table.search().select(["legacy"]).limit(10).to_arrow() + assert "_rowid" not in hits.column_names + + +def test_fetch_blobs_rejects_legacy_v1_column(): + table = _legacy_v1_table("legacy_fetch") + with pytest.raises(ValueError, match="legacy blob column.*blob v2"): + table.fetch_blobs("legacy", [0]) + + +@pytest.mark.asyncio +async def test_async_fetch_blob_files_lazy_read(): + db = await lancedb.connect_async("memory:///") + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = await db.create_table("async_lazy", schema=schema) + payload = b"async-lazy" * 100 + await table.add([{"id": 1, "image": payload}]) + hits = ( + await table.query().select({"image_alias": "image"}).limit(10).to_arrow() + ).combine_chunks() + assert "_rowid" not in hits.column_names + handles = await table.fetch_blob_files("image", hits) + assert len(handles) == 1 + assert await handles[0].aread() == payload + + +def test_fetch_blobs_from_query_result_without_row_id_raises(): + table = _blob_table("no_rowid", [{"id": 1, "image": b"x"}]) + hits = table.search().select(["id"]).to_arrow() + assert "_rowid" not in hits.column_names + with pytest.raises(ValueError, match="_rowid"): + table.fetch_blobs("image", hits) + + +_HYBRID_BLOB_SCHEMA = pa.schema( + [ + pa.field("id", pa.int64()), + pa.field("text", pa.utf8()), + pa.field("vector", pa.list_(pa.float32(), list_size=2)), + lancedb.blob("image"), + ] +) +_HYBRID_BLOB_ROWS = [ + {"id": 1, "text": "hello alpha", "vector": [1.0, 0.0], "image": b"alpha"}, + {"id": 2, "text": "hello beta", "vector": [0.9, 0.1], "image": b"beta"}, + {"id": 3, "text": "other", "vector": [0.0, 1.0], "image": b"other"}, +] + + +def _hybrid_blob_table(db): + table = db.create_table("hybrid_blob_fetch", schema=_HYBRID_BLOB_SCHEMA) + table.add(_HYBRID_BLOB_ROWS) + table.create_index("text", config=FTS(with_position=False)) + return table + + +async def _hybrid_blob_table_async(db): + table = await db.create_table("hybrid_blob_fetch_async", schema=_HYBRID_BLOB_SCHEMA) + await table.add(_HYBRID_BLOB_ROWS) + await table.create_index("text", config=FTS(with_position=False)) + return table + + +def test_blob_v2_hybrid_fetch_blobs(): + table = _hybrid_blob_table(lancedb.connect("memory:///")) + hits = ( + table.search(query_type="hybrid") + .vector([1.0, 0.0]) + .text("hello") + .select(["id", "image"]) + .limit(2) + .to_arrow() + ) + + assert "_rowid" not in hits.column_names + assert "_lance_row_id" in hits.schema.field("image").type.names + blobs = table.fetch_blobs("image", hits) + assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"} + + +@pytest.mark.asyncio +async def test_blob_v2_hybrid_fetch_blobs_async(): + db = await lancedb.connect_async("memory:///hybrid_blob_fetch_async") + table = await _hybrid_blob_table_async(db) + hits = await ( + table.query() + .nearest_to([1.0, 0.0]) + .nearest_to_text("hello") + .select(["id", "image"]) + .limit(2) + .to_arrow() + ) + + assert "_rowid" not in hits.column_names + assert "_lance_row_id" in hits.schema.field("image").type.names + blobs = await table.fetch_blobs("image", hits) + assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"} + + +def test_blob_file_seek_read_and_read_range(): + payload = _identifiable_payload(1024) + table = _blob_table("seek_read", [{"id": 1, "image": payload}]) + by_id = _row_ids_by_id(table) + handle = table.fetch_blob_files("image", [by_id[1]])[0] + + assert handle.seek(100) == 100 + assert handle.read(16) == payload[100:116] + handle.seek(100) + assert handle.read_range(500, 8) == payload[500:508] + assert handle.tell() == 100 + + with pytest.raises(ValueError, match="whence"): + handle.seek(0, 99) + + +def test_fetch_blob_files_from_query_partial_read(): + payload = _identifiable_payload(65536) + table = _blob_table("query_partial", [{"id": 1, "image": payload}]) + hits = table.search().select(["id", "image"]).limit(1).to_arrow() + assert "_rowid" not in hits.column_names + + handle = table.fetch_blob_files("image", hits)[0] + assert handle.size() == 65536 + assert handle.read_range(0, 128) == payload[:128] + assert handle.tell() == 0 + assert handle.seek(40000) == 40000 + assert handle.read(16) == payload[40000:40016] + + +def test_blob_file_buffered_reader(): + payload = _identifiable_payload(4096) + table = _blob_table("buffered_reader", [{"id": 1, "image": payload}]) + hits = table.search().select(["id", "image"]).limit(1).to_arrow() + handle = table.fetch_blob_files("image", hits)[0] + reader = io.BufferedReader(handle) + assert reader.read(8) == payload[:8] + assert reader.read(8) == payload[8:16] + assert reader.read() == payload[16:] + + +def test_fetch_blob_files_cross_fragment_nulls_and_dups(): + db = lancedb.connect("memory:///") + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = db.create_table("cross_fragment", schema=schema) + table.add([{"id": 1, "image": b"alpha"}]) + table.add([{"id": 2, "image": None}, {"id": 3, "image": b"beta"}]) + + by_id = _row_ids_by_id(table) + request = [by_id[3], by_id[2], by_id[1], by_id[3]] + handles = table.fetch_blob_files("image", request) + assert len(handles) == 4 + assert handles[1] is None + assert handles[0].read() == b"beta" + assert handles[2].read() == b"alpha" + assert handles[3].seek(1) == 1 + assert handles[3].read() == b"eta" + + +def test_blob_file_pyav_decode_seek(tmp_path): + av = pytest.importorskip("av") + import fractions + + clip = tmp_path / "clip.mp4" + with av.open(str(clip), mode="w") as container: + stream = container.add_stream("mpeg4", rate=5) + stream.width, stream.height, stream.pix_fmt = 32, 32, "yuv420p" + stream.time_base = fractions.Fraction(1, 5) + for pts in range(5): + frame = av.VideoFrame(32, 32, "yuv420p") + frame.pts = pts + container.mux(stream.encode(frame)) + container.mux(stream.encode(None)) + + table = _blob_table("pyav", [{"id": 1, "image": clip.read_bytes()}]) + hits = table.search().select(["image"]).limit(1).to_arrow() + handle = table.fetch_blob_files("image", hits)[0] + + with av.open(handle) as container: + stream = container.streams.video[0] + container.seek(0) + assert next(container.decode(stream)) is not None + + +def test_blob_v2_hybrid_fetch_blob_files_seek(): + table = _hybrid_blob_table(lancedb.connect("memory:///")) + hits = ( + table.search(query_type="hybrid") + .vector([1.0, 0.0]) + .text("hello") + .select(["id", "image"]) + .limit(2) + .to_arrow() + ) + assert "_rowid" not in hits.column_names + + handles = table.fetch_blob_files("image", hits) + assert len(handles) == 2 + assert {handle.read_range(0, 2) for handle in handles} == {b"al", b"be"} + first = handles[0] + assert first.seek(1) == 1 + assert first.read(2) in {b"lp", b"et"} + + +def test_blob_file_header_sniff_from_search(): + payload = b"%PDF-1.7\n" + bytes(4096) + table = _blob_table("header_sniff", [{"id": 1, "image": payload}]) + hits = table.search().select(["id", "image"]).limit(1).to_arrow() + handle = table.fetch_blob_files("image", hits)[0] + assert handle.read_range(0, 4) == b"%PDF" + assert handle.tell() == 0 + + +def test_blob_file_multiple_handles_independent_cursors(): + table = _blob_table( + "multi_handle", + [{"id": 1, "image": b"first-payload"}, {"id": 2, "image": b"second-payload"}], + ) + by_id = _row_ids_by_id(table) + first, second = table.fetch_blob_files("image", [by_id[1], by_id[2]]) + assert first.seek(6) == 6 + assert second.tell() == 0 + assert first.read(7) == b"payload" + assert second.read(6) == b"second" + + +def test_fetch_blob_files_nested_path_seek(): + db = lancedb.connect("memory:///") + info = pa.StructArray.from_arrays( + [ + pa.array(["first", "second"], type=pa.string()), + _blob_array("blob", [b"nested-alpha", b"nested-beta"]), + ], + names=["name", "blob"], + ) + data = pa.Table.from_arrays( + [pa.array([1, 2], type=pa.int64()), info], + names=["id", "info"], + ) + table = db.create_table("nested_seek", data=data) + by_id = _row_ids_by_id(table) + handle = table.fetch_blob_files("info.blob", [by_id[2]])[0] + assert handle.seek(7) == 7 + assert handle.read() == b"beta" + + +def test_fetch_blobs_survives_sort_after_query(): + table = _blob_table( + "sort_survives", + [{"id": i, "image": f"payload-{i}".encode()} for i in range(5)], + ) + hits = table.search().select(["id", "image"]).to_arrow() + sort_idx = pc.sort_indices(hits["id"], sort_keys=[("id", "descending")]) + sorted_hits = hits.take(sort_idx) + + blobs = table.fetch_blobs("image", sorted_hits) + expected = [f"payload-{i}".encode() for i in sorted_hits["id"].to_pylist()] + assert [blobs[i].as_py() for i in range(len(blobs))] == expected + + +def test_fetch_blobs_survives_filter_and_sort_after_query(): + table = _blob_table( + "filter_sort_survives", + [{"id": i, "image": f"payload-{i}".encode()} for i in range(5)], + ) + hits = table.search().select(["id", "image"]).to_arrow() + filtered = hits.filter(pc.field("id") >= 2) + sort_idx = pc.sort_indices(filtered["id"], sort_keys=[("id", "descending")]) + filtered_sorted = filtered.take(sort_idx) + + blobs = table.fetch_blobs("image", filtered_sorted) + expected = [f"payload-{i}".encode() for i in filtered_sorted["id"].to_pylist()] + assert [blobs[i].as_py() for i in range(len(blobs))] == expected + + +def test_fetch_blob_files_survives_sort_after_query(): + table = _blob_table( + "lazy_sort_survives", + [{"id": i, "image": f"payload-{i}".encode()} for i in range(5)], + ) + hits = table.search().select(["id", "image"]).to_arrow() + sort_idx = pc.sort_indices(hits["id"], sort_keys=[("id", "descending")]) + sorted_hits = hits.take(sort_idx) + + handles = table.fetch_blob_files("image", sorted_hits) + expected = [f"payload-{i}".encode() for i in sorted_hits["id"].to_pylist()] + assert [handle.read() for handle in handles] == expected + + +def test_fetch_blobs_nested_path_survives_sort_after_query(): + db = lancedb.connect("memory:///") + values = [f"payload-{i}".encode() for i in range(4)] + info = pa.StructArray.from_arrays( + [pa.array(["row"] * 4, type=pa.string()), _blob_array("blob", values)], + names=["name", "blob"], + ) + data = pa.Table.from_arrays( + [pa.array(range(4), type=pa.int64()), info], + names=["id", "info"], + ) + table = db.create_table("nested_sort_survives", data=data) + + hits = table.search().to_arrow() + sort_idx = pc.sort_indices(hits["id"], sort_keys=[("id", "descending")]) + sorted_hits = hits.take(sort_idx) + + blobs = table.fetch_blobs("info.blob", sorted_hits) + expected = [f"payload-{i}".encode() for i in sorted_hits["id"].to_pylist()] + assert [blobs[i].as_py() for i in range(len(blobs))] == expected + + +def _identifiable_payload(size: int) -> bytes: + block = 256 + return b"".join(bytes([i % 256]) * block for i in range(size // block)) diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index b60efc68e..62ee6e1eb 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -11,6 +11,7 @@ import lancedb from lancedb.db import AsyncConnection from lancedb.embeddings.base import TextEmbeddingFunction from lancedb.embeddings.registry import get_registry, register +from lancedb.expr import col from lancedb.index import FTS, IvfPq import lancedb.pydantic import numpy as np @@ -63,11 +64,71 @@ def _blob_query_data(): ) +def _create_blob_v2_query_table(db, name): + schema = pa.schema( + [ + pa.field("id", pa.int64()), + pa.field("tag", pa.utf8()), + pa.field("vector", pa.list_(pa.float32(), list_size=2)), + lancedb.blob("blob"), + ] + ) + table = db.create_table(name, schema=schema) + table.add( + [ + {"id": 1, "tag": "drop", "vector": [1.0, 0.0], "blob": b"one"}, + {"id": 2, "tag": "keep", "vector": [2.0, 0.0], "blob": b"two"}, + {"id": 3, "tag": "keep", "vector": [3.0, 0.0], "blob": b"three"}, + {"id": 4, "tag": "keep", "vector": [4.0, 0.0], "blob": b"four"}, + ] + ) + return table + + +async def _create_blob_v2_query_table_async(db, name): + schema = pa.schema( + [ + pa.field("id", pa.int64()), + pa.field("tag", pa.utf8()), + pa.field("vector", pa.list_(pa.float32(), list_size=2)), + lancedb.blob("blob"), + ] + ) + table = await db.create_table(name, schema=schema) + await table.add( + [ + {"id": 1, "tag": "drop", "vector": [1.0, 0.0], "blob": b"one"}, + {"id": 2, "tag": "keep", "vector": [2.0, 0.0], "blob": b"two"}, + {"id": 3, "tag": "keep", "vector": [3.0, 0.0], "blob": b"three"}, + {"id": 4, "tag": "keep", "vector": [4.0, 0.0], "blob": b"four"}, + ] + ) + return table + + def _assert_lazy_blob(value, expected: bytes): assert hasattr(value, "readall") assert value.readall() == expected +def _assert_blob_bytes_projection(df): + assert df["id_alias"].tolist() == [3, 4] + assert df["payload"].tolist() == [b"three", b"four"] + assert df["double_id"].tolist() == [6, 8] + + +def _blob_query_table(db, name, blob_schema): + if blob_schema == "v1": + return db.create_table(name, _blob_query_data()) + return _create_blob_v2_query_table(db, name) + + +async def _blob_query_table_async(db, name, blob_schema): + if blob_schema == "v1": + return await db.create_table(name, _blob_query_data()) + return await _create_blob_v2_query_table_async(db, name) + + @pytest.fixture(scope="module") def table(tmpdir_factory) -> lancedb.table.Table: tmp_path = str(tmpdir_factory.mktemp("data")) @@ -235,10 +296,11 @@ def test_plain_scan_query_to_pandas_blob_modes(tmp_db, blob_mode): assert not hasattr(first, "readall") -def test_plain_scan_query_to_pandas_blob_projection(tmp_db): +@pytest.mark.parametrize("blob_schema", ["v1", "v2"]) +def test_plain_scan_query_to_pandas_blob_bytes_projection(tmp_db, blob_schema): pytest.importorskip("lance") - table = tmp_db.create_table( - "test_query_to_pandas_blob_projection", _blob_query_data() + table = _blob_query_table( + tmp_db, f"test_query_to_pandas_blob_{blob_schema}_bytes", blob_schema ) df = ( @@ -250,9 +312,8 @@ def test_plain_scan_query_to_pandas_blob_projection(tmp_db): .to_pandas(blob_mode="bytes") ) - assert df["id_alias"].tolist() == [3, 4] - assert df["payload"].tolist() == [b"three", b"four"] - assert df["double_id"].tolist() == [6, 8] + _assert_blob_bytes_projection(df) + assert "_rowid" not in df.columns @pytest.mark.parametrize("blob_mode", ["bytes", "descriptions"]) @@ -348,18 +409,6 @@ async def test_async_plain_scan_query_to_pandas_blob_projection(tmp_db_async): assert lazy_df["id"].tolist() == [1] _assert_lazy_blob(lazy_df["blob"].iloc[0], b"one") - bytes_df = await ( - table.query() - .where("id >= 2") - .select({"id_alias": "id", "payload": "blob", "double_id": "id * 2"}) - .limit(2) - .offset(1) - .to_pandas(blob_mode="bytes") - ) - assert bytes_df["id_alias"].tolist() == [3, 4] - assert bytes_df["payload"].tolist() == [b"three", b"four"] - assert bytes_df["double_id"].tolist() == [6, 8] - desc_df = await ( table.query() .where("id = 1") @@ -371,6 +420,31 @@ async def test_async_plain_scan_query_to_pandas_blob_projection(tmp_db_async): assert not hasattr(first, "readall") +@pytest.mark.asyncio +@pytest.mark.parametrize("blob_schema", ["v1", "v2"]) +async def test_async_plain_scan_query_to_pandas_blob_bytes_projection( + tmp_db_async, blob_schema +): + pytest.importorskip("lance") + table = await _blob_query_table_async( + tmp_db_async, + f"test_async_query_to_pandas_blob_{blob_schema}_bytes", + blob_schema, + ) + + df = await ( + table.query() + .where("id >= 2") + .select({"id_alias": "id", "payload": "blob", "double_id": "id * 2"}) + .limit(2) + .offset(1) + .to_pandas(blob_mode="bytes") + ) + + _assert_blob_bytes_projection(df) + assert "_rowid" not in df.columns + + @pytest.mark.asyncio @pytest.mark.parametrize("blob_mode", ["bytes", "descriptions"]) async def test_async_plain_scan_query_to_pandas_blob_mode_does_not_collect_arrow( @@ -502,6 +576,18 @@ def test_with_row_id(table: lancedb.table.Table): assert rs["_rowid"].to_pylist() == [0, 1] +def test_blob_v2_query_omits_auto_row_id(tmp_db): + table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_omits_auto_rowid") + + query_obj = table.search().select(["id", "blob"]).limit(2).to_query_object() + assert query_obj.with_row_id is None + + rs = table.search().select(["id", "blob"]).limit(2).to_arrow() + + assert "_rowid" not in rs.column_names + assert rs["id"].to_pylist() == [1, 2] + + def test_where_repeated_combines_with_and(table: lancedb.table.Table): # Calling where() more than once should AND the filters together instead of # silently replacing the previous one (regression test for #2649). @@ -1946,3 +2032,39 @@ def test_fast_search(tmp_path): # 2. Fast Search -> Should NOT include "LanceScan" (Uses Index) plan = table.search(q).fast_search().explain_plan(True) assert "LanceScan" not in plan + + +def test_blob_v2_with_row_id_bytes_pandas(tmp_db): + table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_rowid_bytes_pandas") + + df = ( + table.search() + .with_row_id(True) + .select(["id", "blob"]) + .to_pandas(blob_mode="bytes") + ) + + assert "_rowid" in df.columns + assert df["id"].tolist() == [1, 2, 3, 4] + assert df["blob"].tolist() == [b"one", b"two", b"three", b"four"] + + +def test_blob_v2_expr_projection_stash(tmp_db): + table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_expr_projection_stash") + + hits = table.search().select({"blob_alias": col("blob")}).limit(2).to_arrow() + + assert "_rowid" not in hits.column_names + assert "_lance_row_id" in hits.schema.field("blob_alias").type.names + blobs = table.fetch_blobs("blob", hits) + assert [blobs[i].as_py() for i in range(len(blobs))] == [b"one", b"two"] + + +def test_blob_v2_to_batches_row_id(tmp_db): + table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_to_batches_rowid") + + hits = table.search().select(["id", "blob"]).limit(2).to_batches().read_all() + + assert "_rowid" in hits.column_names + blobs = table.fetch_blobs("blob", hits) + assert [blobs[i].as_py() for i in range(len(blobs))] == [b"one", b"two"] diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index dae9d5b2e..21cfb3df8 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -45,6 +45,32 @@ def _blob_test_data(): ) +def _blob_v2_table(db: DBConnection, name: str): + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")]) + table = db.create_table(name, schema=schema) + table.add([{"id": 1, "blob": b"hello"}, {"id": 2, "blob": b"world"}]) + return table + + +async def _blob_v2_table_async(db: AsyncConnection, name: str): + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")]) + table = await db.create_table(name, schema=schema) + await table.add([{"id": 1, "blob": b"hello"}, {"id": 2, "blob": b"world"}]) + return table + + +def _blob_table(db: DBConnection, name: str, blob_schema: str): + if blob_schema == "v1": + return db.create_table(name, data=_blob_test_data()) + return _blob_v2_table(db, name) + + +async def _blob_table_async(db: AsyncConnection, name: str, blob_schema: str): + if blob_schema == "v1": + return await db.create_table(name, data=_blob_test_data()) + return await _blob_v2_table_async(db, name) + + def _assert_lazy_blob(value, expected: bytes): assert hasattr(value, "readall") assert value.readall() == expected @@ -107,6 +133,18 @@ def test_table_to_pandas_blob_modes(tmp_db: DBConnection, blob_mode): assert not hasattr(first, "readall") +@pytest.mark.parametrize("blob_schema", ["v1", "v2"]) +def test_table_to_pandas_blob_bytes(tmp_db: DBConnection, blob_schema): + pytest.importorskip("lance") + table = _blob_table(tmp_db, f"test_to_pandas_blob_{blob_schema}_bytes", blob_schema) + + df = table.to_pandas(blob_mode="bytes") + + assert list(df.columns) == ["id", "blob"] + assert df["blob"].tolist() == [b"hello", b"world"] + assert "_rowid" not in df.columns + + def test_table_to_pandas_kwargs(tmp_db: DBConnection): pd = pytest.importorskip("pandas") data = pa.table({"id": pa.array([1, 2], pa.int64())}) @@ -118,15 +156,20 @@ def test_table_to_pandas_kwargs(tmp_db: DBConnection): @pytest.mark.asyncio -async def test_async_table_to_pandas_blob_bytes(tmp_db_async: AsyncConnection): +@pytest.mark.parametrize("blob_schema", ["v1", "v2"]) +async def test_async_table_to_pandas_blob_bytes( + tmp_db_async: AsyncConnection, blob_schema +): pytest.importorskip("lance") - table = await tmp_db_async.create_table( - "test_async_to_pandas_blob_bytes", data=_blob_test_data() + table = await _blob_table_async( + tmp_db_async, f"test_async_to_pandas_blob_{blob_schema}_bytes", blob_schema ) df = await table.to_pandas(blob_mode="bytes") + assert list(df.columns) == ["id", "blob"] assert df["blob"].tolist() == [b"hello", b"world"] + assert "_rowid" not in df.columns @pytest.mark.asyncio diff --git a/python/src/lib.rs b/python/src/lib.rs index 024ef292f..a72386305 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -16,7 +16,7 @@ use query::{FTSQuery, HybridQuery, Query, VectorQuery}; use session::Session; use table::{ AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, LsmWriteSpec, - MergeResult, Table, UpdateFieldMetadataResult, UpdateResult, + MergeResult, PyBlobFile, Table, UpdateFieldMetadataResult, UpdateResult, }; pub mod arrow; @@ -44,6 +44,7 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/python/src/table.rs b/python/src/table.rs index 7bc17ccea..fe2504ec2 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright The LanceDB Authors use std::{collections::HashMap, sync::Arc}; -use crate::runtime::future_into_py; +use crate::runtime::{block_on, future_into_py}; use crate::{ connection::Connection, error::PythonErrorExt, @@ -12,10 +12,12 @@ use crate::{ table::scannable::PyScannable, }; use arrow::{ + array::{Array, LargeBinaryArray}, datatypes::{DataType, Schema}, ffi_stream::ArrowArrayStreamReader, pyarrow::{FromPyArrow, PyArrowType, ToPyArrow}, }; +use lancedb::blob::BlobFile; use lancedb::table::{ AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable, @@ -24,7 +26,7 @@ use pyo3::{ Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python, exceptions::{PyRuntimeError, PyValueError}, pyclass, pymethods, - types::{IntoPyDict, PyAnyMethods, PyDict, PyDictMethods}, + types::{IntoPyDict, PyAnyMethods, PyBytes, PyDict, PyDictMethods}, }; mod scannable; @@ -412,6 +414,78 @@ impl From for DropColumnsResult { } } +/// Lazy blob handle from ``Table.fetch_blob_files``. +#[pyclass(name = "BlobFile")] +pub struct PyBlobFile { + inner: Arc, +} + +#[pymethods] +impl PyBlobFile { + fn read_bytes(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner.clone(); + let bytes = block_on(async move { inner.read().await }) + .map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?; + Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind()) + } + + pub fn read(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + let bytes = inner + .read() + .await + .map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?; + Python::attach(|py| Ok(PyBytes::new(py, bytes.as_ref()).unbind())) + }) + } + + fn close(self_: PyRef<'_, Self>) -> PyResult<()> { + let inner = self_.inner.clone(); + block_on(async move { inner.close().await }) + .map_err(|e| PyRuntimeError::new_err(format!("blob close failed: {e}"))) + } + + fn is_closed(self_: PyRef<'_, Self>) -> bool { + let inner = self_.inner.clone(); + block_on(async move { inner.is_closed().await }) + } + + fn seek(self_: PyRef<'_, Self>, position: u64) -> PyResult<()> { + let inner = self_.inner.clone(); + block_on(async move { inner.seek(position).await }) + .map_err(|e| PyRuntimeError::new_err(format!("blob seek failed: {e}"))) + } + + fn tell(self_: PyRef<'_, Self>) -> PyResult { + let inner = self_.inner.clone(); + block_on(async move { inner.tell().await }) + .map_err(|e| PyRuntimeError::new_err(format!("blob tell failed: {e}"))) + } + + fn size(self_: PyRef<'_, Self>) -> u64 { + self_.inner.size() + } + + /// Read a blob-local byte range without moving the cursor. + fn read_range(self_: PyRef<'_, Self>, offset: u64, length: usize) -> PyResult> { + let end = offset + .checked_add(length as u64) + .ok_or_else(|| PyValueError::new_err("offset + length overflowed"))?; + let inner = self_.inner.clone(); + let bytes = block_on(async move { inner.read_range(offset..end).await }) + .map_err(|e| PyRuntimeError::new_err(format!("blob read_range failed: {e}")))?; + Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind()) + } + + fn read_up_to(self_: PyRef<'_, Self>, length: usize) -> PyResult> { + let inner = self_.inner.clone(); + let bytes = block_on(async move { inner.read_up_to(length).await }) + .map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?; + Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind()) + } +} + #[pyclass] pub struct Table { // We keep a copy of the name to use if the inner table is dropped @@ -901,6 +975,55 @@ impl Table { )) } + /// Names of the blob v2 columns declared on this table, in declaration order. + pub fn blob_columns(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner_ref()?.clone(); + future_into_py(self_.py(), async move { + inner.blob_columns().await.infer_error() + }) + } + + /// Read blob bytes for `row_ids` from blob v2 column `column`. + #[pyo3(signature = (column, row_ids))] + pub fn fetch_blobs( + self_: PyRef<'_, Self>, + column: String, + row_ids: Vec, + ) -> PyResult> { + let inner = self_.inner_ref()?.clone(); + future_into_py(self_.py(), async move { + let blobs: LargeBinaryArray = inner + .fetch_blobs(column.as_str(), &row_ids) + .await + .infer_error()?; + Python::attach(|py| blobs.to_data().to_pyarrow(py).map(|obj| obj.unbind())) + }) + } + + /// Open lazy blob handles for `row_ids` from blob v2 column `column`. + #[pyo3(signature = (column, row_ids))] + pub fn fetch_blob_files( + self_: PyRef<'_, Self>, + column: String, + row_ids: Vec, + ) -> PyResult> { + let inner = self_.inner_ref()?.clone(); + future_into_py(self_.py(), async move { + let handles = inner + .fetch_blob_files(column.as_str(), &row_ids) + .await + .infer_error()?; + Ok(handles + .into_iter() + .map(|handle| { + handle.map(|file| PyBlobFile { + inner: Arc::new(file), + }) + }) + .collect::>()) + }) + } + /// Optimize the on-disk data by compacting and pruning old data, for better performance. #[pyo3(signature = (cleanup_since_ms=None, delete_unverified=None))] pub fn optimize(