mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-16 11:08:24 +00:00
feat(python): blob v2 fetch API (#3578)
Python bindings for blob v2 read on **local** tables. Rust read APIs landed in #3562. This PR wires `fetch_blob_files`, `fetch_blobs`, v2 query/`to_pandas(blob_mode="bytes")`, and hidden `_rowid` metadata so `fetch_*` works from query hits without exposing `_rowid` in the column list. **Cloud:** `RemoteTable.fetch_blobs` / `fetch_blob_files` raise `NotImplementedError` until Phalanx ships the server route (separate track; not blocking local merge). ### Primary path: lazy file handles ```python table = db.create_table("videos", schema=pa.schema([ pa.field("id", pa.int64()), lancedb.blob("video"), ])) table.add([{"id": 1, "video": open("clip.mp4", "rb").read()}]) hits = table.search().select(["id", "video"]).to_arrow() handle = table.fetch_blob_files("video", hits)[0] # seek + partial read — PyAV / decoders can use the handle handle.seek(frame_offset) chunk = handle.read_range(0, 65536) ``` `BlobFile` exposes `seek`, `read`, `read_range`, `read_up_to`, and works with `BufferedReader`. ### When you want full bytes ```python blobs = table.fetch_blobs("video", hits) # eager materialize, null-aligned df = table.to_pandas(blob_mode="bytes") # descriptors → bytes in pandas ``` ### `_rowid` (join key, not user `id`) Fetch needs Lance row ids. For v2 blob queries we auto-inject `_rowid`, stash it in Arrow schema metadata on `to_arrow()`, and drop the visible column unless you pass `.with_row_id(True)`. v1 legacy blobs (`lance-encoding:blob`) unchanged; fetch on v1 raises the migration error. ## Test plan - [x] `./scripts/test-blob.sh python` (105 passed in worktree) - [x] `fetch_blob_files` lazy read, seek, partial read, null alignment, cross-fragment dups - [x] hybrid query → `fetch_blobs` / `fetch_blob_files` - [ ] Will re-review after seek/`BlobFile` commit (`d77ab1a6`) --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -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",
|
||||
|
||||
@@ -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"<BlobFile size={self.size()}>"
|
||||
|
||||
|
||||
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]
|
||||
@@ -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:
|
||||
|
||||
+280
-63
@@ -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)
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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))
|
||||
@@ -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"]
|
||||
|
||||
@@ -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
|
||||
|
||||
+2
-1
@@ -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::<Connection>()?;
|
||||
m.add_class::<Session>()?;
|
||||
m.add_class::<Table>()?;
|
||||
m.add_class::<PyBlobFile>()?;
|
||||
m.add_class::<IndexConfig>()?;
|
||||
m.add_class::<Query>()?;
|
||||
m.add_class::<FTSQuery>()?;
|
||||
|
||||
+125
-2
@@ -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<lancedb::table::DropColumnsResult> for DropColumnsResult {
|
||||
}
|
||||
}
|
||||
|
||||
/// Lazy blob handle from ``Table.fetch_blob_files``.
|
||||
#[pyclass(name = "BlobFile")]
|
||||
pub struct PyBlobFile {
|
||||
inner: Arc<BlobFile>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyBlobFile {
|
||||
fn read_bytes(self_: PyRef<'_, Self>) -> PyResult<Py<PyBytes>> {
|
||||
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<Bound<'_, PyAny>> {
|
||||
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<u64> {
|
||||
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<Py<PyBytes>> {
|
||||
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<Py<PyBytes>> {
|
||||
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<Bound<'_, PyAny>> {
|
||||
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<u64>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
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<u64>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
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::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
/// 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(
|
||||
|
||||
Reference in New Issue
Block a user