mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-15 16:22:27 +00:00
feat(python): add blob_mode to to_arrow and to_list (#4529)
This commit is contained in:
+164
-51
@@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Optional, Union
|
||||
import pyarrow as pa
|
||||
|
||||
from .expr import Expr
|
||||
from .schema import row_addressable_blob_v2_paths
|
||||
from .schema import blob_column_paths, row_addressable_blob_v2_paths
|
||||
from .types import BlobMode, QueryProjection, QueryProjectionSpec
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -27,8 +27,8 @@ BLOB_MODE_TO_HANDLING = {
|
||||
|
||||
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]]
|
||||
FetchBlobsSync = Callable[[str, list[int]], pa.Array | pa.ChunkedArray]
|
||||
FetchBlobsAsync = Callable[[str, list[int]], Awaitable[pa.Array | pa.ChunkedArray]]
|
||||
|
||||
|
||||
class BlobFile(io.RawIOBase):
|
||||
@@ -136,7 +136,7 @@ def projection_includes_blob_column(
|
||||
columns = set(blob_columns)
|
||||
if not columns:
|
||||
return False
|
||||
if projection is None:
|
||||
if not projection:
|
||||
return True
|
||||
for output, source in _iter_projection_pairs(projection):
|
||||
if output in columns or any(
|
||||
@@ -150,18 +150,68 @@ def blob_v2_projection_sources(
|
||||
schema: pa.Schema,
|
||||
projection: QueryProjection,
|
||||
) -> dict[str, str]:
|
||||
blob_columns = row_addressable_blob_v2_paths(schema)
|
||||
if not blob_columns:
|
||||
return {}
|
||||
if projection is None:
|
||||
return _project_blob_paths(row_addressable_blob_v2_paths(schema), projection)
|
||||
|
||||
|
||||
def reject_unfetchable_blob_columns(
|
||||
schema: pa.Schema, projection: QueryProjection, blob_sources: dict[str, str]
|
||||
) -> None:
|
||||
# Rerankers can drop field metadata, so validate against the source schema.
|
||||
unfetchable = [
|
||||
output
|
||||
for output in _project_blob_paths(blob_column_paths(schema), projection)
|
||||
if output not in blob_sources
|
||||
]
|
||||
if unfetchable:
|
||||
raise ValueError(
|
||||
f"blob_mode='bytes' cannot fetch blob columns {unfetchable}. Bytes are "
|
||||
"fetched by row id, which covers blob v2 columns with one blob per "
|
||||
"row but not legacy blob columns or blobs inside lists. Use "
|
||||
"blob_mode='descriptions' or remove these columns from the projection."
|
||||
)
|
||||
|
||||
|
||||
def reject_unmatched_blob_outputs(
|
||||
output_schema: pa.Schema, blob_sources: dict[str, str]
|
||||
) -> None:
|
||||
# Another spelling of a blob column, such as a quoted or recased name, is a
|
||||
# blob in the query's output schema but maps to no source by name.
|
||||
unmatched = [
|
||||
path for path in blob_column_paths(output_schema) if path not in blob_sources
|
||||
]
|
||||
if unmatched:
|
||||
raise ValueError(
|
||||
f"blob_mode='bytes' could not match result columns {unmatched} to a blob "
|
||||
"column of the table. Select blob columns by their plain names, or "
|
||||
"use blob_mode='descriptions'."
|
||||
)
|
||||
|
||||
|
||||
def _project_blob_paths(
|
||||
blob_columns: list[str], projection: QueryProjection
|
||||
) -> dict[str, str]:
|
||||
"""Map each projected output path to the blob column path it reads."""
|
||||
# An empty projection selects all columns.
|
||||
if not projection:
|
||||
return {column: column for column in blob_columns}
|
||||
sources = {}
|
||||
sources: dict[str, str] = {}
|
||||
entries: dict[str, str] = {}
|
||||
for output, source in _iter_projection_pairs(projection):
|
||||
for column in blob_columns:
|
||||
if column == source:
|
||||
sources[output] = column
|
||||
path = output
|
||||
elif column.startswith(f"{source}."):
|
||||
sources[f"{output}{column[len(source) :]}"] = column
|
||||
path = f"{output}{column[len(source) :]}"
|
||||
else:
|
||||
continue
|
||||
if path in sources:
|
||||
raise ValueError(
|
||||
f"projection entries {entries[path]!r} and {output!r} both "
|
||||
f"produce blob column '{path}'. Rename one of them or drop "
|
||||
"the redundant selection."
|
||||
)
|
||||
sources[path] = column
|
||||
entries[path] = output
|
||||
return sources
|
||||
|
||||
|
||||
@@ -212,11 +262,14 @@ async def replace_v2_blob_columns_with_bytes(
|
||||
blob_sources: dict[str, str],
|
||||
fetch_blobs: FetchBlobsAsync,
|
||||
) -> pa.Table:
|
||||
for output_name, source_name in blob_sources.items():
|
||||
if not _has_column_path(tbl.schema, output_name):
|
||||
continue
|
||||
blobs = await fetch_blobs(source_name, tbl)
|
||||
tbl = _set_blob_column(tbl, output_name, blobs)
|
||||
present = _present_blob_sources(tbl, blob_sources)
|
||||
if not present:
|
||||
return tbl
|
||||
row_ids = _hit_row_ids(tbl, present)
|
||||
for output_name, source_name in present.items():
|
||||
tbl = _set_blob_column(
|
||||
tbl, output_name, await fetch_blobs(source_name, row_ids)
|
||||
)
|
||||
return tbl
|
||||
|
||||
|
||||
@@ -225,19 +278,48 @@ def replace_v2_blob_columns_with_bytes_sync(
|
||||
blob_sources: dict[str, str],
|
||||
fetch_blobs: FetchBlobsSync,
|
||||
) -> pa.Table:
|
||||
for output_name, source_name in blob_sources.items():
|
||||
if not _has_column_path(tbl.schema, output_name):
|
||||
continue
|
||||
blobs = fetch_blobs(source_name, tbl)
|
||||
tbl = _set_blob_column(tbl, output_name, blobs)
|
||||
present = _present_blob_sources(tbl, blob_sources)
|
||||
if not present:
|
||||
return tbl
|
||||
row_ids = _hit_row_ids(tbl, present)
|
||||
for output_name, source_name in present.items():
|
||||
tbl = _set_blob_column(tbl, output_name, fetch_blobs(source_name, row_ids))
|
||||
return tbl
|
||||
|
||||
|
||||
def _present_blob_sources(
|
||||
tbl: pa.Table, blob_sources: dict[str, str]
|
||||
) -> dict[str, str]:
|
||||
return {
|
||||
output: source
|
||||
for output, source in blob_sources.items()
|
||||
if _has_column_path(tbl.schema, output)
|
||||
}
|
||||
|
||||
|
||||
def _hit_row_ids(tbl: pa.Table, blob_sources: dict[str, str]) -> list[int]:
|
||||
"""Read row ids before replacing columns, since a swapped alias can overwrite
|
||||
the lookup source."""
|
||||
if "_rowid" in tbl.column_names:
|
||||
return tbl["_rowid"].to_pylist()
|
||||
for output in blob_sources:
|
||||
descriptor = _leaf_struct_column(tbl, output)
|
||||
if pa.types.is_struct(descriptor.type) and (
|
||||
ROW_ID_FIELD_NAME in descriptor.type.names
|
||||
):
|
||||
return descriptor.field(ROW_ID_FIELD_NAME).to_pylist()
|
||||
raise ValueError(
|
||||
f"query result has no '_rowid' column and no '{ROW_ID_FIELD_NAME}' "
|
||||
f"field on blob columns {list(blob_sources)}. Call .with_row_id(True) "
|
||||
"to fetch blob bytes for this query."
|
||||
)
|
||||
|
||||
|
||||
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]
|
||||
present_paths = [p for p in blob_paths if _has_column_path(tbl.schema, p)]
|
||||
if not present_paths:
|
||||
raise ValueError("query result has no blob v2 column to carry a row id")
|
||||
|
||||
@@ -257,7 +339,7 @@ def read_row_ids_from_hits(hits: pa.Table, blob_column: str) -> list[int]:
|
||||
|
||||
try:
|
||||
leaf = _leaf_struct_column(hits, blob_column)
|
||||
if ROW_ID_FIELD_NAME in leaf.type.names:
|
||||
if pa.types.is_struct(leaf.type) and ROW_ID_FIELD_NAME in leaf.type.names:
|
||||
return leaf.field(ROW_ID_FIELD_NAME).to_pylist()
|
||||
except KeyError:
|
||||
pass
|
||||
@@ -321,31 +403,22 @@ def _iter_projection_pairs(
|
||||
|
||||
|
||||
def _has_column_path(schema: pa.Schema, path: str) -> bool:
|
||||
if schema.get_field_index(path) >= 0:
|
||||
return True
|
||||
top_name, *rest = path.split(".")
|
||||
top_name, rest = _split_column_path(schema, path)
|
||||
index = schema.get_field_index(top_name)
|
||||
if index < 0:
|
||||
return False
|
||||
field = schema.field(index)
|
||||
for name in rest:
|
||||
if not pa.types.is_struct(field.type):
|
||||
return False
|
||||
index = field.type.get_field_index(name)
|
||||
if index < 0:
|
||||
return False
|
||||
field = field.type.field(index)
|
||||
return True
|
||||
return (
|
||||
index >= 0 and _walk_struct_fields(schema.field(index).type, rest) is not None
|
||||
)
|
||||
|
||||
|
||||
def _set_blob_column(
|
||||
tbl: pa.Table, output_name: str, blobs: pa.Array | pa.ChunkedArray
|
||||
) -> pa.Table:
|
||||
index = tbl.schema.get_field_index(output_name)
|
||||
if index >= 0:
|
||||
return tbl.set_column(index, pa.field(output_name, blobs.type), [blobs])
|
||||
top_name, rest = _split_column_path(tbl.schema, output_name)
|
||||
if not rest:
|
||||
index = tbl.schema.get_field_index(top_name)
|
||||
return tbl.set_column(index, pa.field(top_name, blobs.type), [blobs])
|
||||
|
||||
parent_path, leaf_name = output_name.rsplit(".", 1)
|
||||
leaf_name = rest[-1]
|
||||
if isinstance(blobs, pa.ChunkedArray):
|
||||
blobs = blobs.combine_chunks()
|
||||
|
||||
@@ -357,9 +430,9 @@ def _set_blob_column(
|
||||
leaf_name, blobs.type, nullable=field.nullable
|
||||
)
|
||||
return
|
||||
raise KeyError(f"Blob field '{leaf_name}' not found in struct '{parent_path}'")
|
||||
raise KeyError(f"Blob field '{leaf_name}' not found in '{output_name}'")
|
||||
|
||||
return _transform_struct_column(tbl, parent_path, replace_blob)
|
||||
return _transform_struct_column(tbl, top_name, rest[:-1], replace_blob)
|
||||
|
||||
|
||||
def _embed_row_id_in_column(tbl: pa.Table, path: str, row_ids: pa.Array) -> pa.Table:
|
||||
@@ -367,7 +440,9 @@ def _embed_row_id_in_column(tbl: pa.Table, path: str, row_ids: pa.Array) -> pa.T
|
||||
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)
|
||||
return _transform_struct_column(
|
||||
tbl, *_split_column_path(tbl.schema, path), add_row_id
|
||||
)
|
||||
|
||||
|
||||
def strip_auto_row_ids(tbl: pa.Table, blob_paths: Iterable[str]) -> pa.Table:
|
||||
@@ -384,16 +459,20 @@ def strip_auto_row_ids(tbl: pa.Table, blob_paths: Iterable[str]) -> pa.Table:
|
||||
return
|
||||
|
||||
for path in blob_paths:
|
||||
if path.split(".")[0] not in tbl.column_names:
|
||||
if not _has_column_path(tbl.schema, path):
|
||||
continue
|
||||
tbl = _transform_struct_column(tbl, path, drop_row_id)
|
||||
tbl = _transform_struct_column(
|
||||
tbl, *_split_column_path(tbl.schema, path), drop_row_id
|
||||
)
|
||||
return tbl
|
||||
|
||||
|
||||
def _transform_struct_column(
|
||||
tbl: pa.Table, path: str, leaf_transform: Callable[[list, list], None]
|
||||
tbl: pa.Table,
|
||||
top_name: str,
|
||||
rest: list[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)
|
||||
@@ -451,15 +530,49 @@ def _rebuild_struct(
|
||||
|
||||
|
||||
def _leaf_struct_column(tbl: pa.Table, path: str) -> pa.StructArray:
|
||||
parts = path.split(".")
|
||||
column = tbl.column(parts[0])
|
||||
top_name, rest = _split_column_path(tbl.schema, path)
|
||||
column = tbl.column(top_name)
|
||||
if isinstance(column, pa.ChunkedArray):
|
||||
column = column.combine_chunks()
|
||||
for part in parts[1:]:
|
||||
for part in rest:
|
||||
if not pa.types.is_struct(column.type):
|
||||
raise KeyError(f"blob path '{path}' crosses non-struct column")
|
||||
column = column.field(part)
|
||||
return column
|
||||
|
||||
|
||||
def _split_column_path(schema: pa.Schema, path: str) -> tuple[str, list[str]]:
|
||||
"""Prefer a path ending in a struct, since a flat alias can shadow a nested
|
||||
blob path."""
|
||||
parts = path.split(".")
|
||||
fallback = None
|
||||
for cut in range(len(parts), 0, -1):
|
||||
top_name, rest = ".".join(parts[:cut]), parts[cut:]
|
||||
index = schema.get_field_index(top_name)
|
||||
if index < 0:
|
||||
continue
|
||||
leaf_type = _walk_struct_fields(schema.field(index).type, rest)
|
||||
if leaf_type is None:
|
||||
continue
|
||||
if pa.types.is_struct(leaf_type):
|
||||
return top_name, rest
|
||||
fallback = fallback or (top_name, rest)
|
||||
return fallback or (parts[0], parts[1:])
|
||||
|
||||
|
||||
def _walk_struct_fields(
|
||||
data_type: pa.DataType, names: list[str]
|
||||
) -> Optional[pa.DataType]:
|
||||
for name in names:
|
||||
if not pa.types.is_struct(data_type):
|
||||
return None
|
||||
index = data_type.get_field_index(name)
|
||||
if index < 0:
|
||||
return None
|
||||
data_type = data_type.field(index).type
|
||||
return data_type
|
||||
|
||||
|
||||
def _normalize_blob_row_ids(
|
||||
row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], blob_column: str
|
||||
) -> list[int]:
|
||||
|
||||
+345
-92
@@ -12,6 +12,7 @@ from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Dict,
|
||||
Iterable,
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
@@ -40,7 +41,7 @@ from .rerankers.base import Reranker
|
||||
from .rerankers.rrf import RRFReranker
|
||||
from .rerankers.util import check_reranker_result
|
||||
from .schema import blob_column_paths, is_blob_like_field, schema_has_blob_field
|
||||
from .util import _validate_query_vector, flatten_columns
|
||||
from .util import _validate_query_vector, flatten_columns, get_uri_scheme
|
||||
from . import _wal_hybrid # WAL-PK-FUSION: delete.
|
||||
from ._blob import (
|
||||
BLOB_MODE_TO_HANDLING,
|
||||
@@ -49,6 +50,8 @@ from ._blob import (
|
||||
blob_auto_row_id_for_scan,
|
||||
blob_v2_projection_sources,
|
||||
finalize_blob_query_table,
|
||||
reject_unfetchable_blob_columns,
|
||||
reject_unmatched_blob_outputs,
|
||||
replace_v2_blob_columns_with_bytes,
|
||||
replace_v2_blob_columns_with_bytes_sync,
|
||||
strip_auto_row_ids,
|
||||
@@ -99,15 +102,45 @@ class _LanceScanner(Protocol):
|
||||
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)
|
||||
|
||||
|
||||
def _unsupported_blob_pandas_error(reason: str) -> RuntimeError:
|
||||
return RuntimeError(
|
||||
"blob columns require Lance native scanner conversion for query "
|
||||
f"to_pandas(), but {reason}. Use a plain scan query or remove blob "
|
||||
"columns from the projection."
|
||||
f"to_pandas(), but {reason}. Use a plain scan query, pass "
|
||||
"blob_mode='bytes' or blob_mode='descriptions', or remove blob columns "
|
||||
"from the projection."
|
||||
)
|
||||
|
||||
|
||||
def _supports_local_blob_fetch(
|
||||
async_table: "AsyncTable", pushdown_operations: Iterable[str]
|
||||
) -> bool:
|
||||
# Remote tables and QueryTable pushdown answer queries on a server, so a
|
||||
# fetch through this handle could read another version.
|
||||
return async_table._inner._is_native() and "QueryTable" not in pushdown_operations
|
||||
|
||||
|
||||
def _blob_bytes_without_snapshot_error() -> NotImplementedError:
|
||||
return NotImplementedError(
|
||||
"blob_mode='bytes' cannot safely materialize blob columns without a "
|
||||
"stable table snapshot. Remote tables and namespace QueryTable pushdown "
|
||||
"run the query on a server, so the blob fetch could read another table "
|
||||
"version. Use blob_mode='descriptions'."
|
||||
)
|
||||
|
||||
|
||||
def _query_needs_lance_scanner(builder: Any) -> bool:
|
||||
"""Whether the query sets options that to_arrow ignores."""
|
||||
to_query_object = getattr(builder, "to_query_object", None)
|
||||
if to_query_object is None:
|
||||
return False
|
||||
try:
|
||||
query = to_query_object()
|
||||
except NotImplementedError:
|
||||
return False
|
||||
return (
|
||||
query.fragments is not None
|
||||
or query.fragment_ids is not None
|
||||
or bool(query.with_row_address)
|
||||
)
|
||||
|
||||
|
||||
@@ -151,7 +184,8 @@ def _combine_where(
|
||||
|
||||
|
||||
def _projection_to_scanner_kwargs(columns: QueryProjection) -> Dict[str, Any]:
|
||||
if columns is None:
|
||||
# An empty projection selects all columns.
|
||||
if not columns:
|
||||
return {}
|
||||
if isinstance(columns, list):
|
||||
if all(isinstance(column, str) for column in columns):
|
||||
@@ -175,7 +209,8 @@ def _projection_to_scanner_kwargs(columns: QueryProjection) -> Dict[str, Any]:
|
||||
|
||||
|
||||
def _query_request_projection(req: "PyQueryRequest") -> QueryProjection:
|
||||
if req.select_source_columns is not None:
|
||||
# Treat empty map as a calculated select
|
||||
if req.select_source_columns:
|
||||
return req.select_source_columns
|
||||
return req.select
|
||||
|
||||
@@ -1072,44 +1107,42 @@ class LanceQueryBuilder(ABC):
|
||||
The maximum time to wait for the query to complete.
|
||||
If None, wait indefinitely.
|
||||
blob_mode: str, default "lazy"
|
||||
Controls how blob columns are returned for plain scan queries.
|
||||
Vector, FTS, hybrid, and other non-native query shapes keep the
|
||||
existing Arrow conversion path and only support blob descriptions.
|
||||
Controls how blob columns are returned. ``"lazy"`` returns
|
||||
:class:`~lancedb._blob.BlobFile` handles and needs a plain scan.
|
||||
``"descriptions"`` and ``"bytes"`` work for every query shape. A
|
||||
plain scan without ``timeout`` reads the local dataset through
|
||||
Lance's scanner, which also returns legacy blob columns as bytes.
|
||||
Other shapes use :meth:`to_arrow` with the same ``blob_mode``.
|
||||
**kwargs
|
||||
Forwarded to pyarrow.Table.to_pandas after query execution and
|
||||
optional flattening.
|
||||
"""
|
||||
validate_blob_mode(blob_mode)
|
||||
output_schema = getattr(self, "output_schema", None)
|
||||
if output_schema is not None:
|
||||
schema = output_schema()
|
||||
if _blob_mode_requires_native_pandas(blob_mode, schema):
|
||||
native_error = None
|
||||
if (flatten is None or blob_mode == "descriptions") and timeout is None:
|
||||
try:
|
||||
df = self._plain_scan_to_pandas(
|
||||
blob_mode, flatten=flatten, **kwargs
|
||||
)
|
||||
if df is not None:
|
||||
return df
|
||||
except Exception as err:
|
||||
native_error = err
|
||||
reason = (
|
||||
if output_schema is not None and schema_has_blob_field(output_schema()):
|
||||
if (flatten is None or blob_mode != "lazy") and timeout is None:
|
||||
df = self._plain_scan_to_pandas(blob_mode, flatten=flatten, **kwargs)
|
||||
if df is not None:
|
||||
return df
|
||||
if blob_mode == "lazy" or _query_needs_lance_scanner(self):
|
||||
raise _unsupported_blob_pandas_error(
|
||||
"this query shape cannot use Lance native pandas conversion"
|
||||
if native_error is None
|
||||
else str(native_error)
|
||||
)
|
||||
raise _unsupported_blob_pandas_error(reason) from native_error
|
||||
|
||||
tbl = flatten_columns(self.to_arrow(timeout=timeout), flatten)
|
||||
if _blob_mode_requires_native_pandas(blob_mode, tbl.schema):
|
||||
tbl = flatten_columns(self._to_arrow_in_mode(timeout, blob_mode), flatten)
|
||||
if blob_mode == "lazy" and schema_has_blob_field(tbl.schema):
|
||||
raise _unsupported_blob_pandas_error(
|
||||
"this query shape cannot use Lance native pandas conversion"
|
||||
)
|
||||
return tbl.to_pandas(**kwargs)
|
||||
|
||||
@abstractmethod
|
||||
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||||
def to_arrow(
|
||||
self,
|
||||
*,
|
||||
timeout: Optional[timedelta] = None,
|
||||
blob_mode: BlobMode = "lazy",
|
||||
) -> pa.Table:
|
||||
"""
|
||||
Execute the query and return the results as an
|
||||
[Apache Arrow Table](https://arrow.apache.org/docs/python/generated/pyarrow.Table.html#pyarrow.Table).
|
||||
@@ -1123,6 +1156,11 @@ class LanceQueryBuilder(ABC):
|
||||
timeout: Optional[timedelta]
|
||||
The maximum time to wait for the query to complete.
|
||||
If None, wait indefinitely.
|
||||
blob_mode: str, default "lazy"
|
||||
``"lazy"`` returns blob v2 descriptors that carry the row id
|
||||
``fetch_blobs`` reads. ``"descriptions"`` returns them without it.
|
||||
``"bytes"`` returns payload bytes and raises for legacy blob columns
|
||||
and blobs inside lists.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@@ -1153,7 +1191,12 @@ class LanceQueryBuilder(ABC):
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def to_list(self, *, timeout: Optional[timedelta] = None) -> List[dict]:
|
||||
def to_list(
|
||||
self,
|
||||
*,
|
||||
timeout: Optional[timedelta] = None,
|
||||
blob_mode: BlobMode = "lazy",
|
||||
) -> List[dict]:
|
||||
"""
|
||||
Execute the query and return the results as a list of dictionaries.
|
||||
|
||||
@@ -1166,8 +1209,19 @@ class LanceQueryBuilder(ABC):
|
||||
timeout: Optional[timedelta]
|
||||
The maximum time to wait for the query to complete.
|
||||
If None, wait indefinitely.
|
||||
blob_mode: str, default "lazy"
|
||||
Passed through to :meth:`to_arrow`.
|
||||
"""
|
||||
return self.to_arrow(timeout=timeout).to_pylist()
|
||||
validate_blob_mode(blob_mode)
|
||||
return self._to_arrow_in_mode(timeout, blob_mode).to_pylist()
|
||||
|
||||
def _to_arrow_in_mode(
|
||||
self, timeout: Optional[timedelta], blob_mode: BlobMode
|
||||
) -> pa.Table:
|
||||
# Older subclasses may not accept blob_mode. Omit it when it has no effect.
|
||||
if blob_mode == "lazy" or not schema_has_blob_field(self._table.schema):
|
||||
return self.to_arrow(timeout=timeout)
|
||||
return self.to_arrow(timeout=timeout, blob_mode=blob_mode)
|
||||
|
||||
def to_pydantic(
|
||||
self, model: type[T], *, timeout: Optional[timedelta] = None
|
||||
@@ -1340,19 +1394,62 @@ class LanceQueryBuilder(ABC):
|
||||
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 ()
|
||||
def _validated_blob_sources(self, blob_mode: BlobMode) -> dict[str, str]:
|
||||
"""Map projected blob v2 outputs to their source columns.
|
||||
|
||||
Bytes mode is checked here, before the query runs.
|
||||
"""
|
||||
validate_blob_mode(blob_mode)
|
||||
schema = self._table.schema
|
||||
blob_sources = blob_v2_projection_sources(schema, self._columns)
|
||||
if blob_mode != "bytes":
|
||||
return blob_sources
|
||||
reject_unfetchable_blob_columns(schema, self._columns, blob_sources)
|
||||
if blob_sources and not _supports_local_blob_fetch(
|
||||
self._table._table, getattr(self._table, "_pushdown_operations", ())
|
||||
):
|
||||
raise _blob_bytes_without_snapshot_error()
|
||||
reject_unmatched_blob_outputs(self._blob_output_schema(), blob_sources)
|
||||
return blob_sources
|
||||
|
||||
def _blob_output_schema(self) -> pa.Schema:
|
||||
return self.output_schema()
|
||||
|
||||
def _apply_arrow_blob_mode(
|
||||
self, tbl: pa.Table, blob_mode: BlobMode, blob_sources: dict[str, str]
|
||||
) -> pa.Table:
|
||||
if blob_mode == "lazy":
|
||||
return tbl
|
||||
if blob_mode == "descriptions":
|
||||
return strip_auto_row_ids(tbl, blob_sources)
|
||||
return replace_v2_blob_columns_with_bytes_sync(
|
||||
tbl, blob_sources, self._table.fetch_blobs
|
||||
)
|
||||
return finalize_blob_query_table(
|
||||
|
||||
def _finalize_blob_query_table(
|
||||
self,
|
||||
tbl: pa.Table,
|
||||
*,
|
||||
blob_mode: BlobMode,
|
||||
blob_sources: dict[str, str],
|
||||
) -> pa.Table:
|
||||
blob_auto_row_id = self._blob_auto_row_id_enabled()
|
||||
finalized = 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,
|
||||
blob_paths=blob_sources.keys() if blob_auto_row_id else (),
|
||||
)
|
||||
finalized = self._apply_arrow_blob_mode(finalized, blob_mode, blob_sources)
|
||||
# with_row_id(False) does not reach the scan, which adds _rowid anyway.
|
||||
# Lazy keeps it because fetch_blobs(hits) reads it there.
|
||||
if (
|
||||
blob_mode != "lazy"
|
||||
and self._with_row_id is False
|
||||
and "_rowid" in finalized.column_names
|
||||
):
|
||||
finalized = finalized.drop_columns(["_rowid"])
|
||||
return finalized
|
||||
|
||||
def with_row_address(self, with_row_address: bool = True) -> Self:
|
||||
"""Set whether to return row addresses.
|
||||
@@ -1565,6 +1662,11 @@ class LanceQueryBuilder(ABC):
|
||||
if not _query_is_plain_scan(query):
|
||||
return None
|
||||
|
||||
# Remote and in-memory tables have no Lance dataset to scan.
|
||||
if not self._table._table._inner._is_native() or (
|
||||
get_uri_scheme(self._table._dataset_path) == "memory"
|
||||
):
|
||||
return None
|
||||
dataset = self._table.to_lance()
|
||||
blob_auto_row_id = self._blob_auto_row_id_enabled()
|
||||
blob_sources = (
|
||||
@@ -1572,6 +1674,10 @@ class LanceQueryBuilder(ABC):
|
||||
if blob_mode == "bytes"
|
||||
else {}
|
||||
)
|
||||
if blob_sources:
|
||||
reject_unfetchable_blob_columns(
|
||||
self._table.schema, query.columns, blob_sources
|
||||
)
|
||||
scanner = dataset.scanner(
|
||||
**_scanner_kwargs_for_query(
|
||||
query,
|
||||
@@ -1585,7 +1691,7 @@ class LanceQueryBuilder(ABC):
|
||||
blob_mode=blob_mode,
|
||||
blob_sources=blob_sources,
|
||||
fetch_blobs=self._table.fetch_blobs,
|
||||
strip_auto_row_id=blob_auto_row_id,
|
||||
strip_auto_row_id=not self._user_requested_row_id(),
|
||||
flatten=flatten,
|
||||
**kwargs,
|
||||
)
|
||||
@@ -1846,7 +1952,12 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
|
||||
"""
|
||||
return self._table._output_schema(self.to_query_object())
|
||||
|
||||
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||||
def to_arrow(
|
||||
self,
|
||||
*,
|
||||
timeout: Optional[timedelta] = None,
|
||||
blob_mode: BlobMode = "lazy",
|
||||
) -> pa.Table:
|
||||
"""
|
||||
Execute the query and return the results as an
|
||||
[Apache Arrow Table](https://arrow.apache.org/docs/python/generated/pyarrow.Table.html#pyarrow.Table).
|
||||
@@ -1860,9 +1971,17 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
|
||||
timeout: Optional[timedelta]
|
||||
The maximum time to wait for the query to complete.
|
||||
If None, wait indefinitely.
|
||||
blob_mode: str, default "lazy"
|
||||
``"lazy"`` returns blob v2 descriptors that carry the row id
|
||||
``fetch_blobs`` reads. ``"descriptions"`` returns them without it.
|
||||
``"bytes"`` returns payload bytes and raises for legacy blob columns
|
||||
and blobs inside lists.
|
||||
"""
|
||||
blob_sources = self._validated_blob_sources(blob_mode)
|
||||
return self._finalize_blob_query_table(
|
||||
self.to_batches(timeout=timeout).read_all()
|
||||
self.to_batches(timeout=timeout).read_all(),
|
||||
blob_mode=blob_mode,
|
||||
blob_sources=blob_sources,
|
||||
)
|
||||
|
||||
def to_query_object(self) -> Query:
|
||||
@@ -2135,16 +2254,23 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
||||
"""
|
||||
return self._table._output_schema(self.to_query_object())
|
||||
|
||||
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||||
def to_arrow(
|
||||
self,
|
||||
*,
|
||||
timeout: Optional[timedelta] = None,
|
||||
blob_mode: BlobMode = "lazy",
|
||||
) -> pa.Table:
|
||||
blob_sources = self._validated_blob_sources(blob_mode)
|
||||
self._table._ensure_no_legacy_fts_index()
|
||||
|
||||
query = self._query_for_scan()
|
||||
results = self._table._execute_query(query, timeout=timeout)
|
||||
results = results.read_all()
|
||||
results = self._table._execute_query(query, timeout=timeout).read_all()
|
||||
if self._reranker is not None:
|
||||
results = self._reranker.rerank_fts(self._query, results)
|
||||
check_reranker_result(results)
|
||||
return self._finalize_blob_query_table(results)
|
||||
return self._finalize_blob_query_table(
|
||||
results, blob_mode=blob_mode, blob_sources=blob_sources
|
||||
)
|
||||
|
||||
def to_batches(
|
||||
self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None
|
||||
@@ -2171,9 +2297,17 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
||||
|
||||
|
||||
class LanceEmptyQueryBuilder(LanceQueryBuilder):
|
||||
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||||
def to_arrow(
|
||||
self,
|
||||
*,
|
||||
timeout: Optional[timedelta] = None,
|
||||
blob_mode: BlobMode = "lazy",
|
||||
) -> pa.Table:
|
||||
blob_sources = self._validated_blob_sources(blob_mode)
|
||||
return self._finalize_blob_query_table(
|
||||
self.to_batches(timeout=timeout).read_all()
|
||||
self.to_batches(timeout=timeout).read_all(),
|
||||
blob_mode=blob_mode,
|
||||
blob_sources=blob_sources,
|
||||
)
|
||||
|
||||
def to_query_object(self) -> Query:
|
||||
@@ -2306,16 +2440,30 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
def to_query_object(self) -> Query:
|
||||
raise NotImplementedError("to_query_object not yet supported on a hybrid query")
|
||||
|
||||
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||||
def to_arrow(
|
||||
self,
|
||||
*,
|
||||
timeout: Optional[timedelta] = None,
|
||||
blob_mode: BlobMode = "lazy",
|
||||
) -> pa.Table:
|
||||
blob_sources = self._validated_blob_sources(blob_mode)
|
||||
# WAL-PK-FUSION: without the fallback, this body is `_run_hybrid`'s.
|
||||
return _wal_hybrid.with_pk_fallback(
|
||||
lambda: self._run_hybrid(timeout=timeout),
|
||||
lambda: self._run_hybrid(
|
||||
timeout=timeout, blob_mode=blob_mode, blob_sources=blob_sources
|
||||
),
|
||||
self._table,
|
||||
fused=lambda: self._pk_fusion is not None,
|
||||
caller_requested_row_id=self._caller_requested_row_id,
|
||||
)
|
||||
|
||||
def _run_hybrid(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||||
def _run_hybrid(
|
||||
self,
|
||||
*,
|
||||
timeout: Optional[timedelta] = None,
|
||||
blob_mode: BlobMode,
|
||||
blob_sources: dict[str, str],
|
||||
) -> pa.Table:
|
||||
self._create_query_builders()
|
||||
fts_query, vector_query = self._fts_query, self._vector_query
|
||||
# WAL-PK-FUSION: without the fallback, both legs always ask for row ids.
|
||||
@@ -2342,18 +2490,33 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
with_row_ids=True,
|
||||
offset=self._offset,
|
||||
)
|
||||
return self._finish_hybrid_results(results)
|
||||
return self._finish_hybrid_results(
|
||||
results, blob_mode=blob_mode, blob_sources=blob_sources
|
||||
)
|
||||
|
||||
def _finish_hybrid_results(self, results: pa.Table) -> pa.Table:
|
||||
def _finish_hybrid_results(
|
||||
self,
|
||||
results: pa.Table,
|
||||
*,
|
||||
blob_mode: BlobMode,
|
||||
blob_sources: dict[str, str],
|
||||
) -> pa.Table:
|
||||
if self._pk_fusion is not None: # WAL-PK-FUSION: delete.
|
||||
return self._pk_fusion.strip(results)
|
||||
return self._apply_arrow_blob_mode(
|
||||
self._pk_fusion.strip(results), blob_mode, blob_sources
|
||||
)
|
||||
if self._user_requested_row_id():
|
||||
return results
|
||||
return self._apply_arrow_blob_mode(results, blob_mode, blob_sources)
|
||||
if self._blob_auto_row_id_enabled():
|
||||
return self._finalize_blob_query_table(results)
|
||||
return self._finalize_blob_query_table(
|
||||
results, blob_mode=blob_mode, blob_sources=blob_sources
|
||||
)
|
||||
if "_rowid" in results.column_names:
|
||||
return results.drop(["_rowid"])
|
||||
return results
|
||||
# Fetch bytes before dropping _rowid.
|
||||
return self._apply_arrow_blob_mode(results, blob_mode, blob_sources).drop(
|
||||
["_rowid"]
|
||||
)
|
||||
return self._apply_arrow_blob_mode(results, blob_mode, blob_sources)
|
||||
|
||||
@staticmethod
|
||||
def _combine_hybrid_results(
|
||||
@@ -2804,6 +2967,10 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
)
|
||||
return "\n".join(results)
|
||||
|
||||
def _blob_output_schema(self) -> pa.Schema:
|
||||
self._create_query_builders()
|
||||
return self._vector_query.output_schema()
|
||||
|
||||
def _create_query_builders(self):
|
||||
"""Set up and configure the vector and FTS query builders."""
|
||||
vector_query, fts_query = self._validate_query(
|
||||
@@ -3064,7 +3231,48 @@ class AsyncQueryBase(object):
|
||||
"""
|
||||
return await self._inner.output_schema()
|
||||
|
||||
async def to_arrow(self, timeout: Optional[timedelta] = None) -> pa.Table:
|
||||
async def _validated_blob_sources(self, blob_mode: BlobMode) -> dict[str, str]:
|
||||
"""Map projected blob v2 outputs to their source columns.
|
||||
|
||||
Bytes mode is checked here, before the query runs.
|
||||
"""
|
||||
validate_blob_mode(blob_mode)
|
||||
# Lazy needs no mapping, and skipping it saves remote tables a request.
|
||||
if blob_mode == "lazy" or self._table is None:
|
||||
return {}
|
||||
schema = await self._table.schema()
|
||||
projection = _query_request_projection(self._inner.to_query_request())
|
||||
blob_sources = blob_v2_projection_sources(schema, projection)
|
||||
if blob_mode != "bytes":
|
||||
return blob_sources
|
||||
reject_unfetchable_blob_columns(schema, projection, blob_sources)
|
||||
if blob_sources and not _supports_local_blob_fetch(
|
||||
self._table, self._table._pushdown_operations
|
||||
):
|
||||
raise _blob_bytes_without_snapshot_error()
|
||||
reject_unmatched_blob_outputs(await self._blob_output_schema(), blob_sources)
|
||||
return blob_sources
|
||||
|
||||
async def _blob_output_schema(self) -> pa.Schema:
|
||||
return await self.output_schema()
|
||||
|
||||
async def _apply_arrow_blob_mode(
|
||||
self, tbl: pa.Table, blob_mode: BlobMode, blob_sources: dict[str, str]
|
||||
) -> pa.Table:
|
||||
if blob_mode == "lazy":
|
||||
return tbl
|
||||
if blob_mode == "descriptions":
|
||||
return strip_auto_row_ids(tbl, blob_sources)
|
||||
return await replace_v2_blob_columns_with_bytes(
|
||||
tbl, blob_sources, self._table.fetch_blobs
|
||||
)
|
||||
|
||||
async def to_arrow(
|
||||
self,
|
||||
timeout: Optional[timedelta] = None,
|
||||
*,
|
||||
blob_mode: BlobMode = "lazy",
|
||||
) -> pa.Table:
|
||||
"""
|
||||
Execute the query and collect the results into an Apache Arrow Table.
|
||||
|
||||
@@ -3078,13 +3286,26 @@ class AsyncQueryBase(object):
|
||||
The maximum time to wait for the query to complete.
|
||||
If not specified, no timeout is applied. If the query does not
|
||||
complete within the specified time, an error will be raised.
|
||||
blob_mode: str, default "lazy"
|
||||
``"lazy"`` returns blob v2 descriptors that carry the row id
|
||||
``fetch_blobs`` reads. ``"descriptions"`` returns them without it.
|
||||
``"bytes"`` returns payload bytes and raises for legacy blob columns
|
||||
and blobs inside lists.
|
||||
Remote tables and namespace query pushdown refuse ``"bytes"``.
|
||||
"""
|
||||
blob_sources = await self._validated_blob_sources(blob_mode)
|
||||
batch_iter = await self.to_batches(timeout=timeout)
|
||||
return self._finalize_blob_query_table(
|
||||
tbl = self._finalize_blob_query_table(
|
||||
pa.Table.from_batches(await batch_iter.read_all(), schema=batch_iter.schema)
|
||||
)
|
||||
return await self._apply_arrow_blob_mode(tbl, blob_mode, blob_sources)
|
||||
|
||||
async def to_list(self, timeout: Optional[timedelta] = None) -> List[dict]:
|
||||
async def to_list(
|
||||
self,
|
||||
timeout: Optional[timedelta] = None,
|
||||
*,
|
||||
blob_mode: BlobMode = "lazy",
|
||||
) -> List[dict]:
|
||||
"""
|
||||
Execute the query and return the results as a list of dictionaries.
|
||||
|
||||
@@ -3098,8 +3319,10 @@ class AsyncQueryBase(object):
|
||||
The maximum time to wait for the query to complete.
|
||||
If not specified, no timeout is applied. If the query does not
|
||||
complete within the specified time, an error will be raised.
|
||||
blob_mode: str, default "lazy"
|
||||
Passed through to :meth:`to_arrow`.
|
||||
"""
|
||||
return (await self.to_arrow(timeout=timeout)).to_pylist()
|
||||
return (await self.to_arrow(timeout=timeout, blob_mode=blob_mode)).to_pylist()
|
||||
|
||||
async def to_pandas(
|
||||
self,
|
||||
@@ -3141,40 +3364,42 @@ 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.
|
||||
blob_mode: str, default "lazy"
|
||||
Controls how blob columns are returned. Remote queries support
|
||||
"descriptions"; "bytes" and "lazy" are not yet supported.
|
||||
Controls how blob columns are returned. ``"lazy"`` returns
|
||||
:class:`~lancedb._blob.BlobFile` handles and needs a plain scan.
|
||||
``"descriptions"`` and ``"bytes"`` work for every query shape. A
|
||||
plain scan without ``timeout`` reads the local dataset through
|
||||
Lance's scanner, which also returns legacy blob columns as bytes.
|
||||
Other shapes use :meth:`to_arrow` with the same ``blob_mode``.
|
||||
Remote queries support ``"descriptions"`` only.
|
||||
**kwargs
|
||||
Forwarded to pyarrow.Table.to_pandas after query execution and
|
||||
optional flattening.
|
||||
"""
|
||||
validate_blob_mode(blob_mode)
|
||||
if self._table is not None and not self._table._inner._is_native():
|
||||
await self._validated_blob_sources(blob_mode)
|
||||
return await self._remote_to_pandas(
|
||||
flatten=flatten, timeout=timeout, blob_mode=blob_mode, **kwargs
|
||||
)
|
||||
|
||||
if hasattr(self._inner, "output_schema"):
|
||||
schema = await self.output_schema()
|
||||
if _blob_mode_requires_native_pandas(blob_mode, schema):
|
||||
native_error = None
|
||||
if (flatten is None or blob_mode == "descriptions") and timeout is None:
|
||||
try:
|
||||
df = await self._plain_scan_to_pandas(
|
||||
blob_mode, flatten=flatten, **kwargs
|
||||
)
|
||||
if df is not None:
|
||||
return df
|
||||
except Exception as err:
|
||||
native_error = err
|
||||
reason = (
|
||||
if hasattr(self._inner, "output_schema") and schema_has_blob_field(
|
||||
await self.output_schema()
|
||||
):
|
||||
if (flatten is None or blob_mode != "lazy") and timeout is None:
|
||||
df = await self._plain_scan_to_pandas(
|
||||
blob_mode, flatten=flatten, **kwargs
|
||||
)
|
||||
if df is not None:
|
||||
return df
|
||||
if blob_mode == "lazy" or _query_needs_lance_scanner(self):
|
||||
raise _unsupported_blob_pandas_error(
|
||||
"this query shape cannot use Lance native pandas conversion"
|
||||
if native_error is None
|
||||
else str(native_error)
|
||||
)
|
||||
raise _unsupported_blob_pandas_error(reason) from native_error
|
||||
|
||||
tbl = flatten_columns(await self.to_arrow(timeout=timeout), flatten)
|
||||
if _blob_mode_requires_native_pandas(blob_mode, tbl.schema):
|
||||
tbl = flatten_columns(
|
||||
await self.to_arrow(timeout=timeout, blob_mode=blob_mode), flatten
|
||||
)
|
||||
if blob_mode == "lazy" and schema_has_blob_field(tbl.schema):
|
||||
raise _unsupported_blob_pandas_error(
|
||||
"this query shape cannot use Lance native pandas conversion"
|
||||
)
|
||||
@@ -3229,6 +3454,13 @@ class AsyncQueryBase(object):
|
||||
if blob_mode == "bytes"
|
||||
else {}
|
||||
)
|
||||
if blob_sources:
|
||||
reject_unfetchable_blob_columns(schema, query.columns, blob_sources)
|
||||
# Remote and in-memory tables have no Lance dataset to scan.
|
||||
if not self._table._inner._is_native() or (
|
||||
get_uri_scheme(await self._table.uri()) == "memory"
|
||||
):
|
||||
return None
|
||||
dataset = await self._table.to_lance()
|
||||
scanner = dataset.scanner(
|
||||
**_scanner_kwargs_for_query(
|
||||
@@ -3243,7 +3475,7 @@ class AsyncQueryBase(object):
|
||||
blob_mode=blob_mode,
|
||||
blob_sources=blob_sources,
|
||||
fetch_blobs=self._table.fetch_blobs,
|
||||
strip_auto_row_id=blob_auto_row_id,
|
||||
strip_auto_row_id=self._with_row_id is not True,
|
||||
flatten=flatten,
|
||||
**kwargs,
|
||||
)
|
||||
@@ -4083,6 +4315,10 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
||||
|
||||
return self
|
||||
|
||||
async def _blob_output_schema(self) -> pa.Schema:
|
||||
vector_query = AsyncVectorQuery(self._inner.to_vector_query(), self._table)
|
||||
return await vector_query.output_schema()
|
||||
|
||||
async def _create_child_queries(
|
||||
self,
|
||||
) -> Tuple["AsyncFTSQuery", "AsyncVectorQuery", int, int]:
|
||||
@@ -4444,7 +4680,12 @@ class BaseQueryBuilder(object):
|
||||
|
||||
return pa.RecordBatchReader.from_batches(async_reader.schema, iter_sync())
|
||||
|
||||
def to_arrow(self, timeout: Optional[timedelta] = None) -> pa.Table:
|
||||
def to_arrow(
|
||||
self,
|
||||
timeout: Optional[timedelta] = None,
|
||||
*,
|
||||
blob_mode: BlobMode = "lazy",
|
||||
) -> pa.Table:
|
||||
"""
|
||||
Execute the query and collect the results into an Apache Arrow Table.
|
||||
|
||||
@@ -4458,10 +4699,20 @@ class BaseQueryBuilder(object):
|
||||
The maximum time to wait for the query to complete.
|
||||
If not specified, no timeout is applied. If the query does not
|
||||
complete within the specified time, an error will be raised.
|
||||
blob_mode: str, default "lazy"
|
||||
``"lazy"`` returns blob v2 descriptors that carry the row id
|
||||
``fetch_blobs`` reads. ``"descriptions"`` returns them without it.
|
||||
``"bytes"`` returns payload bytes and raises for legacy blob columns
|
||||
and blobs inside lists.
|
||||
"""
|
||||
return LOOP.run(self._inner.to_arrow(timeout))
|
||||
return LOOP.run(self._inner.to_arrow(timeout, blob_mode=blob_mode))
|
||||
|
||||
def to_list(self, timeout: Optional[timedelta] = None) -> List[dict]:
|
||||
def to_list(
|
||||
self,
|
||||
timeout: Optional[timedelta] = None,
|
||||
*,
|
||||
blob_mode: BlobMode = "lazy",
|
||||
) -> List[dict]:
|
||||
"""
|
||||
Execute the query and return the results as a list of dictionaries.
|
||||
|
||||
@@ -4475,8 +4726,10 @@ class BaseQueryBuilder(object):
|
||||
The maximum time to wait for the query to complete.
|
||||
If not specified, no timeout is applied. If the query does not
|
||||
complete within the specified time, an error will be raised.
|
||||
blob_mode: str, default "lazy"
|
||||
Passed through to :meth:`to_arrow`.
|
||||
"""
|
||||
return LOOP.run(self._inner.to_list(timeout))
|
||||
return LOOP.run(self._inner.to_list(timeout, blob_mode=blob_mode))
|
||||
|
||||
def to_pandas(
|
||||
self,
|
||||
@@ -4518,7 +4771,7 @@ class BaseQueryBuilder(object):
|
||||
If not specified, no timeout is applied. If the query does not
|
||||
complete within the specified time, an error will be raised.
|
||||
blob_mode: str, default "lazy"
|
||||
Controls how blob columns are returned for plain scan queries.
|
||||
Passed through to :meth:`AsyncQueryBase.to_pandas`.
|
||||
**kwargs
|
||||
Forwarded to pyarrow.Table.to_pandas after query execution and
|
||||
optional flattening.
|
||||
|
||||
@@ -57,6 +57,7 @@ from lancedb.table import IndexConfigType, KNOWN_METRICS
|
||||
import pyarrow as pa
|
||||
|
||||
from lancedb.common import DATA, VEC, VECTOR_COLUMN_NAME
|
||||
from lancedb._blob import validate_blob_mode
|
||||
from lancedb.merge import LanceMergeInsertBuilder
|
||||
from lancedb.embeddings import EmbeddingFunctionRegistry
|
||||
from lancedb.table import _normalize_progress
|
||||
@@ -239,12 +240,14 @@ class RemoteTable(Table):
|
||||
"""List all versions of the table"""
|
||||
return LOOP.run(self._table.list_versions())
|
||||
|
||||
def to_arrow(self) -> pa.Table:
|
||||
def to_arrow(self, *, blob_mode: BlobMode = "lazy") -> pa.Table:
|
||||
"""to_arrow() is not supported for remote tables."""
|
||||
validate_blob_mode(blob_mode)
|
||||
raise NotImplementedError("to_arrow() is not supported for remote tables.")
|
||||
|
||||
def to_pandas(self, blob_mode: BlobMode = "lazy", **kwargs):
|
||||
"""to_pandas() is not supported for remote tables."""
|
||||
validate_blob_mode(blob_mode)
|
||||
raise NotImplementedError("to_pandas() is not supported for remote tables.")
|
||||
|
||||
def checkout(self, version: Union[int, str]):
|
||||
|
||||
@@ -1399,9 +1399,17 @@ class Table(ABC):
|
||||
return self.to_arrow().to_pandas(**kwargs)
|
||||
|
||||
@abstractmethod
|
||||
def to_arrow(self) -> pa.Table:
|
||||
def to_arrow(self, *, blob_mode: BlobMode = "lazy") -> pa.Table:
|
||||
"""Return the table as a pyarrow Table.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
blob_mode: str, default "lazy"
|
||||
``"lazy"`` returns blob v2 descriptors that carry the row id
|
||||
``fetch_blobs`` reads. ``"descriptions"`` returns them without it.
|
||||
``"bytes"`` returns payload bytes and raises for legacy blob columns
|
||||
and blobs inside lists.
|
||||
|
||||
Returns
|
||||
-------
|
||||
pa.Table
|
||||
@@ -2961,6 +2969,8 @@ class LanceTable(Table):
|
||||
managed_versioning=managed_versioning,
|
||||
)
|
||||
)
|
||||
# Take queries read the pushdown set from the inner table.
|
||||
self._table._pushdown_operations = self._pushdown_operations
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
@@ -3339,12 +3349,23 @@ class LanceTable(Table):
|
||||
|
||||
return self.to_lance().to_pandas(blob_mode=blob_mode, **kwargs)
|
||||
|
||||
def to_arrow(self) -> pa.Table:
|
||||
def to_arrow(self, *, blob_mode: BlobMode = "lazy") -> pa.Table:
|
||||
"""Return the table as a pyarrow Table.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
blob_mode: str, default "lazy"
|
||||
``"lazy"`` returns blob v2 descriptors that carry the row id
|
||||
``fetch_blobs`` reads. ``"descriptions"`` returns them without it.
|
||||
``"bytes"`` returns payload bytes and raises for legacy blob columns
|
||||
and blobs inside lists.
|
||||
|
||||
Returns
|
||||
-------
|
||||
pa.Table"""
|
||||
validate_blob_mode(blob_mode)
|
||||
if blob_mode != "lazy" and schema_has_blob_field(self.schema):
|
||||
return self.search().to_arrow(blob_mode=blob_mode)
|
||||
if (
|
||||
_should_push_down_query_table(
|
||||
self._namespace_client, self._pushdown_operations
|
||||
@@ -5788,22 +5809,37 @@ class AsyncTable:
|
||||
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:
|
||||
async def to_arrow(self, *, blob_mode: BlobMode = "lazy") -> pa.Table:
|
||||
"""Return the table as a pyarrow Table.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
blob_mode: str, default "lazy"
|
||||
``"lazy"`` returns blob v2 descriptors that carry the row id
|
||||
``fetch_blobs`` reads. ``"descriptions"`` returns them without it.
|
||||
``"bytes"`` returns payload bytes and raises for legacy blob columns
|
||||
and blobs inside lists.
|
||||
|
||||
Returns
|
||||
-------
|
||||
pa.Table
|
||||
"""
|
||||
validate_blob_mode(blob_mode)
|
||||
schema = await self.schema()
|
||||
fetches_bytes = blob_mode == "bytes" and schema_has_blob_field(schema)
|
||||
if (
|
||||
_should_push_down_query_table(
|
||||
not fetches_bytes
|
||||
and _should_push_down_query_table(
|
||||
self._namespace_client, self._pushdown_operations
|
||||
)
|
||||
and not self._route_pushdown_to_rust
|
||||
):
|
||||
return (await self._execute_query(Query())).read_all()
|
||||
|
||||
return await self.query().to_arrow()
|
||||
tbl = (await self._execute_query(Query())).read_all()
|
||||
if blob_mode == "descriptions":
|
||||
tbl = strip_auto_row_ids(tbl, row_addressable_blob_v2_paths(schema))
|
||||
return tbl
|
||||
# The query builder refuses bytes under pushdown before running the query.
|
||||
return await self.query().to_arrow(blob_mode=blob_mode)
|
||||
|
||||
async def create_index(
|
||||
self,
|
||||
|
||||
@@ -22,6 +22,7 @@ from lancedb._blob import (
|
||||
)
|
||||
from lancedb.expr import col
|
||||
from lancedb.index import FTS
|
||||
from lancedb.rerankers import Reranker
|
||||
from lancedb.schema import blob_column_paths, blob_v2_column_paths
|
||||
|
||||
|
||||
@@ -411,6 +412,158 @@ def test_nested_blob_to_pandas_bytes(tmp_path):
|
||||
].tolist() == [b"hello", None, None]
|
||||
|
||||
|
||||
def test_to_arrow_and_to_list_blob_mode(tmp_path):
|
||||
table = _nested_blob_table(lancedb.connect(tmp_path), "arrow_blob_mode")
|
||||
expected = [
|
||||
{"name": "first", "blob": b"hello"},
|
||||
{"name": "second", "blob": None},
|
||||
None,
|
||||
]
|
||||
|
||||
assert table.to_arrow(blob_mode="bytes")["info"].to_pylist() == expected
|
||||
assert table.search().to_arrow(blob_mode="bytes")["info"].to_pylist() == expected
|
||||
assert table.search().to_list(blob_mode="bytes") == [
|
||||
{"id": 0, "info": expected[0]},
|
||||
{"id": 1, "info": expected[1]},
|
||||
{"id": 2, "info": expected[2]},
|
||||
]
|
||||
assert table.search().select({"b": "info.blob"}).to_list(blob_mode="bytes") == [
|
||||
{"b": b"hello"},
|
||||
{"b": None},
|
||||
{"b": None},
|
||||
]
|
||||
assert table.search().select(["id", "info.blob"]).to_list(blob_mode="bytes") == [
|
||||
{"id": 0, "info.blob": b"hello"},
|
||||
{"id": 1, "info.blob": None},
|
||||
{"id": 2, "info.blob": None},
|
||||
]
|
||||
assert table.search().select({"x.y": "info"}).to_list(blob_mode="bytes") == [
|
||||
{"x.y": expected[0]},
|
||||
{"x.y": expected[1]},
|
||||
{"x.y": expected[2]},
|
||||
]
|
||||
assert [
|
||||
row["info"] for row in table.search().select([]).to_list(blob_mode="bytes")
|
||||
] == expected
|
||||
assert table.search().select([]).to_pandas(blob_mode="bytes")["info"].tolist() == (
|
||||
table.search().select([]).to_arrow(blob_mode="bytes")["info"].to_pylist()
|
||||
)
|
||||
assert "_rowid" not in table.search().select([]).to_arrow().column_names
|
||||
# A flat non-blob column named like the nested blob path keeps its values.
|
||||
shadowed = table.search().select({"info": "info", "info.blob": "id"})
|
||||
assert [row["info.blob"] for row in shadowed.to_list()] == [0, 1, 2]
|
||||
assert [row["info"] for row in shadowed.to_list(blob_mode="bytes")] == expected
|
||||
with pytest.raises(ValueError, match="both produce blob column 'info.blob'"):
|
||||
table.search().select(["info", "info.blob"]).to_list(blob_mode="bytes")
|
||||
|
||||
default_blob = table.search().select({"b": "info.blob"}).to_list()[0]["b"]
|
||||
assert isinstance(default_blob, dict)
|
||||
lazy = table.to_arrow(blob_mode="lazy")["info"].to_pylist()
|
||||
assert table.to_arrow()["info"].to_pylist() == lazy
|
||||
lazy_descriptor = lazy[0]["blob"]
|
||||
assert "_lance_row_id" in lazy_descriptor
|
||||
lazy_descriptor.pop("_lance_row_id")
|
||||
descriptions = table.to_arrow(blob_mode="descriptions")["info"].to_pylist()
|
||||
assert descriptions[0]["blob"] == lazy_descriptor
|
||||
assert (
|
||||
table.search().to_list(blob_mode="descriptions")[0]["info"]["blob"]
|
||||
== lazy_descriptor
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="blob_mode must be one of"):
|
||||
table.to_arrow(blob_mode="handles")
|
||||
with pytest.raises(ValueError, match="blob_mode must be one of"):
|
||||
table.search().to_list(blob_mode="handles")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_to_arrow_and_to_list_blob_mode(tmp_path):
|
||||
db = await lancedb.connect_async(tmp_path)
|
||||
table = await db.create_table(
|
||||
"async_arrow_blob_mode",
|
||||
schema=pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field(
|
||||
"info",
|
||||
pa.struct([pa.field("name", pa.string()), lancedb.blob("blob")]),
|
||||
),
|
||||
]
|
||||
),
|
||||
)
|
||||
await table.add(
|
||||
pa.table(
|
||||
{
|
||||
"id": pa.array([0, 1]),
|
||||
"info": pa.StructArray.from_arrays(
|
||||
[
|
||||
pa.array(["first", "second"]),
|
||||
pa.array([b"hello", None], type=pa.large_binary()),
|
||||
],
|
||||
names=["name", "blob"],
|
||||
),
|
||||
}
|
||||
)
|
||||
)
|
||||
expected = [
|
||||
{"name": "first", "blob": b"hello"},
|
||||
{"name": "second", "blob": None},
|
||||
]
|
||||
|
||||
assert (await table.to_arrow(blob_mode="bytes"))["info"].to_pylist() == expected
|
||||
assert (await table.query().to_arrow(blob_mode="bytes"))[
|
||||
"info"
|
||||
].to_pylist() == expected
|
||||
assert await table.query().to_list(blob_mode="bytes") == [
|
||||
{"id": 0, "info": expected[0]},
|
||||
{"id": 1, "info": expected[1]},
|
||||
]
|
||||
default_blob = (await table.query().select({"b": "info.blob"}).to_list())[0]["b"]
|
||||
assert isinstance(default_blob, dict)
|
||||
descriptions = await table.to_arrow(blob_mode="descriptions")
|
||||
assert (
|
||||
"_lance_row_id"
|
||||
not in descriptions.schema.field("info").type.field("blob").type.names
|
||||
)
|
||||
|
||||
|
||||
def test_vector_query_to_arrow_blob_mode_bytes(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("vector", pa.list_(pa.float32(), 2)),
|
||||
lancedb.blob("blob"),
|
||||
]
|
||||
)
|
||||
table = db.create_table("vector_arrow_blob", schema=schema)
|
||||
table.add(
|
||||
[
|
||||
{"id": 1, "vector": [1.0, 0.0], "blob": b"one"},
|
||||
{"id": 2, "vector": [0.0, 1.0], "blob": b"two"},
|
||||
]
|
||||
)
|
||||
table.add([{"id": 3, "vector": [0.9, 0.1], "blob": b"three"}])
|
||||
table.add([{"id": 4, "vector": [0.1, 0.9], "blob": None}])
|
||||
|
||||
hits = table.search([0.0, 1.0]).limit(4).to_arrow(blob_mode="bytes").to_pylist()
|
||||
assert [(hit["id"], hit["blob"]) for hit in hits] == [
|
||||
(2, b"two"),
|
||||
(4, None),
|
||||
(3, b"three"),
|
||||
(1, b"one"),
|
||||
]
|
||||
arrow = table.search([1.0, 0.0]).limit(1).to_arrow(blob_mode="bytes")
|
||||
assert arrow["blob"].to_pylist() == [b"one"]
|
||||
assert table.search([1.0, 0.0]).limit(1).to_list(blob_mode="bytes")[0]["blob"] == (
|
||||
b"one"
|
||||
)
|
||||
assert isinstance(
|
||||
table.search([1.0, 0.0]).limit(1).to_list()[0]["blob"],
|
||||
dict,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_nested_blob_to_pandas_bytes_async(tmp_path):
|
||||
db = await lancedb.connect_async(tmp_path)
|
||||
@@ -547,6 +700,19 @@ async def test_async_table_to_pandas_descriptions_mode_omits_row_id():
|
||||
assert set(descriptor.keys()) == {"kind", "position", "size", "blob_id", "blob_uri"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_calculated_select_does_not_read_every_blob_column():
|
||||
db = await lancedb.connect_async("memory:///calculated_select")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = await db.create_table("calculated_select", schema=schema)
|
||||
await table.add([{"id": 1, "image": b"hello"}])
|
||||
|
||||
hits = await table.query().select({"next_id": col("id") + 1}).to_arrow()
|
||||
|
||||
assert hits.column_names == ["next_id"]
|
||||
assert hits["next_id"].to_pylist() == [2]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_typed_blob_projection_preserves_source_column():
|
||||
db = await lancedb.connect_async("memory:///typed_blob_projection")
|
||||
@@ -561,6 +727,185 @@ async def test_async_typed_blob_projection_preserves_source_column():
|
||||
assert blobs.to_pylist() == [b"alpha"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("shape", ["scan", "hybrid"])
|
||||
async def test_async_bytes_mode_rejects_blob_spelling_it_cannot_match(shape):
|
||||
db = await lancedb.connect_async("memory:///")
|
||||
table = await _hybrid_blob_table_async(db)
|
||||
query = table.query()
|
||||
if shape == "hybrid":
|
||||
query = query.nearest_to([1.0, 0.0]).nearest_to_text("hello")
|
||||
|
||||
with pytest.raises(ValueError, match=r"could not match result columns \['pic'\]"):
|
||||
await query.select({"id": "id", "pic": "`image`"}).to_arrow(blob_mode="bytes")
|
||||
|
||||
|
||||
def test_bytes_mode_keeps_struct_with_blob_descriptor_field_names():
|
||||
descriptor_like = pa.struct(
|
||||
[
|
||||
pa.field("kind", pa.uint8()),
|
||||
pa.field("position", pa.uint64()),
|
||||
pa.field("size", pa.uint64()),
|
||||
pa.field("blob_id", pa.uint32()),
|
||||
pa.field("blob_uri", pa.utf8()),
|
||||
]
|
||||
)
|
||||
meta = {"kind": 1, "position": 2, "size": 3, "blob_id": 4, "blob_uri": "u"}
|
||||
table = lancedb.connect("memory:///").create_table(
|
||||
"descriptor_like",
|
||||
schema=pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
lancedb.blob("image"),
|
||||
pa.field("meta", descriptor_like),
|
||||
]
|
||||
),
|
||||
)
|
||||
table.add([{"id": 1, "image": b"one", "meta": meta}])
|
||||
|
||||
assert table.search().to_list(blob_mode="bytes") == [
|
||||
{"id": 1, "image": b"one", "meta": meta}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("source", ["`image`", "IMAGE"])
|
||||
def test_bytes_mode_rejects_blob_spelling_it_cannot_match(source):
|
||||
table = _blob_table("blob_spelling", [{"id": 1, "image": b"one"}])
|
||||
|
||||
with pytest.raises(ValueError, match="could not match result columns"):
|
||||
table.search().select({"id": "id", "pic": source}).to_arrow(blob_mode="bytes")
|
||||
|
||||
|
||||
def test_to_pandas_does_not_drop_fragment_selection_on_fallback():
|
||||
table = _blob_table("blob_fragments", [{"id": 1, "image": b"one"}])
|
||||
table.add([{"id": 2, "image": b"two"}])
|
||||
|
||||
query = table.search().fragment_ids([0])
|
||||
with pytest.raises(RuntimeError, match="native pandas conversion"):
|
||||
query.to_pandas(blob_mode="bytes")
|
||||
with pytest.raises(RuntimeError, match="native pandas conversion"):
|
||||
table.search().with_row_address(True).to_pandas(blob_mode="descriptions")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("blob_mode", ["lazy", "descriptions", "bytes"])
|
||||
def test_to_pandas_reports_lance_dataset_failures(tmp_path, blob_mode):
|
||||
table = lancedb.connect(tmp_path).create_table(
|
||||
"dataset_failure",
|
||||
schema=pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]),
|
||||
)
|
||||
table.add([{"id": 1, "image": b"one"}])
|
||||
|
||||
def reject_credentials():
|
||||
raise PermissionError("storage credentials rejected")
|
||||
|
||||
table.to_lance = reject_credentials
|
||||
with pytest.raises(PermissionError, match="storage credentials rejected"):
|
||||
table.search().to_pandas(blob_mode=blob_mode)
|
||||
|
||||
|
||||
def test_to_pandas_keeps_to_arrow_overrides_without_blob_mode():
|
||||
from lancedb.query import LanceEmptyQueryBuilder
|
||||
|
||||
class ToArrowWithoutBlobMode(LanceEmptyQueryBuilder):
|
||||
def to_arrow(self, *, timeout=None):
|
||||
return super().to_arrow(timeout=timeout)
|
||||
|
||||
table = lancedb.connect("memory:///").create_table("no_blob", data=[{"id": 1}])
|
||||
query = ToArrowWithoutBlobMode(table)
|
||||
assert query.to_list(blob_mode="bytes") == [{"id": 1}]
|
||||
assert query.to_pandas(blob_mode="descriptions")["id"].tolist() == [1]
|
||||
with pytest.raises(ValueError, match="blob_mode must be one of"):
|
||||
query.to_list(blob_mode="handles")
|
||||
with pytest.raises(ValueError, match="blob_mode must be one of"):
|
||||
table.search().to_list(blob_mode="handles")
|
||||
|
||||
|
||||
def test_take_queries_honor_blob_mode():
|
||||
table = _blob_table(
|
||||
"take_blob_mode", [{"id": 1, "image": b"one"}, {"id": 2, "image": b"two"}]
|
||||
)
|
||||
|
||||
hits = table.take_offsets([1, 0]).to_arrow(blob_mode="bytes")
|
||||
assert dict(zip(hits["id"].to_pylist(), hits["image"].to_pylist())) == {
|
||||
1: b"one",
|
||||
2: b"two",
|
||||
}
|
||||
assert [
|
||||
row["image"] for row in table.take_row_ids([0]).to_list(blob_mode="bytes")
|
||||
] == [b"one"]
|
||||
assert table.take_offsets([0]).to_pandas(blob_mode="bytes")["image"].tolist() == [
|
||||
b"one"
|
||||
]
|
||||
description = table.take_row_ids([0]).to_list(blob_mode="descriptions")[0]["image"]
|
||||
assert "_lance_row_id" not in description
|
||||
|
||||
|
||||
def test_bytes_mode_follows_swapped_blob_aliases():
|
||||
db = lancedb.connect("memory:///")
|
||||
table = db.create_table(
|
||||
"swapped_aliases",
|
||||
schema=pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("caption", pa.utf8()),
|
||||
lancedb.blob("image"),
|
||||
lancedb.blob("thumb"),
|
||||
]
|
||||
),
|
||||
)
|
||||
table.add([{"id": 0, "caption": "cap", "image": b"img", "thumb": b"th"}])
|
||||
|
||||
swapped = table.search().select({"id": "id", "image": "thumb", "thumb": "image"})
|
||||
assert swapped.to_list(blob_mode="bytes") == [
|
||||
{"id": 0, "image": b"th", "thumb": b"img"}
|
||||
]
|
||||
over_text = table.search().select({"id": "id", "image": "caption", "pic": "image"})
|
||||
assert over_text.to_list(blob_mode="bytes") == [
|
||||
{"id": 0, "image": "cap", "pic": b"img"}
|
||||
]
|
||||
|
||||
|
||||
def test_bytes_mode_refused_under_namespace_query_pushdown(tmp_path):
|
||||
db = lancedb.connect_namespace(
|
||||
"dir",
|
||||
{"root": str(tmp_path)},
|
||||
namespace_client_pushdown_operations=["QueryTable"],
|
||||
)
|
||||
table = db.create_table(
|
||||
"pushdown_blob",
|
||||
schema=pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]),
|
||||
)
|
||||
table.add([{"id": 1, "image": b"one"}])
|
||||
|
||||
with pytest.raises(NotImplementedError, match="stable table snapshot"):
|
||||
table.to_arrow(blob_mode="bytes")
|
||||
with pytest.raises(NotImplementedError, match="stable table snapshot"):
|
||||
table.search().to_list(blob_mode="bytes")
|
||||
with pytest.raises(NotImplementedError, match="stable table snapshot"):
|
||||
table.take_row_ids([0]).to_list(blob_mode="bytes")
|
||||
descriptor = table.to_arrow(blob_mode="descriptions")["image"].to_pylist()[0]
|
||||
assert "_lance_row_id" not in descriptor
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_bytes_mode_refused_under_namespace_query_pushdown(tmp_path):
|
||||
db = lancedb.connect_namespace_async(
|
||||
"dir",
|
||||
{"root": str(tmp_path)},
|
||||
namespace_client_pushdown_operations=["QueryTable"],
|
||||
)
|
||||
table = await db.create_table(
|
||||
"pushdown_blob",
|
||||
schema=pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]),
|
||||
)
|
||||
await table.add([{"id": 1, "image": b"one"}])
|
||||
|
||||
with pytest.raises(NotImplementedError, match="stable table snapshot"):
|
||||
await table.to_arrow(blob_mode="bytes")
|
||||
with pytest.raises(NotImplementedError, match="stable table snapshot"):
|
||||
await table.query().to_list(blob_mode="bytes")
|
||||
|
||||
|
||||
def test_fetch_blobs_round_trip():
|
||||
table = _blob_table(
|
||||
"round_trip",
|
||||
@@ -937,6 +1282,17 @@ def test_list_blob_column_queries_as_raw_descriptors():
|
||||
table.fetch_blob_ranges(column, [(0, 0, 1)])
|
||||
|
||||
|
||||
def test_bytes_mode_rejects_blobs_inside_lists():
|
||||
table = _list_blob_table("list_blob_bytes")
|
||||
|
||||
with pytest.raises(ValueError, match=r"\['images.image'\].*blobs inside lists"):
|
||||
table.to_arrow(blob_mode="bytes")
|
||||
with pytest.raises(ValueError, match="blobs inside lists"):
|
||||
table.search().select({"pics": "images"}).to_list(blob_mode="bytes")
|
||||
descriptions = table.to_arrow(blob_mode="descriptions")["images"].to_pylist()
|
||||
assert descriptions[0][0]["size"] == len(b"before")
|
||||
|
||||
|
||||
def test_row_addressable_paths_exclude_list_children():
|
||||
from lancedb.schema import row_addressable_blob_v2_paths
|
||||
|
||||
@@ -1261,6 +1617,11 @@ def test_fetch_blobs_nested_path():
|
||||
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"]
|
||||
|
||||
# The source path's parent name now holds a string column in the hits.
|
||||
hits = table.search().select({"id": "id", "info": "info.name", "b": "info.blob"})
|
||||
blobs = table.fetch_blobs("info.blob", hits.to_arrow())
|
||||
assert blobs.to_pylist() == [b"nested-alpha", b"nested-beta"]
|
||||
|
||||
|
||||
def test_fetch_blob_files_lazy_read():
|
||||
payload = b"lazy-read" * 100
|
||||
@@ -1438,6 +1799,122 @@ def test_blob_v2_hybrid_fetch_blobs():
|
||||
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"}
|
||||
|
||||
|
||||
def _blob_query_shape(table, shape):
|
||||
if shape == "scan":
|
||||
return table.search().where("id < 3")
|
||||
if shape == "vector":
|
||||
return table.search([1.0, 0.0])
|
||||
if shape == "fts":
|
||||
return table.search("hello", query_type="fts")
|
||||
return table.search(query_type="hybrid").vector([1.0, 0.0]).text("hello")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("with_row_id", [None, True, False])
|
||||
@pytest.mark.parametrize("blob_mode", ["bytes", "descriptions"])
|
||||
@pytest.mark.parametrize("shape", ["scan", "vector", "fts", "hybrid"])
|
||||
def test_blob_mode_for_every_query_shape(shape, blob_mode, with_row_id):
|
||||
table = _hybrid_blob_table(lancedb.connect("memory:///"))
|
||||
query = _blob_query_shape(table, shape).select(["id", "image"]).limit(2)
|
||||
if with_row_id is not None:
|
||||
query = query.with_row_id(with_row_id)
|
||||
|
||||
hits = query.to_arrow(blob_mode=blob_mode)
|
||||
|
||||
assert ("_rowid" in hits.column_names) == (with_row_id is True)
|
||||
images = {row["id"]: row["image"] for row in hits.to_pylist()}
|
||||
if blob_mode == "bytes":
|
||||
assert images == {1: b"alpha", 2: b"beta"}
|
||||
else:
|
||||
assert sorted(images) == [1, 2]
|
||||
assert all("_lance_row_id" not in image for image in images.values())
|
||||
df = query.to_pandas(blob_mode=blob_mode)
|
||||
assert dict(zip(df["id"], df["image"])) == images
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("shape", ["scan", "vector", "fts", "hybrid"])
|
||||
async def test_async_blob_bytes_for_every_query_shape(shape):
|
||||
db = await lancedb.connect_async("memory:///")
|
||||
table = await _hybrid_blob_table_async(db)
|
||||
query = table.query()
|
||||
if shape in ("vector", "hybrid"):
|
||||
query = query.nearest_to([1.0, 0.0])
|
||||
if shape in ("fts", "hybrid"):
|
||||
query = query.nearest_to_text("hello")
|
||||
if shape == "scan":
|
||||
query = query.where("id < 3")
|
||||
query = query.select(["id", "image"]).limit(2)
|
||||
|
||||
hits = await query.to_arrow(blob_mode="bytes")
|
||||
|
||||
assert {row["id"]: row["image"] for row in hits.to_pylist()} == {
|
||||
1: b"alpha",
|
||||
2: b"beta",
|
||||
}
|
||||
df = await query.to_pandas(blob_mode="bytes")
|
||||
assert dict(zip(df["id"], df["image"])) == {1: b"alpha", 2: b"beta"}
|
||||
|
||||
|
||||
class _KeepAllScoresReranker(Reranker):
|
||||
"""Rebuilds the merged table like model rerankers with return_score="all"."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(return_score="all")
|
||||
|
||||
def rerank_hybrid(self, query, vector_results, fts_results):
|
||||
combined = self._merge_and_keep_scores(vector_results, fts_results)
|
||||
scores = pa.array([1.0] * len(combined), type=pa.float32())
|
||||
return combined.append_column("_relevance_score", scores)
|
||||
|
||||
|
||||
def test_bytes_mode_rejects_legacy_blob_after_reranker_drops_metadata(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table(
|
||||
"legacy_rerank",
|
||||
pa.table(
|
||||
{
|
||||
"id": pa.array([1, 2], pa.int64()),
|
||||
"text": ["hello alpha", "hello beta"],
|
||||
"vector": pa.array(
|
||||
[[1.0, 0.0], [0.9, 0.1]], type=pa.list_(pa.float32(), 2)
|
||||
),
|
||||
"blob": pa.array([b"one", b"two"], pa.large_binary()),
|
||||
},
|
||||
schema=_HYBRID_BLOB_SCHEMA.remove(3).append(
|
||||
pa.field(
|
||||
"blob", pa.large_binary(), metadata={"lance-encoding:blob": "true"}
|
||||
)
|
||||
),
|
||||
),
|
||||
storage_options={"new_table_data_storage_version": "2.1"},
|
||||
)
|
||||
table.create_index("text", config=FTS(with_position=False))
|
||||
query = (
|
||||
table.search(query_type="hybrid")
|
||||
.vector([1.0, 0.0])
|
||||
.text("hello")
|
||||
.rerank(_KeepAllScoresReranker())
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match=r"\['blob'\].*legacy blob columns"):
|
||||
query.to_arrow(blob_mode="bytes")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("source", ["`image`", "IMAGE"])
|
||||
def test_bytes_mode_rejects_blob_spelling_after_reranker_drops_metadata(source):
|
||||
table = _hybrid_blob_table(lancedb.connect("memory:///"))
|
||||
query = (
|
||||
table.search(query_type="hybrid")
|
||||
.vector([1.0, 0.0])
|
||||
.text("hello")
|
||||
.select({"id": "id", "pic": source})
|
||||
.rerank(_KeepAllScoresReranker())
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match=r"could not match result columns \['pic'\]"):
|
||||
query.to_arrow(blob_mode="bytes")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_blob_v2_hybrid_fetch_blobs_async():
|
||||
db = await lancedb.connect_async("memory:///hybrid_blob_fetch_async")
|
||||
|
||||
@@ -509,18 +509,30 @@ def test_vector_query_to_pandas_blob_mode_requires_native_path(tmp_db):
|
||||
)
|
||||
|
||||
|
||||
def test_vector_query_to_pandas_blob_descriptions_requires_plain_scan(tmp_db):
|
||||
def test_vector_query_to_pandas_legacy_blob_descriptions_and_bytes(tmp_db):
|
||||
pytest.importorskip("lance")
|
||||
table = tmp_db.create_table(
|
||||
"test_vector_query_blob_descriptions",
|
||||
_blob_query_data(),
|
||||
storage_options=LEGACY_BLOB_STORAGE_OPTIONS,
|
||||
)
|
||||
query = table.search([1.0, 0.0]).select(["blob", "vector"]).limit(1)
|
||||
|
||||
with pytest.raises(RuntimeError, match="plain scan query"):
|
||||
table.search([1.0, 0.0]).select(["blob", "vector"]).limit(1).to_pandas(
|
||||
blob_mode="descriptions"
|
||||
)
|
||||
df = query.to_pandas(blob_mode="descriptions")
|
||||
assert df["blob"].tolist() == [{"position": 0, "size": 3}]
|
||||
|
||||
with pytest.raises(ValueError, match="legacy blob columns"):
|
||||
query.to_pandas(blob_mode="bytes")
|
||||
with pytest.raises(ValueError, match="legacy blob columns"):
|
||||
query.to_arrow(blob_mode="bytes")
|
||||
with pytest.raises(ValueError, match="legacy blob columns"):
|
||||
table.to_arrow(blob_mode="bytes")
|
||||
assert table.to_pandas(blob_mode="bytes")["blob"].tolist() == [
|
||||
b"one",
|
||||
b"two",
|
||||
b"three",
|
||||
b"four",
|
||||
]
|
||||
|
||||
|
||||
def test_order_by_plain_query(mem_db):
|
||||
|
||||
@@ -3469,6 +3469,45 @@ async def test_async_remote_blob_pandas_rejects_unpinned_fetch(monkeypatch, blob
|
||||
await table.to_pandas(**pandas_kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["to_arrow", "to_list"])
|
||||
def test_remote_search_bytes_refused_before_query(monkeypatch, method):
|
||||
with blob_remote_table() as table:
|
||||
query = table.search().select(["id", "image"])
|
||||
descriptions = getattr(query, method)(blob_mode="descriptions")
|
||||
images = (
|
||||
descriptions["image"].to_pylist()
|
||||
if method == "to_arrow"
|
||||
else [row["image"] for row in descriptions]
|
||||
)
|
||||
assert "_lance_row_id" not in images[0]
|
||||
|
||||
def fail_query(*args, **kwargs):
|
||||
raise AssertionError("bytes must be refused before the query runs")
|
||||
|
||||
monkeypatch.setattr(table, "_execute_query", fail_query)
|
||||
with pytest.raises(NotImplementedError, match="stable table snapshot"):
|
||||
getattr(query, method)(blob_mode="bytes")
|
||||
with pytest.raises(ValueError, match="blob_mode must be one of"):
|
||||
table.to_pandas(blob_mode="handles")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_remote_to_arrow_bytes_refused_before_query(monkeypatch):
|
||||
with blob_remote_table() as remote_table:
|
||||
table = remote_table._table
|
||||
|
||||
async def fail_fetch(*args, **kwargs):
|
||||
raise AssertionError("blob fetch must not use an unpinned remote table")
|
||||
|
||||
monkeypatch.setattr(table, "fetch_blobs", fail_fetch)
|
||||
query = table.query().select(["id", "image"])
|
||||
monkeypatch.setattr(query, "to_batches", fail_fetch)
|
||||
with pytest.raises(NotImplementedError, match="stable table snapshot"):
|
||||
await query.to_arrow(blob_mode="bytes")
|
||||
with pytest.raises(NotImplementedError, match="stable table snapshot"):
|
||||
await table.to_arrow(blob_mode="bytes")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("blob_mode", ["bytes", "lazy"])
|
||||
async def test_async_remote_nonblob_projection_to_pandas(blob_mode):
|
||||
|
||||
@@ -119,7 +119,9 @@ def test_wal_hybrid_fuses_on_the_key_and_drops_it_again():
|
||||
limit=10,
|
||||
with_row_ids=True,
|
||||
)
|
||||
results = builder._finish_hybrid_results(combined)
|
||||
results = builder._finish_hybrid_results(
|
||||
combined, blob_mode="lazy", blob_sources={}
|
||||
)
|
||||
|
||||
assert "_rowid" not in results.column_names, "the surrogate is internal"
|
||||
assert "id" not in results.column_names, "the caller selected only `text`"
|
||||
|
||||
Reference in New Issue
Block a user