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:
Drew Gallardo
2026-07-10 12:54:16 -07:00
committed by GitHub
parent 104fc5a08e
commit a548e59d49
13 changed files with 1833 additions and 119 deletions
+3 -1
View File
@@ -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",
+420
View File
@@ -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]
+18
View File
@@ -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
View File
@@ -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)
+13
View File
@@ -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.
+125 -1
View File
@@ -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.
+85 -29
View File
@@ -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.
+14 -1
View File
@@ -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"]
+562
View File
@@ -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))
+140 -18
View File
@@ -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"]
+46 -3
View File
@@ -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
View File
@@ -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
View File
@@ -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(