mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-17 01:02:28 +00:00
Merge main and retain sliced Arrow validity fix
This commit is contained in:
+222
-66
@@ -6,13 +6,14 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import operator
|
||||
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 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:
|
||||
@@ -26,21 +27,24 @@ 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):
|
||||
"""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.
|
||||
Bytes load on ``read``, ``read1``, ``read_range``, or ``read_ranges``, not
|
||||
when the handle is opened. Range reads leave the cursor unchanged. Use
|
||||
:meth:`aread` from async code.
|
||||
"""
|
||||
|
||||
def __init__(self, inner) -> None:
|
||||
self._inner = inner
|
||||
|
||||
async def aread(self) -> bytes:
|
||||
if self.closed:
|
||||
raise ValueError("read of closed file")
|
||||
return await self._inner.read()
|
||||
|
||||
def close(self) -> None:
|
||||
@@ -58,13 +62,16 @@ class BlobFile(io.RawIOBase):
|
||||
|
||||
def seek(self, offset: int, whence: int = io.SEEK_SET) -> int:
|
||||
if whence == io.SEEK_SET:
|
||||
self._inner.seek(offset)
|
||||
position = offset
|
||||
elif whence == io.SEEK_CUR:
|
||||
self._inner.seek(self._inner.tell() + offset)
|
||||
position = self._inner.tell() + offset
|
||||
elif whence == io.SEEK_END:
|
||||
self._inner.seek(self._inner.size() + offset)
|
||||
position = self._inner.size() + offset
|
||||
else:
|
||||
raise ValueError(f"invalid whence: {whence}")
|
||||
if position < 0:
|
||||
raise ValueError(f"negative seek value {position}")
|
||||
self._inner.seek(position)
|
||||
return self._inner.tell()
|
||||
|
||||
def tell(self) -> int:
|
||||
@@ -74,17 +81,39 @@ class BlobFile(io.RawIOBase):
|
||||
return self._inner.size()
|
||||
|
||||
def readall(self) -> bytes:
|
||||
if self.closed:
|
||||
raise ValueError("read of closed file")
|
||||
return self._inner.read_bytes()
|
||||
|
||||
def read(self, size: int = -1) -> bytes:
|
||||
if size == -1:
|
||||
def read(self, size: Optional[int] = -1) -> bytes:
|
||||
if self.closed:
|
||||
raise ValueError("read of closed file")
|
||||
if size is None or size < 0:
|
||||
return self._inner.read_bytes()
|
||||
return super().read(size)
|
||||
|
||||
def read1(self, size: Optional[int] = -1) -> bytes:
|
||||
"""Read up to ``size`` bytes with one underlying read."""
|
||||
if self.closed:
|
||||
raise ValueError("read of closed file")
|
||||
if size is None or size < 0:
|
||||
size = io.DEFAULT_BUFFER_SIZE
|
||||
return self._inner.read_up_to(size)
|
||||
|
||||
def read_range(self, offset: int, length: int) -> bytes:
|
||||
if self.closed:
|
||||
raise ValueError("read of closed file")
|
||||
return self._inner.read_range(offset, length)
|
||||
|
||||
def read_ranges(self, ranges: list[tuple[int, int]]) -> list[bytes]:
|
||||
"""Read ``(offset, length)`` pairs without moving the cursor."""
|
||||
if self.closed:
|
||||
raise ValueError("read of closed file")
|
||||
return self._inner.read_ranges(ranges)
|
||||
|
||||
def readinto(self, b: WriteableBuffer) -> int:
|
||||
if self.closed:
|
||||
raise ValueError("readinto of closed file")
|
||||
view = memoryview(b).cast("B")
|
||||
chunk = self._inner.read_up_to(len(view))
|
||||
view[: len(chunk)] = chunk
|
||||
@@ -107,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(
|
||||
@@ -121,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
|
||||
|
||||
|
||||
@@ -183,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
|
||||
|
||||
|
||||
@@ -196,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")
|
||||
|
||||
@@ -228,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
|
||||
@@ -292,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()
|
||||
|
||||
@@ -328,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:
|
||||
@@ -338,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:
|
||||
@@ -355,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)
|
||||
@@ -422,26 +530,74 @@ 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.Table], blob_column: str
|
||||
row_ids: Union[list[int], pa.Array, pa.ChunkedArray, 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)
|
||||
values = read_row_ids_from_hits(row_ids, blob_column)
|
||||
elif isinstance(row_ids, (pa.Array, pa.ChunkedArray)):
|
||||
if not pa.types.is_integer(row_ids.type):
|
||||
raise ValueError(
|
||||
"row id arrays must have an integer type; pass a query table with "
|
||||
"_rowid when using a blob descriptor column"
|
||||
)
|
||||
values = row_ids.to_pylist()
|
||||
else:
|
||||
values = list(row_ids)
|
||||
|
||||
normalized = []
|
||||
for value in values:
|
||||
try:
|
||||
row_id = operator.index(value)
|
||||
except TypeError as exc:
|
||||
raise ValueError("row ids must be integers without nulls") from exc
|
||||
if row_id < 0:
|
||||
raise ValueError("row ids must be non-negative")
|
||||
normalized.append(row_id)
|
||||
return normalized
|
||||
|
||||
|
||||
def _wrap_blob_files(handles: Iterable[object]) -> list[Optional[BlobFile]]:
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
from datetime import date, datetime, timedelta
|
||||
from decimal import Decimal
|
||||
from typing import Dict, List, Optional, Tuple, Any, TypedDict, Union, Literal
|
||||
from typing import Dict, List, Optional, Sequence, Tuple, Any, TypedDict, Union, Literal
|
||||
from uuid import UUID
|
||||
|
||||
import pyarrow as pa
|
||||
@@ -311,6 +311,7 @@ class BlobFile:
|
||||
def tell(self) -> int: ...
|
||||
def size(self) -> int: ...
|
||||
def read_range(self, offset: int, length: int) -> bytes: ...
|
||||
def read_ranges(self, ranges: list[tuple[int, int]]) -> list[bytes]: ...
|
||||
def read_up_to(self, length: int) -> bytes: ...
|
||||
|
||||
class Job:
|
||||
@@ -563,7 +564,7 @@ class Table:
|
||||
async def fetch_blob_ranges(
|
||||
self,
|
||||
column: str,
|
||||
requests: List[Tuple[int, int, int]],
|
||||
requests: Sequence[Sequence[int]],
|
||||
) -> pa.LargeBinaryArray: ...
|
||||
async def fetch_blob_files(
|
||||
self, column: str, row_ids: list[int]
|
||||
@@ -955,6 +956,8 @@ class PermutationReader:
|
||||
def __init__(self, base_table: Table, permutation_table: Table): ...
|
||||
|
||||
class Catalog:
|
||||
@property
|
||||
def authz(self) -> AuthorizationClient: ...
|
||||
@property
|
||||
def uri(self) -> str: ...
|
||||
async def create_database(
|
||||
@@ -1001,3 +1004,60 @@ class JobListing:
|
||||
def page_token(self) -> Optional[str]: ...
|
||||
def __aiter__(self) -> JobListing: ...
|
||||
async def __anext__(self) -> JobInfo: ...
|
||||
|
||||
class AuthorizationClient:
|
||||
async def list_acls(self, request: str) -> str: ...
|
||||
async def add_acl(self, request: str) -> None: ...
|
||||
async def delete_acl(self, request: str) -> None: ...
|
||||
async def list_role_bindings(self, request: str) -> str: ...
|
||||
async def add_role_binding(self, request: str) -> None: ...
|
||||
async def delete_role_binding(self, request: str) -> None: ...
|
||||
async def get_principal(self, request: str) -> str: ...
|
||||
async def get_group(self, request: str) -> str: ...
|
||||
|
||||
class AuthzSubject:
|
||||
def __init__(self, kind: str, value: str): ...
|
||||
@staticmethod
|
||||
def from_canonical(value: str) -> AuthzSubject: ...
|
||||
@property
|
||||
def kind(self) -> str: ...
|
||||
@property
|
||||
def value(self) -> str: ...
|
||||
@property
|
||||
def display_value(self) -> str: ...
|
||||
def to_wire(self, context: Optional[str] = None) -> str: ...
|
||||
def __str__(self) -> str: ...
|
||||
def __repr__(self) -> str: ...
|
||||
|
||||
class AuthzPrivilege:
|
||||
def __init__(self, value: str): ...
|
||||
@staticmethod
|
||||
def known() -> list[str]: ...
|
||||
def __str__(self) -> str: ...
|
||||
|
||||
class AuthzObject:
|
||||
def __init__(self, value: str): ...
|
||||
@staticmethod
|
||||
def system() -> AuthzObject: ...
|
||||
@staticmethod
|
||||
def database(name: str) -> AuthzObject: ...
|
||||
@staticmethod
|
||||
def namespace(database: str, namespace_path: Sequence[str]) -> AuthzObject: ...
|
||||
@staticmethod
|
||||
def table(
|
||||
database: str, namespace_path: Sequence[str], name: str
|
||||
) -> AuthzObject: ...
|
||||
@staticmethod
|
||||
def view(
|
||||
database: str, namespace_path: Sequence[str], name: str
|
||||
) -> AuthzObject: ...
|
||||
@staticmethod
|
||||
def secret(
|
||||
database: str, namespace_path: Sequence[str], name: str
|
||||
) -> AuthzObject: ...
|
||||
@staticmethod
|
||||
def function(
|
||||
database: str, namespace_path: Sequence[str], name: str
|
||||
) -> AuthzObject: ...
|
||||
def __str__(self) -> str: ...
|
||||
def __repr__(self) -> str: ...
|
||||
|
||||
@@ -0,0 +1,588 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Deployment-wide authorization through ``catalog.authz``.
|
||||
|
||||
ACLs grant one privilege on an object to a principal, group, or role. Role
|
||||
bindings grant a role to a principal or group. Identity-provider group membership
|
||||
is read-only. Name resolution and permission checks are performed by the server.
|
||||
"""
|
||||
|
||||
from collections.abc import AsyncIterator, Iterator, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
import json
|
||||
from typing import Literal, Optional, TypeVar, Union, cast
|
||||
|
||||
from . import _lancedb
|
||||
from .background_loop import LOOP
|
||||
|
||||
__all__ = [
|
||||
"Authorization",
|
||||
"AsyncAuthorization",
|
||||
"Subject",
|
||||
"Object",
|
||||
"Privilege",
|
||||
"AccessControlEntry",
|
||||
"RoleBinding",
|
||||
"IdentityInfo",
|
||||
"PrincipalInfo",
|
||||
"GroupInfo",
|
||||
]
|
||||
|
||||
_PAGE_SIZE = 100
|
||||
_T = TypeVar("_T")
|
||||
_SubjectKind = Literal[
|
||||
"principal_id",
|
||||
"principal_name",
|
||||
"group_id",
|
||||
"group_name",
|
||||
"role",
|
||||
"principal_api_key",
|
||||
"unknown",
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Subject:
|
||||
"""An immutable identity selector. Prefer the named factory methods.
|
||||
|
||||
IDs are stable; names must resolve unambiguously on the server. API-key
|
||||
selectors identify the subject of an operation, not its authenticated caller.
|
||||
Secret values are omitted from representations. Unknown canonical subject
|
||||
kinds returned by newer servers use ``kind="unknown"`` and preserve the full
|
||||
wire string in ``value``; they can be passed back to ACL methods.
|
||||
"""
|
||||
|
||||
kind: _SubjectKind
|
||||
_value: str = field(repr=False)
|
||||
|
||||
_inner: _lancedb.AuthzSubject = field(init=False, repr=False, compare=False)
|
||||
|
||||
def __post_init__(self):
|
||||
object.__setattr__(
|
||||
self, "_inner", _lancedb.AuthzSubject(self.kind, self._value)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def principal_id(cls, id: str) -> "Subject":
|
||||
"""Select a principal by canonical identity-provider ID."""
|
||||
return cls("principal_id", id)
|
||||
|
||||
@classmethod
|
||||
def principal_name(cls, name: str) -> "Subject":
|
||||
"""Select a principal by an unambiguous provider name."""
|
||||
return cls("principal_name", name)
|
||||
|
||||
@classmethod
|
||||
def group_id(cls, id: str) -> "Subject":
|
||||
"""Select a group by canonical identity-provider ID."""
|
||||
return cls("group_id", id)
|
||||
|
||||
@classmethod
|
||||
def group_name(cls, name: str) -> "Subject":
|
||||
"""Select a group by an unambiguous provider name."""
|
||||
return cls("group_name", name)
|
||||
|
||||
@classmethod
|
||||
def role(cls, name: str) -> "Subject":
|
||||
"""Select a LanceDB-managed role, without an ``r:`` prefix."""
|
||||
return cls("role", name)
|
||||
|
||||
@classmethod
|
||||
def principal_api_key(cls, api_key: str) -> "Subject":
|
||||
"""Select the principal holding this secret API key."""
|
||||
return cls("principal_api_key", api_key)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
value = self._inner.display_value
|
||||
return f"Subject.{self.kind}({value!r})"
|
||||
|
||||
@property
|
||||
def value(self) -> str:
|
||||
"""The selected value, or full wire string for unknown subject kinds.
|
||||
|
||||
API-key values are secrets.
|
||||
"""
|
||||
return self._value
|
||||
|
||||
def __reduce__(self):
|
||||
return type(self), (self.kind, self._value)
|
||||
|
||||
def _wire(self) -> str:
|
||||
return self._inner.to_wire()
|
||||
|
||||
@classmethod
|
||||
def _canonical(cls, value: str) -> "Subject":
|
||||
inner = _lancedb.AuthzSubject.from_canonical(value)
|
||||
return cls(cast(_SubjectKind, inner.kind), inner.value)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Object:
|
||||
"""An immutable resource name in the authorization syntax.
|
||||
|
||||
Use the builders to avoid interpreting separators as part of a name.
|
||||
Methods also accept raw object strings, validated by the same Rust parser.
|
||||
Unknown resource types are preserved for compatibility with newer servers.
|
||||
Omitted or empty namespace paths select the database's public namespace.
|
||||
"""
|
||||
|
||||
_value: str
|
||||
_inner: _lancedb.AuthzObject = field(init=False, repr=False, compare=False)
|
||||
|
||||
def __post_init__(self):
|
||||
object.__setattr__(self, "_inner", _lancedb.AuthzObject(self._value))
|
||||
|
||||
@classmethod
|
||||
def system(cls) -> "Object":
|
||||
"""The deployment-wide system object."""
|
||||
return cls(str(_lancedb.AuthzObject.system()))
|
||||
|
||||
@classmethod
|
||||
def database(cls, name: str) -> "Object":
|
||||
"""A database; slash-containing names remain one logical database."""
|
||||
return cls(str(_lancedb.AuthzObject.database(name)))
|
||||
|
||||
@classmethod
|
||||
def namespace(
|
||||
cls, *, database: str, namespace_path: Sequence[str] = ("public",)
|
||||
) -> "Object":
|
||||
"""A namespace path, defaulting to the database's public namespace."""
|
||||
return cls(str(_lancedb.AuthzObject.namespace(database, namespace_path)))
|
||||
|
||||
@classmethod
|
||||
def table(
|
||||
cls, *, database: str, name: str, namespace_path: Sequence[str] = ("public",)
|
||||
) -> "Object":
|
||||
"""A table, defaulting to the database's public namespace."""
|
||||
return cls(str(_lancedb.AuthzObject.table(database, namespace_path, name)))
|
||||
|
||||
@classmethod
|
||||
def view(
|
||||
cls, *, database: str, name: str, namespace_path: Sequence[str] = ("public",)
|
||||
) -> "Object":
|
||||
"""A view, defaulting to the database's public namespace."""
|
||||
return cls(str(_lancedb.AuthzObject.view(database, namespace_path, name)))
|
||||
|
||||
@classmethod
|
||||
def secret(
|
||||
cls, *, database: str, name: str, namespace_path: Sequence[str] = ("public",)
|
||||
) -> "Object":
|
||||
"""A secret, defaulting to the database's public namespace."""
|
||||
return cls(str(_lancedb.AuthzObject.secret(database, namespace_path, name)))
|
||||
|
||||
@classmethod
|
||||
def function(
|
||||
cls, *, database: str, name: str, namespace_path: Sequence[str] = ("public",)
|
||||
) -> "Object":
|
||||
"""A function, defaulting to the database's public namespace."""
|
||||
return cls(str(_lancedb.AuthzObject.function(database, namespace_path, name)))
|
||||
|
||||
def __str__(self) -> str:
|
||||
return str(self._inner)
|
||||
|
||||
def __reduce__(self):
|
||||
return type(self), (self._value,)
|
||||
|
||||
|
||||
class Privilege(str, Enum):
|
||||
"""Known privileges, validated and serialized by Rust.
|
||||
|
||||
ACL methods also accept strings such as ``"SELECT"`` and future privilege
|
||||
names such as ``"PRIVILEGE20"``. Future names are preserved verbatim;
|
||||
their meaning and validity are determined by the server.
|
||||
|
||||
Grant ``CREATE_SECRET`` or ``CREATE_FUNCTION`` on a namespace to delegate
|
||||
creation of those resources without transferring namespace ownership.
|
||||
"""
|
||||
|
||||
def __new__(cls, value: str):
|
||||
canonical = str(_lancedb.AuthzPrivilege(value))
|
||||
member = str.__new__(cls, canonical)
|
||||
member._value_ = canonical
|
||||
return member
|
||||
|
||||
OWNERSHIP = "OWNERSHIP"
|
||||
USAGE = "USAGE"
|
||||
SELECT = "SELECT"
|
||||
INSERT = "INSERT"
|
||||
UPDATE = "UPDATE"
|
||||
DELETE = "DELETE"
|
||||
CREATE_DATABASE = "CREATE_DATABASE"
|
||||
CREATE_NAMESPACE = "CREATE_NAMESPACE"
|
||||
CREATE_TABLE = "CREATE_TABLE"
|
||||
CREATE_MATERIALIZED_VIEW = "CREATE_MATERIALIZED_VIEW"
|
||||
CREATE_VIEW = "CREATE_VIEW"
|
||||
CREATE_SECRET = "CREATE_SECRET"
|
||||
CREATE_FUNCTION = "CREATE_FUNCTION"
|
||||
OPERATE = "OPERATE"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AccessControlEntry:
|
||||
"""One object/subject/privilege grant; subject contains a canonical ID or role."""
|
||||
|
||||
object: str
|
||||
subject: Subject
|
||||
subject_name: Optional[str]
|
||||
privilege: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RoleBinding:
|
||||
"""A direct principal/group role binding, not an expanded effective role."""
|
||||
|
||||
subject: Subject
|
||||
subject_name: Optional[str]
|
||||
role: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class IdentityInfo:
|
||||
"""A canonical identity-provider ID and its optional display name."""
|
||||
|
||||
id: str
|
||||
name: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PrincipalInfo:
|
||||
"""A principal's groups and effective roles, including group-inherited roles."""
|
||||
|
||||
principal_id: str
|
||||
principal_name: Optional[str]
|
||||
groups: list[IdentityInfo]
|
||||
roles: list[str]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GroupInfo:
|
||||
"""A group's provider-managed members and LanceDB-managed direct roles."""
|
||||
|
||||
group_id: str
|
||||
group_name: Optional[str]
|
||||
principals: list[IdentityInfo]
|
||||
roles: list[str]
|
||||
|
||||
|
||||
def _subject(subject: Subject, context: Optional[str] = None) -> str:
|
||||
if not isinstance(subject, Subject):
|
||||
raise ValueError("subject must be a Subject selector")
|
||||
return subject._inner.to_wire(context)
|
||||
|
||||
|
||||
def _object(value: Union[str, Object]) -> str:
|
||||
if not isinstance(value, (str, Object)) or not str(value):
|
||||
raise ValueError("object must be a nonempty string or Object")
|
||||
return str(value) if isinstance(value, Object) else str(_lancedb.AuthzObject(value))
|
||||
|
||||
|
||||
def _limit(limit: Optional[int]) -> None:
|
||||
if limit is not None and (
|
||||
isinstance(limit, bool) or not isinstance(limit, int) or limit < 0
|
||||
):
|
||||
raise ValueError("limit must be a nonnegative integer or None")
|
||||
|
||||
|
||||
def _privilege(value: Union[str, Privilege]) -> str:
|
||||
if not isinstance(value, str):
|
||||
raise ValueError("privilege must be a string or Privilege")
|
||||
return str(_lancedb.AuthzPrivilege(value))
|
||||
|
||||
|
||||
def _entry(value: dict) -> AccessControlEntry:
|
||||
return AccessControlEntry(
|
||||
value["object"],
|
||||
Subject._canonical(value["subject_id"]),
|
||||
value.get("subject_name"),
|
||||
value["privilege"],
|
||||
)
|
||||
|
||||
|
||||
def _binding(value: dict) -> RoleBinding:
|
||||
return RoleBinding(
|
||||
Subject._canonical(value["subject_id"]),
|
||||
value.get("subject_name"),
|
||||
value["role"],
|
||||
)
|
||||
|
||||
|
||||
class AsyncAuthorization:
|
||||
"""Async authorization returned by ``AsyncCatalog.authz``.
|
||||
|
||||
HTTP failures raise [HttpError][lancedb.remote.errors.HttpError], preserving
|
||||
status and request ID. Invalid local arguments raise ValueError. Lists are
|
||||
lazy async iterators: failures can occur during iteration after some results
|
||||
have been yielded. Iteration does not provide a consistent snapshot.
|
||||
"""
|
||||
|
||||
def __init__(self, inner: _lancedb.AuthorizationClient):
|
||||
self._inner = inner
|
||||
|
||||
async def _pages(self, method, request: dict, key: str, limit: Optional[int]):
|
||||
yielded = 0
|
||||
while limit is None or yielded < limit:
|
||||
request["max_results"] = (
|
||||
_PAGE_SIZE if limit is None else min(_PAGE_SIZE, limit - yielded)
|
||||
)
|
||||
response = json.loads(await method(json.dumps(request)))
|
||||
for value in response[key]:
|
||||
yield value
|
||||
yielded += 1
|
||||
if limit is not None and yielded >= limit:
|
||||
return
|
||||
token = response.get("continuation_token")
|
||||
if not token:
|
||||
return
|
||||
if token == request.get("continuation_token"):
|
||||
raise RuntimeError(
|
||||
"Authorization listing returned a repeated continuation token"
|
||||
)
|
||||
request["continuation_token"] = token
|
||||
|
||||
def list_acls(
|
||||
self,
|
||||
*,
|
||||
object: Optional[Union[str, Object]] = None,
|
||||
subject: Optional[Subject] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> AsyncIterator[AccessControlEntry]:
|
||||
"""Iterate grants, optionally filtered by exact object and/or subject.
|
||||
|
||||
Requires USAGE on the supplied object, or on system if object is omitted.
|
||||
``limit`` bounds yielded entries across all pages; None lists all and zero
|
||||
makes no requests. Missing matches produce an empty iterator.
|
||||
"""
|
||||
_limit(limit)
|
||||
request = {}
|
||||
if object is not None:
|
||||
request["object"] = _object(object)
|
||||
if subject is not None:
|
||||
request["subject"] = _subject(subject)
|
||||
|
||||
async def entries():
|
||||
async for value in self._pages(
|
||||
self._inner.list_acls, request, "entries", limit
|
||||
):
|
||||
yield _entry(value)
|
||||
|
||||
return entries()
|
||||
|
||||
async def add_acl(
|
||||
self,
|
||||
*,
|
||||
object: Union[str, Object],
|
||||
subject: Subject,
|
||||
privilege: Union[str, Privilege],
|
||||
) -> None:
|
||||
"""Grant one privilege; requires OWNERSHIP on the object.
|
||||
|
||||
A duplicate grant raises HTTP 409. Granting OWNERSHIP transfers ownership.
|
||||
Writes are not automatically retried after ambiguous failures.
|
||||
"""
|
||||
await self._inner.add_acl(
|
||||
json.dumps(
|
||||
{
|
||||
"object": _object(object),
|
||||
"subject": _subject(subject),
|
||||
"privilege": _privilege(privilege),
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
async def delete_acl(
|
||||
self,
|
||||
*,
|
||||
object: Union[str, Object],
|
||||
subject: Subject,
|
||||
privilege: Union[str, Privilege],
|
||||
) -> None:
|
||||
"""Revoke one privilege; requires OWNERSHIP on the object.
|
||||
|
||||
A missing grant raises HTTP 404. OWNERSHIP itself cannot be deleted.
|
||||
"""
|
||||
await self._inner.delete_acl(
|
||||
json.dumps(
|
||||
{
|
||||
"object": _object(object),
|
||||
"subject": _subject(subject),
|
||||
"privilege": _privilege(privilege),
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
def list_role_bindings(
|
||||
self,
|
||||
*,
|
||||
subject: Optional[Subject] = None,
|
||||
role: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> AsyncIterator[RoleBinding]:
|
||||
"""Iterate direct bindings, optionally filtered by principal/group and role.
|
||||
|
||||
Requires USAGE on system. Limit semantics match ``list_acls``. Principals'
|
||||
roles inherited through groups are available through ``get_principal``.
|
||||
"""
|
||||
_limit(limit)
|
||||
request = {}
|
||||
if subject is not None:
|
||||
request["subject"] = _subject(subject, "role_member")
|
||||
if role is not None:
|
||||
request["role"] = role
|
||||
|
||||
async def bindings():
|
||||
async for value in self._pages(
|
||||
self._inner.list_role_bindings, request, "bindings", limit
|
||||
):
|
||||
yield _binding(value)
|
||||
|
||||
return bindings()
|
||||
|
||||
async def add_role_binding(self, *, subject: Subject, role: str) -> None:
|
||||
"""Grant a role idempotently; requires OPERATE on system."""
|
||||
await self._inner.add_role_binding(
|
||||
json.dumps({"subject": _subject(subject, "role_member"), "role": role})
|
||||
)
|
||||
|
||||
async def delete_role_binding(self, *, subject: Subject, role: str) -> None:
|
||||
"""Idempotently revoke a role; requires OPERATE on system."""
|
||||
await self._inner.delete_role_binding(
|
||||
json.dumps({"subject": _subject(subject, "role_member"), "role": role})
|
||||
)
|
||||
|
||||
async def get_principal(self, subject: Subject) -> PrincipalInfo:
|
||||
"""Look up groups and effective roles; requires USAGE on system.
|
||||
|
||||
Select by ID, name, or API key. Missing/ambiguous identities raise HTTP
|
||||
404/409; unavailable or unconfigured resolution raises HTTP 503/501.
|
||||
"""
|
||||
value = json.loads(
|
||||
await self._inner.get_principal(
|
||||
json.dumps({"principal": _subject(subject, "principal")})
|
||||
)
|
||||
)
|
||||
return PrincipalInfo(
|
||||
value["principal_id"],
|
||||
value.get("principal_name"),
|
||||
[IdentityInfo(**group) for group in value["groups"]],
|
||||
value["roles"],
|
||||
)
|
||||
|
||||
async def get_group(self, subject: Subject) -> GroupInfo:
|
||||
"""Look up a group's members and roles by ID/name; requires USAGE on system.
|
||||
|
||||
Membership is read-only and comes from the identity provider. Lookup
|
||||
errors have the same HTTP semantics as ``get_principal``.
|
||||
"""
|
||||
value = json.loads(
|
||||
await self._inner.get_group(
|
||||
json.dumps({"group": _subject(subject, "group")})
|
||||
)
|
||||
)
|
||||
return GroupInfo(
|
||||
value["group_id"],
|
||||
value.get("group_name"),
|
||||
[IdentityInfo(**principal) for principal in value["principals"]],
|
||||
value["roles"],
|
||||
)
|
||||
|
||||
|
||||
async def _next(iterator: AsyncIterator[_T]) -> _T:
|
||||
return await iterator.__anext__()
|
||||
|
||||
|
||||
def _sync_iter(iterator: AsyncIterator[_T]) -> Iterator[_T]:
|
||||
try:
|
||||
while True:
|
||||
yield LOOP.run(_next(iterator))
|
||||
except StopAsyncIteration:
|
||||
return
|
||||
finally:
|
||||
close = getattr(iterator, "aclose", None)
|
||||
if close is not None:
|
||||
LOOP.run(close())
|
||||
|
||||
|
||||
class Authorization:
|
||||
"""Synchronous authorization returned by ``Catalog.authz``.
|
||||
|
||||
Shares the async API's permissions, errors, and lazy pagination behavior.
|
||||
Group membership comes from the identity provider and cannot be edited here.
|
||||
"""
|
||||
|
||||
def __init__(self, inner: AsyncAuthorization):
|
||||
self._inner = inner
|
||||
|
||||
def list_acls(
|
||||
self,
|
||||
*,
|
||||
object: Optional[Union[str, Object]] = None,
|
||||
subject: Optional[Subject] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> Iterator[AccessControlEntry]:
|
||||
"""Iterate grants.
|
||||
|
||||
See [list_acls][lancedb.authz.AsyncAuthorization.list_acls] for permissions.
|
||||
"""
|
||||
return _sync_iter(
|
||||
self._inner.list_acls(object=object, subject=subject, limit=limit)
|
||||
)
|
||||
|
||||
def add_acl(
|
||||
self,
|
||||
*,
|
||||
object: Union[str, Object],
|
||||
subject: Subject,
|
||||
privilege: Union[str, Privilege],
|
||||
) -> None:
|
||||
"""Grant a privilege; duplicates raise HTTP 409.
|
||||
|
||||
OWNERSHIP transfers ownership.
|
||||
"""
|
||||
LOOP.run(
|
||||
self._inner.add_acl(object=object, subject=subject, privilege=privilege)
|
||||
)
|
||||
|
||||
def delete_acl(
|
||||
self,
|
||||
*,
|
||||
object: Union[str, Object],
|
||||
subject: Subject,
|
||||
privilege: Union[str, Privilege],
|
||||
) -> None:
|
||||
"""Revoke a privilege. Missing grants raise HTTP 404.
|
||||
|
||||
OWNERSHIP cannot be deleted.
|
||||
"""
|
||||
LOOP.run(
|
||||
self._inner.delete_acl(object=object, subject=subject, privilege=privilege)
|
||||
)
|
||||
|
||||
def list_role_bindings(
|
||||
self,
|
||||
*,
|
||||
subject: Optional[Subject] = None,
|
||||
role: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> Iterator[RoleBinding]:
|
||||
"""Iterate direct role bindings, using the same limit semantics as list_acls."""
|
||||
return _sync_iter(
|
||||
self._inner.list_role_bindings(subject=subject, role=role, limit=limit)
|
||||
)
|
||||
|
||||
def add_role_binding(self, *, subject: Subject, role: str) -> None:
|
||||
"""Grant a role idempotently; requires OPERATE on system."""
|
||||
LOOP.run(self._inner.add_role_binding(subject=subject, role=role))
|
||||
|
||||
def delete_role_binding(self, *, subject: Subject, role: str) -> None:
|
||||
"""Idempotently revoke a role; requires OPERATE on system."""
|
||||
LOOP.run(self._inner.delete_role_binding(subject=subject, role=role))
|
||||
|
||||
def get_principal(self, subject: Subject) -> PrincipalInfo:
|
||||
"""Look up groups and effective roles; requires USAGE on system."""
|
||||
return LOOP.run(self._inner.get_principal(subject))
|
||||
|
||||
def get_group(self, subject: Subject) -> GroupInfo:
|
||||
"""Look up a group's members and roles; requires USAGE on system."""
|
||||
return LOOP.run(self._inner.get_group(subject))
|
||||
@@ -3,11 +3,13 @@
|
||||
|
||||
"""Remote catalogs manage databases through a server's root namespace."""
|
||||
|
||||
from functools import cached_property
|
||||
from datetime import timedelta
|
||||
from typing import Any, AsyncIterator, Iterator, Optional, Union
|
||||
|
||||
from . import _lancedb
|
||||
from .background_loop import LOOP
|
||||
from .authz import AsyncAuthorization, Authorization
|
||||
from .db import AsyncConnection, DBConnection
|
||||
from .remote import ClientConfig, OAuthConfig
|
||||
from .remote.db import RemoteDBConnection
|
||||
@@ -86,6 +88,11 @@ class AsyncCatalog:
|
||||
def __init__(self, inner: _lancedb.Catalog):
|
||||
self._inner = inner
|
||||
|
||||
@cached_property
|
||||
def authz(self) -> AsyncAuthorization:
|
||||
"""Authorization sharing this catalog's scope and credentials."""
|
||||
return AsyncAuthorization(self._inner.authz)
|
||||
|
||||
@property
|
||||
def uri(self) -> str:
|
||||
"""The catalog's root namespace endpoint."""
|
||||
@@ -157,6 +164,11 @@ class Catalog:
|
||||
self._sql_host_override = sql_host_override
|
||||
self._oauth_config = oauth_config
|
||||
|
||||
@cached_property
|
||||
def authz(self) -> Authorization:
|
||||
"""Authorization sharing this catalog's scope and credentials."""
|
||||
return Authorization(self._inner.authz)
|
||||
|
||||
@property
|
||||
def uri(self) -> str:
|
||||
"""The catalog's root namespace endpoint."""
|
||||
|
||||
+406
-115
@@ -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,50 @@ 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 _is_remote_query_table(table: Any) -> bool:
|
||||
async_table = getattr(table, "_table", table)
|
||||
return not async_table._inner._is_native()
|
||||
|
||||
|
||||
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 +189,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 +214,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
|
||||
|
||||
@@ -312,6 +352,28 @@ def _finish_plain_scan_pandas(
|
||||
return df
|
||||
|
||||
|
||||
def _finish_remote_pandas(
|
||||
tbl: pa.Table,
|
||||
*,
|
||||
blob_mode: BlobMode,
|
||||
blob_paths: Iterable[str],
|
||||
flatten: Optional[Union[int, bool]],
|
||||
**kwargs,
|
||||
) -> pd.DataFrame:
|
||||
# A live remote table can be replaced between query and blob fetch.
|
||||
# Row ids and version numbers do not identify the producing table.
|
||||
if not blob_column_paths(tbl.schema):
|
||||
return flatten_columns(tbl, flatten).to_pandas(**kwargs)
|
||||
if blob_mode == "descriptions":
|
||||
tbl = strip_auto_row_ids(tbl, blob_paths)
|
||||
return flatten_columns(tbl, flatten).to_pandas(**kwargs)
|
||||
raise NotImplementedError(
|
||||
f"remote to_pandas(blob_mode={blob_mode!r}) cannot safely materialize "
|
||||
"blob columns without a stable table snapshot; "
|
||||
"use blob_mode='descriptions'"
|
||||
)
|
||||
|
||||
|
||||
async def _finish_plain_scan_pandas_async(
|
||||
scanner: _LanceScanner,
|
||||
*,
|
||||
@@ -1072,44 +1134,58 @@ 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. Eligible local plain scans
|
||||
use Lance's scanner, which returns lazy blob handles in ``"lazy"``
|
||||
mode and also supports legacy blob columns. Other local blob v2
|
||||
queries return descriptors in ``"lazy"`` and ``"descriptions"``
|
||||
modes, or payloads in ``"bytes"`` mode, through :meth:`to_arrow`.
|
||||
Remote blob queries support only ``"descriptions"`` until a stable
|
||||
table snapshot can bind the query and blob fetches.
|
||||
**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 = (
|
||||
"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):
|
||||
raise _unsupported_blob_pandas_error(
|
||||
"this query shape cannot use Lance native pandas conversion"
|
||||
if _is_remote_query_table(self._table):
|
||||
blob_sources = self._validated_blob_sources(blob_mode)
|
||||
return _finish_remote_pandas(
|
||||
self.to_arrow(timeout=timeout),
|
||||
blob_mode=blob_mode,
|
||||
blob_paths=blob_sources,
|
||||
flatten=flatten,
|
||||
**kwargs,
|
||||
)
|
||||
return tbl.to_pandas(**kwargs)
|
||||
|
||||
output_schema = getattr(self, "output_schema", None)
|
||||
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 _query_needs_lance_scanner(self):
|
||||
raise _unsupported_blob_pandas_error(
|
||||
"this query shape cannot use Lance native pandas conversion"
|
||||
)
|
||||
|
||||
arrow_mode = "descriptions" if blob_mode == "lazy" else blob_mode
|
||||
tbl = self._to_arrow_in_mode(timeout, arrow_mode)
|
||||
if (
|
||||
blob_mode == "lazy"
|
||||
and schema_has_blob_field(tbl.schema)
|
||||
and not blob_v2_projection_sources(self._table.schema, self._columns)
|
||||
):
|
||||
raise _unsupported_blob_pandas_error(
|
||||
"the query has no fetchable blob v2 source"
|
||||
)
|
||||
return flatten_columns(tbl, flatten).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 +1199,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 +1234,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 +1252,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 +1437,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 +1705,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 +1717,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 +1734,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 +1995,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 +2014,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 +2297,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 +2340,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 +2483,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 +2533,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 +3010,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 +3274,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 +3329,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 +3362,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,44 +3407,49 @@ 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. Eligible local plain scans
|
||||
use Lance's scanner, which returns lazy blob handles in ``"lazy"``
|
||||
mode and also supports legacy blob columns. Other local blob v2
|
||||
queries return descriptors in ``"lazy"`` and ``"descriptions"``
|
||||
modes, or payloads in ``"bytes"`` mode, through :meth:`to_arrow`.
|
||||
Remote blob queries support only ``"descriptions"`` until a stable
|
||||
table snapshot can bind the query and blob fetches.
|
||||
**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 _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):
|
||||
arrow_mode = "descriptions" if blob_mode == "lazy" else blob_mode
|
||||
tbl = await self.to_arrow(timeout=timeout, blob_mode=arrow_mode)
|
||||
if (
|
||||
blob_mode == "lazy"
|
||||
and schema_has_blob_field(tbl.schema)
|
||||
and not await self._validated_blob_sources("descriptions")
|
||||
):
|
||||
raise _unsupported_blob_pandas_error(
|
||||
"this query shape cannot use Lance native pandas conversion"
|
||||
"the query has no fetchable blob v2 source"
|
||||
)
|
||||
return tbl.to_pandas(**kwargs)
|
||||
return flatten_columns(tbl, flatten).to_pandas(**kwargs)
|
||||
|
||||
async def _remote_to_pandas(
|
||||
self,
|
||||
@@ -3188,21 +3459,13 @@ class AsyncQueryBase(object):
|
||||
blob_mode: BlobMode,
|
||||
**kwargs,
|
||||
) -> "pd.DataFrame":
|
||||
# A live remote table can be replaced between query and blob fetch.
|
||||
# Row ids and version numbers do not identify the producing table.
|
||||
tbl = await self.to_arrow(timeout=timeout)
|
||||
projected_blob_paths = set(blob_column_paths(tbl.schema))
|
||||
if not projected_blob_paths:
|
||||
return flatten_columns(tbl, flatten).to_pandas(**kwargs)
|
||||
|
||||
if blob_mode == "descriptions":
|
||||
tbl = strip_auto_row_ids(tbl, self._blob_paths)
|
||||
return flatten_columns(tbl, flatten).to_pandas(**kwargs)
|
||||
|
||||
raise NotImplementedError(
|
||||
f"remote to_pandas(blob_mode={blob_mode!r}) cannot safely materialize "
|
||||
"blob columns without a stable table snapshot; "
|
||||
"use blob_mode='descriptions'"
|
||||
return _finish_remote_pandas(
|
||||
tbl,
|
||||
blob_mode=blob_mode,
|
||||
blob_paths=self._blob_paths,
|
||||
flatten=flatten,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
async def _plain_scan_to_pandas(
|
||||
@@ -3229,6 +3492,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 +3513,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 +4353,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 +4718,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 +4737,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 +4764,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 +4809,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.
|
||||
|
||||
@@ -170,6 +170,10 @@ class ClientConfig:
|
||||
Configuration for retrying failed requests.
|
||||
timeout_config: Optional[TimeoutConfig]
|
||||
Configuration for request timeouts.
|
||||
blob_request_concurrency: Optional[int]
|
||||
Maximum concurrent blob HTTP requests across all handles on this
|
||||
connection. Defaults to 8. Set a positive integer here or via
|
||||
`LANCE_CLIENT_BLOB_REQUEST_CONCURRENCY`.
|
||||
extra_headers: Optional[dict]
|
||||
Additional headers to include in requests.
|
||||
id_delimiter: Optional[str]
|
||||
@@ -193,6 +197,7 @@ class ClientConfig:
|
||||
user_agent: str = f"LanceDB-Python-Client/{__version__}"
|
||||
retry_config: RetryConfig = field(default_factory=RetryConfig)
|
||||
timeout_config: Optional[TimeoutConfig] = field(default_factory=TimeoutConfig)
|
||||
blob_request_concurrency: Optional[int] = None
|
||||
extra_headers: Optional[dict] = None
|
||||
id_delimiter: Optional[str] = None
|
||||
tls_config: Optional[TlsConfig] = None
|
||||
|
||||
@@ -112,6 +112,7 @@ def _client_config_to_dict(config: ClientConfig) -> dict[str, Any]:
|
||||
"user_agent": config.user_agent,
|
||||
"retry_config": _retry_config_to_dict(config.retry_config),
|
||||
"timeout_config": _timeout_config_to_dict(config.timeout_config),
|
||||
"blob_request_concurrency": config.blob_request_concurrency,
|
||||
"extra_headers": config.extra_headers,
|
||||
"id_delimiter": config.id_delimiter,
|
||||
"tls_config": _tls_config_to_dict(config.tls_config),
|
||||
|
||||
@@ -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]):
|
||||
@@ -1154,7 +1157,9 @@ class RemoteTable(Table):
|
||||
return LOOP.run(self._table.blob_columns())
|
||||
|
||||
def fetch_blobs(
|
||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
||||
self,
|
||||
column: str,
|
||||
row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table],
|
||||
) -> pa.LargeBinaryArray:
|
||||
return LOOP.run(self._table.fetch_blobs(column, row_ids))
|
||||
|
||||
@@ -1162,7 +1167,9 @@ class RemoteTable(Table):
|
||||
return LOOP.run(self._table.fetch_blob_ranges(column, list(requests)))
|
||||
|
||||
def fetch_blob_files(
|
||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
||||
self,
|
||||
column: str,
|
||||
row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table],
|
||||
) -> "list[Optional[BlobFile]]":
|
||||
return LOOP.run(self._table.fetch_blob_files(column, row_ids))
|
||||
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
"""Schema helpers for Lance blob columns."""
|
||||
|
||||
import importlib
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import pyarrow as pa
|
||||
import pyarrow.ipc
|
||||
@@ -181,13 +181,60 @@ def _resolve_blob_type():
|
||||
return _resolved_blob_type
|
||||
|
||||
|
||||
def blob(name: str, nullable: bool = True) -> pa.Field:
|
||||
def blob(
|
||||
name: str,
|
||||
nullable: bool = True,
|
||||
*,
|
||||
inline_size_threshold: Optional[int] = None,
|
||||
dedicated_size_threshold: Optional[int] = None,
|
||||
pack_file_size_threshold: Optional[int] = None,
|
||||
) -> pa.Field:
|
||||
"""Create a Lance blob v2 column field.
|
||||
|
||||
When pylance is installed this is ``lance.blob.BlobType``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name:
|
||||
Field name.
|
||||
nullable:
|
||||
Whether the column accepts null values.
|
||||
inline_size_threshold:
|
||||
Maximum payload size in bytes to keep inline in the data file
|
||||
before using packed blob storage. Zero is allowed.
|
||||
dedicated_size_threshold:
|
||||
Maximum payload size in bytes to store in packed blob storage
|
||||
before using dedicated blob storage. Must be positive.
|
||||
pack_file_size_threshold:
|
||||
Maximum size in bytes of a single packed blob sidecar file before
|
||||
a new one is started. Must be positive.
|
||||
|
||||
The thresholds are forwarded to ``lance.blob_field``, which validates
|
||||
them and stores them as field metadata, so pylance is required whenever
|
||||
any threshold is set. When all thresholds are omitted the behavior is
|
||||
unchanged.
|
||||
"""
|
||||
blob_type = _resolve_blob_type()
|
||||
return pa.field(name, blob_type(), nullable=nullable)
|
||||
if (
|
||||
inline_size_threshold is None
|
||||
and dedicated_size_threshold is None
|
||||
and pack_file_size_threshold is None
|
||||
):
|
||||
blob_type = _resolve_blob_type()
|
||||
return pa.field(name, blob_type(), nullable=nullable)
|
||||
try:
|
||||
lance = importlib.import_module("lance")
|
||||
except ModuleNotFoundError as err:
|
||||
raise ValueError(
|
||||
"Setting blob tier thresholds requires pylance to be installed"
|
||||
) from err
|
||||
_resolve_blob_type()
|
||||
return lance.blob_field(
|
||||
name,
|
||||
nullable=nullable,
|
||||
inline_size_threshold=inline_size_threshold,
|
||||
dedicated_size_threshold=dedicated_size_threshold,
|
||||
pack_file_size_threshold=pack_file_size_threshold,
|
||||
)
|
||||
|
||||
|
||||
def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataType:
|
||||
|
||||
+184
-37
@@ -576,6 +576,55 @@ def _cast_to_target_schema(
|
||||
return pa.RecordBatchReader.from_batches(reordered_schema, gen())
|
||||
|
||||
|
||||
def _preserve_blob_v2_fields(
|
||||
inferred_schema: pa.Schema, existing_schema: pa.Schema
|
||||
) -> pa.Schema:
|
||||
"""Keep existing blob fields when an overwrite infers a replacement schema."""
|
||||
existing_fields = {field.name: field for field in existing_schema}
|
||||
fields = [
|
||||
_preserve_blob_v2_field(field, existing_fields[field.name])
|
||||
if field.name in existing_fields
|
||||
else field
|
||||
for field in inferred_schema
|
||||
]
|
||||
return pa.schema(fields, metadata=inferred_schema.metadata)
|
||||
|
||||
|
||||
def _preserve_blob_v2_field(inferred: pa.Field, existing: pa.Field) -> pa.Field:
|
||||
if is_blob_v2_field(existing):
|
||||
return existing
|
||||
|
||||
inferred_type, existing_type = inferred.type, existing.type
|
||||
if pa.types.is_struct(inferred_type) and pa.types.is_struct(existing_type):
|
||||
existing_children = {field.name: field for field in existing_type}
|
||||
children = [
|
||||
_preserve_blob_v2_field(field, existing_children[field.name])
|
||||
if field.name in existing_children
|
||||
else field
|
||||
for field in inferred_type
|
||||
]
|
||||
return inferred.with_type(pa.struct(children))
|
||||
|
||||
if _is_list_like(inferred_type) and _is_list_like(existing_type):
|
||||
value_field = _preserve_blob_v2_field(
|
||||
inferred_type.value_field, existing_type.value_field
|
||||
)
|
||||
if value_field == inferred_type.value_field:
|
||||
return inferred
|
||||
value_field = value_field.with_name(existing_type.value_field.name)
|
||||
if pa.types.is_fixed_size_list(inferred_type):
|
||||
value_type = pa.list_(value_field, inferred_type.list_size)
|
||||
elif pa.types.is_large_list(inferred_type):
|
||||
value_type = pa.large_list(value_field)
|
||||
else:
|
||||
value_type = pa.list_(value_field)
|
||||
return inferred.with_type(value_type)
|
||||
|
||||
if blob_v2_column_paths(pa.schema([existing])):
|
||||
return existing
|
||||
return inferred
|
||||
|
||||
|
||||
def _coerce_blob_write_columns(
|
||||
batch: pa.RecordBatch, target_schema: pa.Schema
|
||||
) -> pa.RecordBatch:
|
||||
@@ -691,12 +740,7 @@ def _coerce_blob_list_values(
|
||||
|
||||
|
||||
def _coerce_value_to_blob(values: pa.Array, target_field: pa.Field) -> pa.Array:
|
||||
if pa.types.is_null(values.type):
|
||||
data = pa.nulls(len(values), type=pa.large_binary())
|
||||
elif pa.types.is_large_binary(values.type):
|
||||
data = values
|
||||
else:
|
||||
data = values.cast(pa.large_binary())
|
||||
value_field = "uri" if _is_string_like(values.type) else "data"
|
||||
length = len(values)
|
||||
storage_type = target_field.type
|
||||
if isinstance(storage_type, pa.ExtensionType):
|
||||
@@ -704,8 +748,11 @@ def _coerce_value_to_blob(values: pa.Array, target_field: pa.Field) -> pa.Array:
|
||||
storage_fields = list(storage_type)
|
||||
children = []
|
||||
for storage_field in storage_fields:
|
||||
if storage_field.name == "data":
|
||||
children.append(data)
|
||||
if storage_field.name == value_field and not pa.types.is_null(values.type):
|
||||
if values.type == storage_field.type:
|
||||
children.append(values)
|
||||
else:
|
||||
children.append(values.cast(storage_field.type))
|
||||
else:
|
||||
children.append(pa.nulls(length, type=storage_field.type))
|
||||
storage = pa.StructArray.from_arrays(
|
||||
@@ -725,7 +772,19 @@ def _physical_array_and_type(array: pa.Array) -> tuple[pa.Array, pa.DataType]:
|
||||
|
||||
|
||||
def _can_coerce_to_blob(data_type: pa.DataType) -> bool:
|
||||
return _is_binary_like(data_type) or pa.types.is_null(data_type)
|
||||
return (
|
||||
_is_binary_like(data_type)
|
||||
or _is_string_like(data_type)
|
||||
or pa.types.is_null(data_type)
|
||||
)
|
||||
|
||||
|
||||
def _is_string_like(data_type: pa.DataType) -> bool:
|
||||
return (
|
||||
pa.types.is_string(data_type)
|
||||
or pa.types.is_large_string(data_type)
|
||||
or pa.types.is_string_view(data_type)
|
||||
)
|
||||
|
||||
|
||||
def _is_binary_like(data_type: pa.DataType) -> bool:
|
||||
@@ -887,6 +946,24 @@ def _arrow_json_storage_type(input_type: pa.DataType) -> Optional[pa.DataType]:
|
||||
|
||||
|
||||
def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field:
|
||||
if is_blob_v2_field(target_field):
|
||||
source_type = field.type
|
||||
if pa.types.is_struct(source_type):
|
||||
if not any(child.name in ("data", "uri") for child in source_type):
|
||||
raise ValueError(
|
||||
f"blob struct input for column '{field.name}' must contain "
|
||||
"a 'data' or 'uri' child"
|
||||
)
|
||||
elif not _can_coerce_to_blob(source_type) and not isinstance(
|
||||
source_type, pa.ExtensionType
|
||||
):
|
||||
raise ValueError(
|
||||
f"cannot coerce column '{field.name}' with type {source_type} "
|
||||
"into a blob v2 struct. expected binary bytes "
|
||||
"(Binary, LargeBinary, BinaryView), strings "
|
||||
"(Utf8, LargeUtf8, Utf8View), or a Struct with a "
|
||||
"'data' or 'uri' child"
|
||||
)
|
||||
# LanceDB exposes stored JSON columns as lance.json (JSONB-backed LargeBinary), but
|
||||
# casting the input to that storage type here merely relabels the raw JSON bytes as
|
||||
# JSONB. Lance must see arrow.json so it can perform the JSONB encoding.
|
||||
@@ -1322,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
|
||||
@@ -1738,7 +1823,8 @@ class Table(ABC):
|
||||
- pyarrow.Table or pyarrow.RecordBatch
|
||||
mode: str
|
||||
The mode to use when writing the data. Valid values are
|
||||
"append" and "overwrite".
|
||||
"append" and "overwrite". Overwrite infers a replacement schema
|
||||
but preserves existing blob v2 fields present in the new data.
|
||||
on_bad_vectors: str, default "error"
|
||||
What to do if any of the vectors are not the same size or contains NaNs.
|
||||
One of "error", "drop", "fill".
|
||||
@@ -2054,19 +2140,30 @@ class Table(ABC):
|
||||
|
||||
@abstractmethod
|
||||
def blob_columns(self) -> list[str]:
|
||||
"""Names of the blob v2 columns declared on this table."""
|
||||
"""Names of blob v2 columns that can be fetched by row id.
|
||||
|
||||
Struct-nested blobs use dotted paths. Blobs inside lists are omitted
|
||||
because a row id does not identify an individual list element.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def fetch_blobs(
|
||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
||||
self,
|
||||
column: str,
|
||||
row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table],
|
||||
) -> pa.LargeBinaryArray:
|
||||
"""Materialize full blob bytes for ``column`` at the given rows.
|
||||
|
||||
The result has the same length and order as ``row_ids``. Null blobs
|
||||
produce null slots; valid empty blobs produce ``b""``.
|
||||
|
||||
``_rowid`` values stay valid after compaction when the table has stable
|
||||
row ids.
|
||||
``row_ids`` may be a list, an integer Arrow array or chunked array,
|
||||
or a query table with row ids. ``_rowid`` values stay valid after
|
||||
compaction when the table has stable row ids.
|
||||
One missing or deleted row ID rejects the entire batch; no partial
|
||||
results are returned.
|
||||
|
||||
Blobs inside lists cannot be fetched by row id yet.
|
||||
|
||||
Remote servers limit each request to 1024 row IDs and 64 MiB of blob
|
||||
bytes. The client splits requests automatically and reads an individual
|
||||
@@ -2078,17 +2175,20 @@ class Table(ABC):
|
||||
def fetch_blob_ranges(
|
||||
self,
|
||||
column: str,
|
||||
requests: Sequence[Tuple[int, int, int]],
|
||||
requests: Sequence[Sequence[int]],
|
||||
) -> pa.LargeBinaryArray:
|
||||
"""Materialize row-specific byte ranges from a blob v2 column.
|
||||
|
||||
Each request is a ``(row_id, offset, length)`` tuple. Requests may be
|
||||
repeated or reordered, including multiple ranges for the same blob.
|
||||
Each request is a three-item ``(row_id, offset, length)`` sequence.
|
||||
Requests may be repeated or reordered, including multiple ranges for
|
||||
the same blob.
|
||||
The result has the same length and order as ``requests``; null blobs
|
||||
produce null slots and empty ranges on non-null blobs produce ``b""``.
|
||||
|
||||
``_rowid`` values stay valid after compaction when the table has stable
|
||||
row ids.
|
||||
One missing or deleted row ID rejects the entire batch; no partial
|
||||
results are returned.
|
||||
|
||||
Row IDs can be obtained from a query with ``with_row_id(True)``. This
|
||||
API is currently supported only by local tables.
|
||||
@@ -2096,18 +2196,22 @@ class Table(ABC):
|
||||
|
||||
@abstractmethod
|
||||
def fetch_blob_files(
|
||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
||||
self,
|
||||
column: str,
|
||||
row_ids: Union[list[int], pa.Array, pa.ChunkedArray, 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 a query ``pyarrow.Table`` carrying row identity via
|
||||
``_rowid`` or a ``_lance_row_id`` field on the blob descriptor. Null
|
||||
rows are ``None``. Remote tables require LanceDB Cloud server 0.5.0 or
|
||||
newer.
|
||||
a ``list[int]``, an integer Arrow array or chunked array, or a query
|
||||
``pyarrow.Table`` carrying row identity via ``_rowid`` or a
|
||||
``_lance_row_id`` field on the blob descriptor. Null blob values return
|
||||
``None``. Remote tables require LanceDB Cloud server 0.5.0 or newer.
|
||||
|
||||
``_rowid`` values stay valid after compaction when the table has stable
|
||||
row ids.
|
||||
One missing or deleted row ID rejects the entire batch; no partial
|
||||
handles are returned.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
@@ -2865,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:
|
||||
@@ -3013,19 +3119,23 @@ class LanceTable(Table):
|
||||
return LOOP.run(self._table.blob_columns())
|
||||
|
||||
def fetch_blobs(
|
||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
||||
self,
|
||||
column: str,
|
||||
row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table],
|
||||
) -> pa.LargeBinaryArray:
|
||||
return LOOP.run(self._table.fetch_blobs(column, row_ids))
|
||||
|
||||
def fetch_blob_ranges(
|
||||
self,
|
||||
column: str,
|
||||
requests: Sequence[Tuple[int, int, int]],
|
||||
requests: Sequence[Sequence[int]],
|
||||
) -> pa.LargeBinaryArray:
|
||||
return LOOP.run(self._table.fetch_blob_ranges(column, list(requests)))
|
||||
|
||||
def fetch_blob_files(
|
||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
||||
self,
|
||||
column: str,
|
||||
row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table],
|
||||
) -> "list[Optional[BlobFile]]":
|
||||
return LOOP.run(self._table.fetch_blob_files(column, row_ids))
|
||||
|
||||
@@ -3239,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
|
||||
@@ -3980,7 +4101,8 @@ class LanceTable(Table):
|
||||
The data to insert into the table.
|
||||
mode: str
|
||||
The mode to use when writing the data. Valid values are
|
||||
"append" and "overwrite".
|
||||
"append" and "overwrite". Overwrite infers a replacement schema
|
||||
but preserves existing blob v2 fields present in the new data.
|
||||
on_bad_vectors: str, default "error"
|
||||
What to do if any of the vectors are not the same size or contains NaNs.
|
||||
One of "error", "drop", "fill", "null".
|
||||
@@ -5687,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,
|
||||
@@ -5998,7 +6135,8 @@ class AsyncTable:
|
||||
- pyarrow.Table or pyarrow.RecordBatch
|
||||
mode: str
|
||||
The mode to use when writing the data. Valid values are
|
||||
"append" and "overwrite".
|
||||
"append" and "overwrite". Overwrite infers a replacement schema
|
||||
but preserves existing blob v2 fields present in the new data.
|
||||
on_bad_vectors: str, default "error"
|
||||
What to do if any of the vectors are not the same size or contains NaNs.
|
||||
One of "error", "drop", "fill", "null".
|
||||
@@ -6036,6 +6174,11 @@ class AsyncTable:
|
||||
data, _ = sanitize_create_table(
|
||||
data, None, on_bad_vectors=on_bad_vectors, fill_value=fill_value
|
||||
)
|
||||
if blob_v2_column_paths(schema):
|
||||
# Plain bytes cannot be inferred as blob v2. Keep matching blob
|
||||
# fields while allowing other columns to use the inferred schema.
|
||||
target_schema = _preserve_blob_v2_fields(data.schema, schema)
|
||||
data = _cast_to_target_schema(data, target_schema)
|
||||
elif on_bad_vectors != "error" or (
|
||||
schema.metadata is not None and b"embedding_functions" in schema.metadata
|
||||
):
|
||||
@@ -7079,7 +7222,9 @@ class AsyncTable:
|
||||
return await self._inner.blob_columns()
|
||||
|
||||
async def fetch_blobs(
|
||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
||||
self,
|
||||
column: str,
|
||||
row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table],
|
||||
) -> pa.LargeBinaryArray:
|
||||
return await self._inner.fetch_blobs(
|
||||
column, _normalize_blob_row_ids(row_ids, column)
|
||||
@@ -7088,12 +7233,14 @@ class AsyncTable:
|
||||
async def fetch_blob_ranges(
|
||||
self,
|
||||
column: str,
|
||||
requests: Sequence[Tuple[int, int, int]],
|
||||
requests: Sequence[Sequence[int]],
|
||||
) -> pa.LargeBinaryArray:
|
||||
return await self._inner.fetch_blob_ranges(column, list(requests))
|
||||
|
||||
async def fetch_blob_files(
|
||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
||||
self,
|
||||
column: str,
|
||||
row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table],
|
||||
) -> "list[Optional[BlobFile]]":
|
||||
handles = await self._inner.fetch_blob_files(
|
||||
column, _normalize_blob_row_ids(row_ids, column)
|
||||
|
||||
@@ -236,6 +236,19 @@ def inf_vector_column_query(schema: pa.Schema, dim: Optional[int] = None) -> str
|
||||
vector_col_names = []
|
||||
for field in schema:
|
||||
vector_col_names.extend(_iter_vector_columns(field, [], dim))
|
||||
|
||||
if len(vector_col_names) == 0 and dim is not None:
|
||||
# The query's dimension matched no column. If the schema is
|
||||
# unambiguous without that filter, the single vector column is still
|
||||
# the right one to search: a dimension mismatch is a query-execution
|
||||
# error, not a missing-column error.
|
||||
unfiltered = []
|
||||
for field in schema:
|
||||
unfiltered.extend(_iter_vector_columns(field, []))
|
||||
if len(unfiltered) == 1:
|
||||
return unfiltered[0]
|
||||
vector_col_names = unfiltered
|
||||
|
||||
if len(vector_col_names) > 1:
|
||||
raise ValueError(
|
||||
"Schema has more than one vector column. "
|
||||
|
||||
@@ -0,0 +1,821 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import asyncio
|
||||
from dataclasses import FrozenInstanceError
|
||||
import json
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
import lancedb
|
||||
from lancedb.authz import (
|
||||
AccessControlEntry,
|
||||
AsyncAuthorization,
|
||||
Authorization,
|
||||
GroupInfo,
|
||||
IdentityInfo,
|
||||
Object,
|
||||
PrincipalInfo,
|
||||
Privilege,
|
||||
RoleBinding,
|
||||
Subject,
|
||||
)
|
||||
from lancedb.remote import HeaderProvider
|
||||
from lancedb.remote.errors import HttpError
|
||||
from test_catalog import catalog_server as _catalog_server
|
||||
|
||||
catalog_server = _catalog_server
|
||||
|
||||
|
||||
class DynamicHeaders(HeaderProvider):
|
||||
def get_headers(self):
|
||||
return {
|
||||
"Authorization": "Bearer rotating-token",
|
||||
"X-LanceDB-Database": "wrong-dynamic-db",
|
||||
"X-LanceDB-Database-Prefix": "wrong-prefix",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
def test_authz_operations_and_scope(catalog_server, asynchronous):
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.extend(
|
||||
[
|
||||
(200, {}),
|
||||
(200, {}),
|
||||
(200, {}),
|
||||
(200, {}),
|
||||
(
|
||||
200,
|
||||
{
|
||||
"entries": [
|
||||
{
|
||||
"subject_id": "g:team-id",
|
||||
"object": "table:tenant/db:a$b$events",
|
||||
"privilege": "PRIVILEGE20",
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(
|
||||
200,
|
||||
{
|
||||
"bindings": [
|
||||
{
|
||||
"subject_id": "p:alice-id",
|
||||
"subject_name": "Alice",
|
||||
"role": "reader",
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(
|
||||
200,
|
||||
{
|
||||
"principal_id": "alice-id",
|
||||
"groups": [{"id": "team-id"}],
|
||||
"roles": ["reader"],
|
||||
},
|
||||
),
|
||||
(
|
||||
200,
|
||||
{
|
||||
"group_id": "team-id",
|
||||
"group_name": "Engineering",
|
||||
"principals": [{"id": "alice-id", "name": "Alice"}],
|
||||
"roles": ["reader"],
|
||||
},
|
||||
),
|
||||
]
|
||||
)
|
||||
config = {
|
||||
"header_provider": DynamicHeaders(),
|
||||
"extra_headers": {"x-lancedb-database": "wrong-static-db"},
|
||||
}
|
||||
obj = Object.table(database="tenant/db", namespace_path=["a", "b"], name="events")
|
||||
|
||||
async def exercise():
|
||||
if asynchronous:
|
||||
catalog = await lancedb.connect_catalog_async(
|
||||
endpoint, client_config=config
|
||||
)
|
||||
else:
|
||||
catalog = lancedb.connect_catalog(endpoint, client_config=config)
|
||||
authz = catalog.authz
|
||||
assert authz is catalog.authz
|
||||
assert requests == [] # Constructing handles does not probe a database.
|
||||
|
||||
async def call(method, *args, **kwargs):
|
||||
value = method(*args, **kwargs)
|
||||
return await value if asynchronous else value
|
||||
|
||||
await call(
|
||||
authz.add_acl,
|
||||
object=obj,
|
||||
subject=Subject.group_name("Engineering"),
|
||||
privilege=Privilege.SELECT,
|
||||
)
|
||||
await call(
|
||||
authz.delete_acl,
|
||||
object=obj,
|
||||
subject=Subject.group_id("team-id"),
|
||||
privilege="SELECT",
|
||||
)
|
||||
await call(
|
||||
authz.add_role_binding,
|
||||
subject=Subject.principal_name("Alice"),
|
||||
role="reader",
|
||||
)
|
||||
await call(
|
||||
authz.delete_role_binding,
|
||||
subject=Subject.principal_api_key("subject-secret"),
|
||||
role="reader",
|
||||
)
|
||||
iterator = authz.list_acls(object=obj, subject=Subject.role("reader"))
|
||||
entries = (
|
||||
[entry async for entry in iterator] if asynchronous else list(iterator)
|
||||
)
|
||||
assert entries == [
|
||||
AccessControlEntry(
|
||||
str(obj), Subject.group_id("team-id"), None, "PRIVILEGE20"
|
||||
)
|
||||
]
|
||||
iterator = authz.list_role_bindings(
|
||||
subject=Subject.group_id("team-id"), role="reader"
|
||||
)
|
||||
bindings = (
|
||||
[binding async for binding in iterator] if asynchronous else list(iterator)
|
||||
)
|
||||
assert bindings == [
|
||||
RoleBinding(Subject.principal_id("alice-id"), "Alice", "reader")
|
||||
]
|
||||
assert await call(
|
||||
authz.get_principal, Subject.principal_id("alice-id")
|
||||
) == PrincipalInfo("alice-id", None, [IdentityInfo("team-id")], ["reader"])
|
||||
assert await call(
|
||||
authz.get_group, Subject.group_name("Engineering")
|
||||
) == GroupInfo(
|
||||
"team-id", "Engineering", [IdentityInfo("alice-id", "Alice")], ["reader"]
|
||||
)
|
||||
|
||||
asyncio.run(exercise())
|
||||
assert [path for path, _, _ in requests] == [
|
||||
"/admin/authz/acl/add",
|
||||
"/admin/authz/acl/delete",
|
||||
"/admin/authz/role_binding/add",
|
||||
"/admin/authz/role_binding/delete",
|
||||
"/admin/authz/acl/list",
|
||||
"/admin/authz/role_binding/list",
|
||||
"/admin/authz/principal",
|
||||
"/admin/authz/group",
|
||||
]
|
||||
assert [body for _, _, body in requests] == [
|
||||
{"object": str(obj), "subject": "G:Engineering", "privilege": "SELECT"},
|
||||
{"object": str(obj), "subject": "g:team-id", "privilege": "SELECT"},
|
||||
{"subject": "P:Alice", "role": "reader"},
|
||||
{"subject": "a:subject-secret", "role": "reader"},
|
||||
{"object": str(obj), "subject": "r:reader", "max_results": 100},
|
||||
{"subject": "g:team-id", "role": "reader", "max_results": 100},
|
||||
{"principal": "p:alice-id"},
|
||||
{"group": "G:Engineering"},
|
||||
]
|
||||
for _, headers, _ in requests:
|
||||
headers = {key.lower(): value for key, value in headers.items()}
|
||||
assert "x-lancedb-database" not in headers
|
||||
assert "x-lancedb-database-prefix" not in headers
|
||||
assert headers["authorization"] == "Bearer rotating-token"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
@pytest.mark.parametrize("limit", [None, 0, 1, 3])
|
||||
def test_lazy_expanded_acl_pagination(catalog_server, asynchronous, limit):
|
||||
endpoint, requests, responses = catalog_server
|
||||
|
||||
def entry(privilege):
|
||||
return {
|
||||
"object": "database:db",
|
||||
"subject_id": "r:reader",
|
||||
"privilege": privilege,
|
||||
}
|
||||
|
||||
responses.extend(
|
||||
[
|
||||
(
|
||||
200,
|
||||
{
|
||||
"entries": [entry("SELECT"), entry("INSERT")],
|
||||
"continuation_token": "opaque/first",
|
||||
},
|
||||
),
|
||||
(200, {"entries": [], "continuation_token": "opaque/second"}),
|
||||
(200, {"entries": [entry("UPDATE"), entry("DELETE")]}),
|
||||
]
|
||||
)
|
||||
|
||||
async def exercise():
|
||||
catalog = (
|
||||
await lancedb.connect_catalog_async(endpoint, api_key="caller-secret")
|
||||
if asynchronous
|
||||
else lancedb.connect_catalog(endpoint, api_key="caller-secret")
|
||||
)
|
||||
iterator = catalog.authz.list_acls(subject=Subject.role("reader"), limit=limit)
|
||||
assert requests == []
|
||||
return [entry async for entry in iterator] if asynchronous else list(iterator)
|
||||
|
||||
values = asyncio.run(exercise())
|
||||
expected = ["SELECT", "INSERT", "UPDATE", "DELETE"]
|
||||
assert [entry.privilege for entry in values] == (
|
||||
expected if limit is None else expected[:limit]
|
||||
)
|
||||
assert len(requests) == (0 if limit == 0 else 1 if limit == 1 else 3)
|
||||
for _, headers, body in requests:
|
||||
assert headers["x-api-key"] == "caller-secret"
|
||||
assert body["subject"] == "r:reader"
|
||||
if len(requests) == 3:
|
||||
assert requests[1][2]["continuation_token"] == "opaque/first"
|
||||
assert requests[2][2]["continuation_token"] == "opaque/second"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", [400, 401, 403, 404, 409, 501, 503])
|
||||
def test_http_errors_preserve_status_without_replaying_writes(catalog_server, status):
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.append((status, {"error": "authorization failure"}))
|
||||
authz = lancedb.connect_catalog(endpoint).authz
|
||||
with pytest.raises(HttpError) as error:
|
||||
authz.add_acl(
|
||||
object="system", subject=Subject.role("operator"), privilege="OPERATE"
|
||||
)
|
||||
assert error.value.status_code == status
|
||||
assert error.value.request_id
|
||||
assert len(requests) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
def test_role_pagination_and_late_failure(catalog_server, asynchronous):
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.extend(
|
||||
[
|
||||
(
|
||||
200,
|
||||
{
|
||||
"bindings": [{"subject_id": "g:team", "role": "reader"}],
|
||||
"continuation_token": "next",
|
||||
},
|
||||
),
|
||||
(403, {"error": "denied"}),
|
||||
]
|
||||
)
|
||||
|
||||
async def exercise():
|
||||
catalog = (
|
||||
await lancedb.connect_catalog_async(endpoint)
|
||||
if asynchronous
|
||||
else lancedb.connect_catalog(endpoint)
|
||||
)
|
||||
iterator = catalog.authz.list_role_bindings(role="reader")
|
||||
first = await anext(iterator) if asynchronous else next(iterator)
|
||||
assert first.subject == Subject.group_id("team")
|
||||
assert len(requests) == 1
|
||||
with pytest.raises(HttpError) as error:
|
||||
if asynchronous:
|
||||
await anext(iterator)
|
||||
else:
|
||||
next(iterator)
|
||||
assert error.value.status_code == 403
|
||||
|
||||
asyncio.run(exercise())
|
||||
assert requests[1][2] == {
|
||||
"role": "reader",
|
||||
"max_results": 100,
|
||||
"continuation_token": "next",
|
||||
}
|
||||
|
||||
|
||||
def test_early_close_does_not_fetch_next_page(catalog_server):
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.append(
|
||||
(
|
||||
200,
|
||||
{
|
||||
"bindings": [{"subject_id": "g:team", "role": "reader"}],
|
||||
"continuation_token": "next",
|
||||
},
|
||||
)
|
||||
)
|
||||
iterator = lancedb.connect_catalog(endpoint).authz.list_role_bindings()
|
||||
next(iterator)
|
||||
iterator.close()
|
||||
assert len(requests) == 1
|
||||
|
||||
|
||||
def test_selectors_objects_and_redaction():
|
||||
secret = Subject.principal_api_key("do-not-print-this")
|
||||
assert "do-not-print-this" not in repr(secret)
|
||||
assert "do-not-print-this" not in str(secret)
|
||||
with pytest.raises(FrozenInstanceError):
|
||||
secret.kind = "role"
|
||||
assert str(Object.system()) == "system"
|
||||
assert str(Object.database("tenant/db")) == "database:tenant/db"
|
||||
assert (
|
||||
str(Object.namespace(database="db", namespace_path=["a", "b"]))
|
||||
== "namespace:db:a$b"
|
||||
)
|
||||
assert str(Object.table(database="db", name="events")) == "table:db:public$events"
|
||||
assert str(Object.view(database="db", name="recent")) == "view:db:public$recent"
|
||||
assert Subject.principal_id("p:id")._wire() == "p:p:id"
|
||||
assert Subject.group_id("team-id").value == "team-id"
|
||||
for namespace in (["a$b"], [""], "public", [".."]):
|
||||
with pytest.raises(ValueError):
|
||||
Object.table(database="db", namespace_path=namespace, name="events")
|
||||
for database in ("", "a//b", "a/../b", "a:b", "a%2Fb", "a$b"):
|
||||
with pytest.raises(ValueError):
|
||||
Object.database(database)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validation_before_io_and_empty_lists():
|
||||
inner = AsyncMock()
|
||||
inner.list_acls.return_value = json.dumps({"entries": []})
|
||||
inner.list_role_bindings.return_value = json.dumps({"bindings": []})
|
||||
authz = AsyncAuthorization(inner)
|
||||
for limit in (-1, 1.5, True):
|
||||
with pytest.raises(ValueError):
|
||||
authz.list_acls(limit=limit)
|
||||
with pytest.raises(ValueError):
|
||||
Authorization(authz).list_role_bindings(limit=limit)
|
||||
with pytest.raises(ValueError):
|
||||
await authz.get_group(Subject.principal_id("alice"))
|
||||
with pytest.raises(ValueError):
|
||||
await authz.get_principal(Subject.group_id("team"))
|
||||
for method in (authz.add_role_binding, authz.delete_role_binding):
|
||||
with pytest.raises(ValueError):
|
||||
await method(subject=Subject.role("reader"), role="writer")
|
||||
with pytest.raises(ValueError):
|
||||
authz.list_role_bindings(subject=Subject.role("reader"))
|
||||
assert [entry async for entry in authz.list_acls(limit=0)] == []
|
||||
assert inner.mock_calls == []
|
||||
assert [entry async for entry in authz.list_acls()] == []
|
||||
assert [entry async for entry in authz.list_role_bindings()] == []
|
||||
|
||||
|
||||
def test_api_key_subject_is_not_logged(catalog_server):
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.append((200, {"principal_id": "public-id", "groups": [], "roles": []}))
|
||||
result = subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-c",
|
||||
"import lancedb, sys; from lancedb.authz import Subject; "
|
||||
"lancedb.connect_catalog(sys.argv[1]).authz.get_principal("
|
||||
"Subject.principal_api_key('test-subject-secret-never-log'))",
|
||||
endpoint,
|
||||
],
|
||||
env={**os.environ, "LANCEDB_LOG": "debug"},
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
timeout=30,
|
||||
)
|
||||
assert requests[0][2] == {"principal": "a:test-subject-secret-never-log"}
|
||||
assert "test-subject-secret-never-log" not in result.stdout + result.stderr
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"builder, kwargs, wire",
|
||||
[
|
||||
("system", {}, "system"),
|
||||
("secret", {"database": "db", "name": "api-key"}, "secret:db:public$api-key"),
|
||||
(
|
||||
"function",
|
||||
{"database": "db", "name": "caption"},
|
||||
"function:db:public$caption",
|
||||
),
|
||||
(
|
||||
"secret",
|
||||
{"database": "tenant/db", "namespace_path": ["a", "b"], "name": "api-key"},
|
||||
"secret:tenant/db:a$b$api-key",
|
||||
),
|
||||
(
|
||||
"function",
|
||||
{"database": "tenant/db", "namespace_path": ["a", "b"], "name": "caption"},
|
||||
"function:tenant/db:a$b$caption",
|
||||
),
|
||||
("database", {"name": "tenant/db"}, "database:tenant/db"),
|
||||
(
|
||||
"namespace",
|
||||
{"database": "tenant/db", "namespace_path": ["a", "b"]},
|
||||
"namespace:tenant/db:a$b",
|
||||
),
|
||||
(
|
||||
"table",
|
||||
{"database": "tenant/db", "namespace_path": ["a", "b"], "name": "events"},
|
||||
"table:tenant/db:a$b$events",
|
||||
),
|
||||
(
|
||||
"view",
|
||||
{"database": "db", "namespace_path": ["public"], "name": "recent"},
|
||||
"view:db:public$recent",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_rust_backed_object_builders(builder, kwargs, wire):
|
||||
import pickle
|
||||
from lancedb import _lancedb
|
||||
|
||||
value = getattr(Object, builder)(**kwargs)
|
||||
assert str(value) == wire
|
||||
assert str(value._inner) == wire
|
||||
assert str(_lancedb.AuthzObject(wire)) == wire
|
||||
assert value == Object(wire)
|
||||
assert hash(value) == hash(Object(wire))
|
||||
assert pickle.loads(pickle.dumps(value)) == value
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kind, prefix",
|
||||
[
|
||||
("principal_id", "p:"),
|
||||
("principal_name", "P:"),
|
||||
("group_id", "g:"),
|
||||
("group_name", "G:"),
|
||||
("role", "r:"),
|
||||
("principal_api_key", "a:"),
|
||||
],
|
||||
)
|
||||
def test_rust_backed_subject_selectors(kind, prefix):
|
||||
import pickle
|
||||
from lancedb import _lancedb
|
||||
|
||||
value = "p:literal@example.com / 名前"
|
||||
subject = getattr(Subject, kind)(value)
|
||||
assert subject.kind == kind
|
||||
assert subject.value == value
|
||||
assert subject._wire() == prefix + value
|
||||
assert _lancedb.AuthzSubject(kind, value).to_wire() == prefix + value
|
||||
assert pickle.loads(pickle.dumps(subject)) == subject
|
||||
assert hash(subject) == hash(getattr(Subject, kind)(value))
|
||||
with pytest.raises(ValueError):
|
||||
getattr(Subject, kind)("")
|
||||
if kind == "principal_api_key":
|
||||
assert value not in repr(subject._inner)
|
||||
assert value not in str(subject._inner)
|
||||
if kind in ("principal_id", "group_id", "role"):
|
||||
assert Subject._canonical(prefix + value) == subject
|
||||
else:
|
||||
with pytest.raises(ValueError) as error:
|
||||
Subject._canonical(prefix + value)
|
||||
assert value not in str(error.value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"wire",
|
||||
[
|
||||
"",
|
||||
"database:",
|
||||
"database:a//b",
|
||||
"database:a/../b",
|
||||
"database:a%2Fb",
|
||||
"namespace:db:a$$b",
|
||||
"table:db:events",
|
||||
"table:db:public$",
|
||||
"table:db:public$..",
|
||||
"table:db:a/b$events",
|
||||
"view:db:public$café",
|
||||
"secret:db:public$",
|
||||
"function:db:public$..",
|
||||
"secret:db:a$$key",
|
||||
"function:db:caption",
|
||||
"secret:db:public$bad/name",
|
||||
"function:db:public$bad:name",
|
||||
"system:db",
|
||||
],
|
||||
)
|
||||
def test_invalid_raw_objects_fail_before_io(catalog_server, wire):
|
||||
endpoint, requests, _ = catalog_server
|
||||
authz = lancedb.connect_catalog(endpoint).authz
|
||||
with pytest.raises(ValueError):
|
||||
Object(wire)
|
||||
with pytest.raises(ValueError):
|
||||
authz.list_acls(object=wire)
|
||||
with pytest.raises(ValueError):
|
||||
authz.add_acl(object=wire, subject=Subject.role("reader"), privilege="SELECT")
|
||||
assert requests == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"wire",
|
||||
[
|
||||
"future:db:resource",
|
||||
"table_extension:db:resource",
|
||||
"secret:tenant/db:a$b$key",
|
||||
"function:tenant/db:a$b$caption",
|
||||
],
|
||||
)
|
||||
def test_objects_round_trip_through_acl_requests(catalog_server, wire):
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.extend(
|
||||
[
|
||||
(200, {}),
|
||||
(
|
||||
200,
|
||||
{
|
||||
"entries": [
|
||||
{
|
||||
"object": wire,
|
||||
"subject_id": "r:reader",
|
||||
"privilege": "SELECT",
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(200, {}),
|
||||
]
|
||||
)
|
||||
authz = lancedb.connect_catalog(endpoint).authz
|
||||
obj = Object(wire)
|
||||
assert str(obj) == wire
|
||||
authz.add_acl(object=obj, subject=Subject.role("reader"), privilege="SELECT")
|
||||
entries = list(authz.list_acls(object=wire))
|
||||
assert entries[0].object == wire
|
||||
authz.delete_acl(
|
||||
object=entries[0].object, subject=Subject.role("reader"), privilege="SELECT"
|
||||
)
|
||||
assert len(requests) == 3
|
||||
for _, _, body in requests:
|
||||
assert body["object"] == wire
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_requests_validate_selector_context_before_io(catalog_server):
|
||||
endpoint, requests, _ = catalog_server
|
||||
catalog = await lancedb.connect_catalog_async(endpoint)
|
||||
native = catalog.authz._inner
|
||||
for method, request in [
|
||||
(native.list_role_bindings, {"subject": "r:reader"}),
|
||||
(native.add_role_binding, {"subject": "r:reader", "role": "writer"}),
|
||||
(native.delete_role_binding, {"subject": "r:reader", "role": "writer"}),
|
||||
(native.get_principal, {"principal": "g:group"}),
|
||||
(native.get_group, {"group": "a:never-log-this-key"}),
|
||||
]:
|
||||
with pytest.raises(ValueError) as error:
|
||||
await method(json.dumps(request))
|
||||
assert "never-log-this-key" not in str(error.value)
|
||||
assert requests == []
|
||||
|
||||
|
||||
def test_privilege_values_use_rust_serialization():
|
||||
import pickle
|
||||
from lancedb import _lancedb
|
||||
|
||||
names = [value.value for value in Privilege]
|
||||
assert names == _lancedb.AuthzPrivilege.known()
|
||||
for privilege in Privilege:
|
||||
assert json.loads(json.dumps(privilege)) == privilege.value
|
||||
assert str(_lancedb.AuthzPrivilege(privilege)) == privilege.value
|
||||
assert pickle.loads(pickle.dumps(privilege)) is privilege
|
||||
for index in range(64):
|
||||
alias = f"PRIVILEGE{index}"
|
||||
expected = names[index] if index < len(names) else alias
|
||||
assert str(_lancedb.AuthzPrivilege(alias)) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"creation_privilege", [None, "CREATE_SECRET", "CREATE_FUNCTION"]
|
||||
)
|
||||
def test_privilege_names_and_future_values_round_trip(
|
||||
catalog_server, asynchronous, creation_privilege
|
||||
):
|
||||
endpoint, requests, responses = catalog_server
|
||||
inputs = [
|
||||
Privilege.SELECT,
|
||||
"SELECT",
|
||||
"PRIVILEGE2",
|
||||
"PRIVILEGE20",
|
||||
"PRIVILEGE63",
|
||||
"PRIVILEGE64",
|
||||
"PRIVILEGE256",
|
||||
"PRIVILEGE18446744073709551616",
|
||||
]
|
||||
expected = ["SELECT", "SELECT", "SELECT", *inputs[3:]]
|
||||
object_id = "table:db:public$t"
|
||||
if creation_privilege is not None:
|
||||
inputs = [Privilege[creation_privilege], creation_privilege]
|
||||
expected = [creation_privilege, creation_privilege]
|
||||
object_id = "namespace:db:public"
|
||||
for privilege in inputs:
|
||||
responses.extend(
|
||||
[
|
||||
(200, {}),
|
||||
(
|
||||
200,
|
||||
{
|
||||
"entries": [
|
||||
{
|
||||
"object": object_id,
|
||||
"subject_id": "r:reader",
|
||||
"privilege": privilege,
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(200, {}),
|
||||
]
|
||||
)
|
||||
|
||||
async def exercise():
|
||||
catalog = (
|
||||
await lancedb.connect_catalog_async(endpoint)
|
||||
if asynchronous
|
||||
else lancedb.connect_catalog(endpoint)
|
||||
)
|
||||
authz = catalog.authz
|
||||
for privilege, canonical in zip(inputs, expected):
|
||||
added = authz.add_acl(
|
||||
object=object_id,
|
||||
subject=Subject.role("reader"),
|
||||
privilege=privilege,
|
||||
)
|
||||
if asynchronous:
|
||||
await added
|
||||
iterator = authz.list_acls()
|
||||
entries = (
|
||||
[entry async for entry in iterator] if asynchronous else list(iterator)
|
||||
)
|
||||
assert entries[0].privilege == canonical
|
||||
deleted = authz.delete_acl(
|
||||
object=entries[0].object,
|
||||
subject=entries[0].subject,
|
||||
privilege=entries[0].privilege,
|
||||
)
|
||||
if asynchronous:
|
||||
await deleted
|
||||
|
||||
asyncio.run(exercise())
|
||||
assert len(requests) == 3 * len(inputs)
|
||||
for index, canonical in enumerate(expected):
|
||||
assert requests[index * 3][2]["object"] == object_id
|
||||
assert requests[index * 3][2]["privilege"] == canonical
|
||||
assert requests[index * 3 + 2][2]["object"] == object_id
|
||||
assert requests[index * 3 + 2][2]["privilege"] == canonical
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
def test_invalid_privileges_fail_before_io(catalog_server, asynchronous):
|
||||
endpoint, requests, _ = catalog_server
|
||||
invalid = [
|
||||
"",
|
||||
"select",
|
||||
"SELECT ",
|
||||
"UNKNOWN",
|
||||
"PRIVILEGE",
|
||||
"PRIVILEGE-1",
|
||||
"PRIVILEGE+1",
|
||||
"PRIVILEGE1x",
|
||||
None,
|
||||
2,
|
||||
True,
|
||||
["SELECT"],
|
||||
]
|
||||
|
||||
async def exercise():
|
||||
catalog = (
|
||||
await lancedb.connect_catalog_async(endpoint)
|
||||
if asynchronous
|
||||
else lancedb.connect_catalog(endpoint)
|
||||
)
|
||||
async_authz = catalog.authz if asynchronous else catalog.authz._inner
|
||||
native = async_authz._inner
|
||||
for privilege in invalid:
|
||||
for method in [catalog.authz.add_acl, catalog.authz.delete_acl]:
|
||||
with pytest.raises(ValueError, match="privilege"):
|
||||
result = method(
|
||||
object="system",
|
||||
subject=Subject.role("reader"),
|
||||
privilege=privilege,
|
||||
)
|
||||
if asynchronous:
|
||||
await result
|
||||
for method in [
|
||||
native.add_acl,
|
||||
native.delete_acl,
|
||||
]:
|
||||
with pytest.raises(ValueError):
|
||||
await method(
|
||||
json.dumps(
|
||||
{
|
||||
"object": "system",
|
||||
"subject": "r:reader",
|
||||
"privilege": privilege,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
asyncio.run(exercise())
|
||||
assert requests == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
def test_acl_page_preserves_unknown_subjects(catalog_server, asynchronous):
|
||||
import pickle
|
||||
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.extend(
|
||||
[
|
||||
(
|
||||
200,
|
||||
{
|
||||
"entries": [
|
||||
{
|
||||
"subject_id": "p:known",
|
||||
"object": "system",
|
||||
"privilege": "USAGE",
|
||||
},
|
||||
{
|
||||
"subject_id": "x:future-subject",
|
||||
"object": "system",
|
||||
"privilege": "USAGE",
|
||||
},
|
||||
]
|
||||
},
|
||||
),
|
||||
(200, {}),
|
||||
]
|
||||
)
|
||||
|
||||
async def exercise():
|
||||
catalog = (
|
||||
await lancedb.connect_catalog_async(endpoint)
|
||||
if asynchronous
|
||||
else lancedb.connect_catalog(endpoint)
|
||||
)
|
||||
iterator = catalog.authz.list_acls()
|
||||
entries = (
|
||||
[entry async for entry in iterator] if asynchronous else list(iterator)
|
||||
)
|
||||
assert len(entries) == 2
|
||||
assert entries[0].subject == Subject.principal_id("known")
|
||||
subject = entries[1].subject
|
||||
assert subject.kind == "unknown"
|
||||
assert subject.value == "x:future-subject"
|
||||
assert subject._wire() == "x:future-subject"
|
||||
assert pickle.loads(pickle.dumps(subject)) == subject
|
||||
result = catalog.authz.delete_acl(
|
||||
object=entries[1].object, subject=subject, privilege=entries[1].privilege
|
||||
)
|
||||
if asynchronous:
|
||||
await result
|
||||
|
||||
asyncio.run(exercise())
|
||||
assert len(requests) == 2
|
||||
assert requests[1][2]["subject"] == "x:future-subject"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("builder", ["secret", "function"])
|
||||
def test_secret_and_function_builders_validate_components(builder):
|
||||
from lancedb import _lancedb
|
||||
|
||||
for database, namespace, name in [
|
||||
("a//b", ["public"], "valid"),
|
||||
("db", ["a$b"], "valid"),
|
||||
("db", "public", "valid"),
|
||||
("db", ["public"], "bad$name"),
|
||||
("db", ["public"], ".."),
|
||||
]:
|
||||
with pytest.raises(ValueError):
|
||||
getattr(Object, builder)(
|
||||
database=database, namespace_path=namespace, name=name
|
||||
)
|
||||
with pytest.raises(ValueError):
|
||||
getattr(_lancedb.AuthzObject, builder)(database, namespace, name)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"builder", ["namespace", "table", "view", "secret", "function"]
|
||||
)
|
||||
def test_empty_and_omitted_namespace_paths_select_public(builder):
|
||||
from lancedb import _lancedb
|
||||
|
||||
kwargs = {"database": "db"}
|
||||
native_args = ()
|
||||
wire = "namespace:db:public"
|
||||
if builder != "namespace":
|
||||
kwargs["name"] = "item"
|
||||
native_args = ("item",)
|
||||
wire = f"{builder}:db:public$item"
|
||||
|
||||
build = getattr(Object, builder)
|
||||
expected = build(**kwargs, namespace_path=["public"])
|
||||
assert str(expected) == wire
|
||||
assert build(**kwargs) == expected
|
||||
for namespace in [[], ()]:
|
||||
assert build(**kwargs, namespace_path=namespace) == expected
|
||||
assert (
|
||||
str(getattr(_lancedb.AuthzObject, builder)("db", namespace, *native_args))
|
||||
== wire
|
||||
)
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -70,21 +71,25 @@ def _row_ids_by_id(table):
|
||||
|
||||
def _assert_missing_blob_row_ids(exc_info):
|
||||
message = str(exc_info.value)
|
||||
assert "row ids" in message
|
||||
assert "row ids but some do not exist in the table" in message
|
||||
assert "rowaddr" not in message
|
||||
assert "fragment" not in message
|
||||
|
||||
|
||||
def _assert_fetch_apis_reject_missing_row_ids(table, row_ids):
|
||||
expected = f"first missing row ids: [{row_ids[0]}]"
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
table.fetch_blobs("image", row_ids)
|
||||
_assert_missing_blob_row_ids(exc_info)
|
||||
assert expected in str(exc_info.value)
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
table.fetch_blob_files("image", row_ids)
|
||||
_assert_missing_blob_row_ids(exc_info)
|
||||
assert expected in str(exc_info.value)
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
table.fetch_blob_ranges("image", [(row_id, 0, 1) for row_id in row_ids])
|
||||
_assert_missing_blob_row_ids(exc_info)
|
||||
assert expected in str(exc_info.value)
|
||||
|
||||
|
||||
def test_blob_factory_declares_v2_field():
|
||||
@@ -95,6 +100,37 @@ def test_blob_factory_declares_v2_field():
|
||||
assert type(field.type) is LanceBlobType
|
||||
|
||||
|
||||
def test_blob_thresholds_land_in_field_metadata():
|
||||
field = lancedb.blob(
|
||||
"payload",
|
||||
inline_size_threshold=8 * 1024,
|
||||
dedicated_size_threshold=1 << 20,
|
||||
pack_file_size_threshold=1 << 26,
|
||||
)
|
||||
assert field.type.extension_name == "lance.blob.v2"
|
||||
metadata = field.metadata
|
||||
assert metadata[b"lance-encoding:blob-inline-size-threshold"] == b"8192"
|
||||
assert metadata[b"lance-encoding:blob-dedicated-size-threshold"] == b"1048576"
|
||||
assert metadata[b"lance-encoding:blob-pack-file-size-threshold"] == b"67108864"
|
||||
|
||||
|
||||
def test_blob_without_thresholds_has_no_threshold_metadata():
|
||||
field = lancedb.blob("image")
|
||||
metadata = field.metadata or {}
|
||||
assert not any(b"threshold" in key for key in metadata)
|
||||
|
||||
|
||||
def test_blob_thresholds_are_validated():
|
||||
with pytest.raises(ValueError, match="inline_size_threshold must be non-negative"):
|
||||
lancedb.blob("payload", inline_size_threshold=-1)
|
||||
with pytest.raises(ValueError, match="dedicated_size_threshold must be positive"):
|
||||
lancedb.blob("payload", dedicated_size_threshold=0)
|
||||
with pytest.raises(ValueError, match="pack_file_size_threshold must be positive"):
|
||||
lancedb.blob("payload", pack_file_size_threshold=0)
|
||||
with pytest.raises(TypeError, match="inline_size_threshold must be an int"):
|
||||
lancedb.blob("payload", inline_size_threshold="8192")
|
||||
|
||||
|
||||
def test_blob_type_works_without_pylance():
|
||||
script = _HIDE_LANCE_BLOB + textwrap.dedent(
|
||||
"""\
|
||||
@@ -376,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)
|
||||
@@ -512,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")
|
||||
@@ -526,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",
|
||||
@@ -536,6 +916,73 @@ def test_fetch_blobs_round_trip():
|
||||
assert [blobs[0].as_py(), blobs[1].as_py()] == [b"alpha", b"beta"]
|
||||
|
||||
|
||||
def test_add_overwrite_preserves_blob_v2(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table("overwrite_blob", schema=schema)
|
||||
table.add([{"id": 1, "image": b"old"}])
|
||||
|
||||
payload = b"x" * 200_000
|
||||
table.add([{"id": 2, "image": payload}], mode="overwrite")
|
||||
|
||||
assert table.count_rows() == 1
|
||||
assert blob_v2_column_paths(table.schema) == ["image"]
|
||||
assert table.blob_columns() == ["image"]
|
||||
assert table.fetch_blobs("image", [_row_ids_by_id(table)[2]]).to_pylist() == [
|
||||
payload
|
||||
]
|
||||
assert list(tmp_path.rglob("*.blob"))
|
||||
|
||||
|
||||
def test_add_overwrite_with_blob_still_infers_vector_schema():
|
||||
db = lancedb.connect("memory:///overwrite_blob_vector")
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("vector", pa.list_(pa.float32(), 2)),
|
||||
lancedb.blob("image"),
|
||||
]
|
||||
)
|
||||
table = db.create_table("overwrite_blob_vector", schema=schema)
|
||||
|
||||
table.add([{"vector": [1.0, 2.0, 3.0], "image": b"data"}], mode="overwrite")
|
||||
|
||||
assert table.schema.field("vector").type == pa.list_(pa.float32(), 3)
|
||||
assert table.blob_columns() == ["image"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_overwrite_preserves_blob_v2_async():
|
||||
db = await lancedb.connect_async("memory:///overwrite_blob_async")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = await db.create_table("overwrite_blob_async", schema=schema)
|
||||
|
||||
await table.add([{"id": 1, "image": b"updated"}], mode="overwrite")
|
||||
|
||||
assert blob_v2_column_paths(await table.schema()) == ["image"]
|
||||
assert await table.blob_columns() == ["image"]
|
||||
hits = await table.query().with_row_id().to_arrow()
|
||||
assert (await table.fetch_blobs("image", hits)).to_pylist() == [b"updated"]
|
||||
|
||||
|
||||
def test_add_overwrite_preserves_nested_blob_v2():
|
||||
db = lancedb.connect("memory:///overwrite_nested_blob")
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("info", pa.struct([lancedb.blob("image")])),
|
||||
]
|
||||
)
|
||||
table = db.create_table("overwrite_nested_blob", schema=schema)
|
||||
|
||||
table.add([{"id": 1, "info": {"image": b"nested"}}], mode="overwrite")
|
||||
|
||||
assert blob_v2_column_paths(table.schema) == ["info.image"]
|
||||
assert table.blob_columns() == ["info.image"]
|
||||
assert table.fetch_blobs("info.image", [_row_ids_by_id(table)[1]]).to_pylist() == [
|
||||
b"nested"
|
||||
]
|
||||
|
||||
|
||||
def test_merge_insert_writes_python_bytes():
|
||||
table = _blob_table("merge_bytes", [{"id": 1, "image": b"before"}])
|
||||
result = (
|
||||
@@ -551,6 +998,50 @@ def test_merge_insert_writes_python_bytes():
|
||||
assert blobs.to_pylist() == [b"updated", b"inserted"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("write", ["merge_insert", "create_table"])
|
||||
def test_uri_string_blob_write_reaches_external_base_validation(tmp_path, write):
|
||||
blob_path = tmp_path / "payload.bin"
|
||||
blob_path.write_bytes(b"payload")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
rows = pa.table(
|
||||
{"id": pa.array([1], pa.int64()), "image": pa.array([blob_path.as_uri()])}
|
||||
)
|
||||
db = lancedb.connect(tmp_path)
|
||||
|
||||
# These APIs do not expose add's outside-base opt-in. The URI should reach
|
||||
# the blob policy check instead of failing at PyArrow's string-to-struct cast.
|
||||
with pytest.raises(ValueError, match="outside registered external bases"):
|
||||
if write == "merge_insert":
|
||||
table = db.create_table("merge_uri", schema=schema)
|
||||
table.merge_insert("id").when_not_matched_insert_all().execute(rows)
|
||||
else:
|
||||
db.create_table("create_uri", data=rows, schema=schema)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("write", ["merge_insert", "create_table"])
|
||||
@pytest.mark.parametrize(
|
||||
"source_type, value",
|
||||
[
|
||||
(pa.binary(4), b"abcd"),
|
||||
(pa.list_(pa.uint8()), [1, 2]),
|
||||
(pa.int64(), 42),
|
||||
],
|
||||
)
|
||||
def test_blob_write_rejects_unsupported_input(write, source_type, value):
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
rows = pa.table(
|
||||
{"id": pa.array([1], pa.int64()), "image": pa.array([value], source_type)}
|
||||
)
|
||||
db = lancedb.connect("memory:///")
|
||||
|
||||
with pytest.raises(ValueError, match="cannot coerce column 'image' with type"):
|
||||
if write == "merge_insert":
|
||||
table = db.create_table("merge_invalid", schema=schema)
|
||||
table.merge_insert("id").when_not_matched_insert_all().execute(rows)
|
||||
else:
|
||||
db.create_table("create_invalid", data=rows, schema=schema)
|
||||
|
||||
|
||||
def test_merge_insert_bytes_after_reopen_without_touching_blob_type(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
@@ -777,12 +1268,29 @@ def test_merge_insert_list_blob_dicts():
|
||||
|
||||
def test_list_blob_column_queries_as_raw_descriptors():
|
||||
table = _list_blob_table("list_query")
|
||||
assert table.blob_columns() == []
|
||||
hits = table.search().limit(10).to_arrow()
|
||||
element = hits.schema.field("images").type.value_type
|
||||
assert pa.types.is_struct(element)
|
||||
assert "_lance_row_id" not in element.names
|
||||
with pytest.raises(ValueError, match="expected struct before segment"):
|
||||
table.fetch_blobs("images.image", [0])
|
||||
for column in ("images", "images.image"):
|
||||
with pytest.raises(ValueError, match="blobs inside lists cannot be fetched"):
|
||||
table.fetch_blobs(column, [0])
|
||||
with pytest.raises(ValueError, match="blobs inside lists cannot be fetched"):
|
||||
table.fetch_blob_files(column, [0])
|
||||
with pytest.raises(ValueError, match="blobs inside lists cannot be fetched"):
|
||||
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():
|
||||
@@ -830,6 +1338,38 @@ def test_fetch_blobs_accepts_query_result():
|
||||
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"gamma"}
|
||||
|
||||
|
||||
def test_fetch_blob_apis_accept_integer_arrow_row_ids():
|
||||
table = _blob_table(
|
||||
"arrow_row_ids",
|
||||
[{"id": 1, "image": b"alpha"}, {"id": 2, "image": b"beta"}],
|
||||
)
|
||||
by_id = _row_ids_by_id(table)
|
||||
ids = [by_id[2], by_id[1]]
|
||||
|
||||
for row_ids in (
|
||||
pa.array(ids, type=pa.uint64()),
|
||||
pa.chunked_array([ids[:1], ids[1:]], type=pa.int64()),
|
||||
):
|
||||
assert table.fetch_blobs("image", row_ids).to_pylist() == [b"beta", b"alpha"]
|
||||
files = table.fetch_blob_files("image", row_ids)
|
||||
assert [file.read() for file in files] == [b"beta", b"alpha"]
|
||||
|
||||
|
||||
def test_fetch_blob_apis_validate_arrow_and_negative_row_ids():
|
||||
table = _blob_table("invalid_arrow_row_ids", [{"id": 1, "image": b"alpha"}])
|
||||
hits = table.search().with_row_id(True).to_arrow()
|
||||
|
||||
for fetch in (table.fetch_blobs, table.fetch_blob_files):
|
||||
with pytest.raises(ValueError, match="non-negative"):
|
||||
fetch("image", [-1])
|
||||
with pytest.raises(ValueError, match="non-negative"):
|
||||
fetch("image", pa.array([-1], type=pa.int64()))
|
||||
with pytest.raises(ValueError, match="without nulls"):
|
||||
fetch("image", pa.array([None], type=pa.int64()))
|
||||
with pytest.raises(ValueError, match="integer type"):
|
||||
fetch("image", hits["image"])
|
||||
|
||||
|
||||
def test_fetch_blobs_after_compact_with_stable_row_ids(tmp_path):
|
||||
db = lancedb.connect(
|
||||
tmp_path, storage_options={"new_table_enable_stable_row_ids": "true"}
|
||||
@@ -932,13 +1472,14 @@ def test_add_list_of_dicts_to_blob_list_column(large_list):
|
||||
table = db.create_table(f"blob_{large_list}_list_add", data=seed)
|
||||
|
||||
table.add([{"id": 1, "items": [None]}])
|
||||
table.add([{"id": 2, "items": [b"a", None, b""]}])
|
||||
table.add(
|
||||
[{"id": 2, "items": [b"a", None]}],
|
||||
[{"id": 3, "items": [b"b", None]}],
|
||||
on_bad_vectors="fill",
|
||||
)
|
||||
|
||||
ids = table.search().select(["id"]).to_arrow()["id"].to_pylist()
|
||||
assert sorted(ids) == [0, 1, 2]
|
||||
assert sorted(ids) == [0, 1, 2, 3]
|
||||
assert pa.types.is_large_list(table.schema.field("items").type) is large_list
|
||||
|
||||
|
||||
@@ -961,6 +1502,15 @@ def test_fetch_blob_ranges_aligns_repeated_ranges_and_nulls():
|
||||
assert ranges.to_pylist() == [b"cde", None, b"ab", b"cde", b""]
|
||||
|
||||
|
||||
def test_fetch_blob_ranges_accepts_three_item_sequences():
|
||||
table = _blob_table("range_sequences", [{"id": 1, "image": b"abcdef"}])
|
||||
row_id = _row_ids_by_id(table)[1]
|
||||
|
||||
ranges = table.fetch_blob_ranges("image", [[row_id, 1, 2], (row_id, 4, 2)])
|
||||
|
||||
assert ranges.to_pylist() == [b"bc", b"ef"]
|
||||
|
||||
|
||||
def test_fetch_blob_ranges_validates_requests():
|
||||
table = _blob_table("range_validation", [{"id": 1, "image": b"abc"}])
|
||||
row_id = _row_ids_by_id(table)[1]
|
||||
@@ -975,6 +1525,30 @@ def test_fetch_blob_ranges_validates_requests():
|
||||
table.fetch_blob_ranges("image", [(2**64 - 1, 0, 1)])
|
||||
_assert_missing_blob_row_ids(exc_info)
|
||||
|
||||
for index, name in enumerate(("row id", "offset", "length")):
|
||||
request = [row_id, 0, 1]
|
||||
request[index] = -1
|
||||
with pytest.raises(
|
||||
ValueError, match=f"blob range {name} must be a non-negative"
|
||||
):
|
||||
table.fetch_blob_ranges("image", [request])
|
||||
|
||||
with pytest.raises(ValueError, match="three values"):
|
||||
table.fetch_blob_ranges("image", [[row_id, 0]])
|
||||
|
||||
|
||||
def test_fetch_blob_apis_reject_row_past_fragment_end():
|
||||
table = _blob_table(
|
||||
"past_fragment_end",
|
||||
[
|
||||
{"id": 1, "image": b"a"},
|
||||
{"id": 2, "image": b"b"},
|
||||
{"id": 3, "image": b"c"},
|
||||
],
|
||||
)
|
||||
_assert_fetch_apis_reject_missing_row_ids(table, [3])
|
||||
_assert_fetch_apis_reject_missing_row_ids(table, [3, 0])
|
||||
|
||||
|
||||
def test_fetch_blob_apis_reject_missing_fragment_row_addr():
|
||||
table = _blob_table("missing_frag", [{"id": 1, "image": b"x"}])
|
||||
@@ -992,6 +1566,19 @@ def test_fetch_blob_apis_reject_deleted_row_ids():
|
||||
_assert_fetch_apis_reject_missing_row_ids(table, [by_id[2], by_id[1]])
|
||||
|
||||
|
||||
def test_fetch_blob_apis_reject_deleted_stable_row_ids(tmp_path):
|
||||
db = lancedb.connect(
|
||||
tmp_path, storage_options={"new_table_enable_stable_row_ids": "true"}
|
||||
)
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table("stable_deleted", schema=schema)
|
||||
table.add([{"id": 1, "image": b"one"}, {"id": 2, "image": b"two"}])
|
||||
by_id = _row_ids_by_id(table)
|
||||
table.delete("id = 2")
|
||||
|
||||
_assert_fetch_apis_reject_missing_row_ids(table, [by_id[2], by_id[1]])
|
||||
|
||||
|
||||
def test_fetch_blob_ranges_empty_requests_returns_empty_array():
|
||||
table = _blob_table("range_empty", [{"id": 1, "image": b"x"}])
|
||||
assert table.fetch_blob_ranges("image", []).to_pylist() == []
|
||||
@@ -1030,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
|
||||
@@ -1040,7 +1632,7 @@ def test_fetch_blob_files_lazy_read():
|
||||
assert handles[0].read() == payload
|
||||
handles[0].close()
|
||||
assert handles[0].closed
|
||||
with pytest.raises(RuntimeError, match="already closed"):
|
||||
with pytest.raises(ValueError, match="read of closed file"):
|
||||
handles[0].read()
|
||||
|
||||
|
||||
@@ -1089,8 +1681,8 @@ def test_blob_file_finalizer_and_close_on_runtime_worker():
|
||||
explicit_handle.close() # Close remains idempotent after the worker call.
|
||||
try:
|
||||
explicit_handle.read()
|
||||
except RuntimeError as error:
|
||||
assert "already closed" in str(error)
|
||||
except ValueError as error:
|
||||
assert "read of closed file" in str(error)
|
||||
else:
|
||||
raise AssertionError("a closed blob file was readable")
|
||||
print("worker close completed")
|
||||
@@ -1148,6 +1740,9 @@ async def test_async_fetch_blob_files_lazy_read():
|
||||
handles = await table.fetch_blob_files("image", hits)
|
||||
assert len(handles) == 1
|
||||
assert await handles[0].aread() == payload
|
||||
handles[0].close()
|
||||
with pytest.raises(ValueError, match="read of closed file"):
|
||||
await handles[0].aread()
|
||||
|
||||
|
||||
def test_fetch_blobs_from_query_result_without_row_id_raises():
|
||||
@@ -1204,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")
|
||||
@@ -1283,6 +1994,49 @@ def test_blob_file_seek_read_and_read_range():
|
||||
handle.seek(0, 99)
|
||||
|
||||
|
||||
def test_blob_file_matches_file_error_and_range_contract():
|
||||
payload = b"abcdef"
|
||||
table = _blob_table("file_contract", [{"id": 1, "image": payload}])
|
||||
row_id = _row_ids_by_id(table)[1]
|
||||
handle = table.fetch_blob_files("image", [row_id])[0]
|
||||
|
||||
handle.seek(2)
|
||||
for offset, whence in [(-1, io.SEEK_SET), (-3, io.SEEK_CUR), (-7, io.SEEK_END)]:
|
||||
with pytest.raises(ValueError, match="negative seek value -1"):
|
||||
handle.seek(offset, whence)
|
||||
assert handle.tell() == 2
|
||||
|
||||
assert handle.read(None) == payload[2:]
|
||||
handle.seek(2)
|
||||
assert handle.read1(2) == payload[2:4]
|
||||
assert handle.read1() == payload[4:]
|
||||
handle.seek(2)
|
||||
assert handle.read_ranges([(0, 2), (4, 2), (2, 0)]) == [b"ab", b"ef", b""]
|
||||
assert handle.tell() == 2
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
handle.read_range(5, 2)
|
||||
with pytest.raises(ValueError):
|
||||
handle.read_ranges([(0, 1), (5, 2)])
|
||||
assert handle.tell() == 2
|
||||
|
||||
handle.close()
|
||||
with pytest.raises(ValueError, match="read of closed file"):
|
||||
handle.read(1)
|
||||
with pytest.raises(ValueError, match="read of closed file"):
|
||||
handle.read(None)
|
||||
with pytest.raises(ValueError, match="read of closed file"):
|
||||
handle.read1(1)
|
||||
with pytest.raises(ValueError, match="read of closed file"):
|
||||
handle.readall()
|
||||
with pytest.raises(ValueError, match="readinto of closed file"):
|
||||
handle.readinto(bytearray(1))
|
||||
with pytest.raises(ValueError):
|
||||
handle.read_range(0, 1)
|
||||
with pytest.raises(ValueError):
|
||||
handle.read_ranges([])
|
||||
|
||||
|
||||
def test_fetch_blob_files_from_query_partial_read():
|
||||
payload = _identifiable_payload(65536)
|
||||
table = _blob_table("query_partial", [{"id": 1, "image": payload}])
|
||||
|
||||
@@ -263,3 +263,45 @@ async def test_database_names_invalid_limit(catalog_server, asynchronous, page_l
|
||||
next(names)
|
||||
assert names.num_page_results() == 0
|
||||
assert requests == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
async def test_catalog_rejects_invalid_database_names_before_io(
|
||||
catalog_server, asynchronous
|
||||
):
|
||||
endpoint, requests, _ = catalog_server
|
||||
catalog = (
|
||||
await lancedb.connect_catalog_async(endpoint)
|
||||
if asynchronous
|
||||
else lancedb.connect_catalog(endpoint)
|
||||
)
|
||||
for name in [
|
||||
"",
|
||||
".",
|
||||
"..",
|
||||
"/db",
|
||||
"db/",
|
||||
"a//b",
|
||||
"a/./b",
|
||||
"a/../b",
|
||||
"a b",
|
||||
"a:b",
|
||||
"a$b",
|
||||
"a%b",
|
||||
"a?b",
|
||||
"a#b",
|
||||
"a\\b",
|
||||
"café",
|
||||
"a\nb",
|
||||
]:
|
||||
for method in [
|
||||
catalog.create_database,
|
||||
catalog.connect_database,
|
||||
catalog.drop_database,
|
||||
]:
|
||||
with pytest.raises(ValueError, match="Invalid database name"):
|
||||
result = method(name)
|
||||
if asynchronous:
|
||||
await result
|
||||
assert requests == []
|
||||
|
||||
@@ -111,6 +111,26 @@ async def _create_blob_v2_query_table_async(db, name):
|
||||
return table
|
||||
|
||||
|
||||
def _nested_blob_v2_query_schema():
|
||||
return pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("vector", pa.list_(pa.float32(), list_size=2)),
|
||||
pa.field(
|
||||
"info",
|
||||
pa.struct([pa.field("label", pa.string()), lancedb.blob("blob")]),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _nested_blob_v2_query_rows():
|
||||
return [
|
||||
{"id": 1, "vector": [1.0, 0.0], "info": {"label": "one", "blob": b"one"}},
|
||||
{"id": 2, "vector": [2.0, 0.0], "info": {"label": "two", "blob": b"two"}},
|
||||
]
|
||||
|
||||
|
||||
def _assert_lazy_blob(value, expected: bytes):
|
||||
assert hasattr(value, "readall")
|
||||
assert value.readall() == expected
|
||||
@@ -503,24 +523,223 @@ def test_vector_query_to_pandas_blob_mode_requires_native_path(tmp_db):
|
||||
storage_options=LEGACY_BLOB_STORAGE_OPTIONS,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="Lance native pandas conversion"):
|
||||
with pytest.raises(RuntimeError, match="no fetchable blob v2 source"):
|
||||
table.search([1.0, 0.0]).select(["blob", "vector"]).limit(1).to_pandas(
|
||||
blob_mode="lazy"
|
||||
)
|
||||
|
||||
|
||||
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",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("blob_mode", ["default", "lazy", "bytes", "descriptions"])
|
||||
@pytest.mark.parametrize("with_row_id", [None, True, False])
|
||||
def test_vector_query_to_pandas_blob_v2_modes(tmp_db, blob_mode, with_row_id):
|
||||
pytest.importorskip("lance")
|
||||
table = _create_blob_v2_query_table(tmp_db, f"test_vector_blob_v2_{blob_mode}")
|
||||
|
||||
query = table.search([1.0, 0.0]).select(["id", "blob", "vector"]).limit(1)
|
||||
if with_row_id is not None:
|
||||
query = query.with_row_id(with_row_id)
|
||||
df = (
|
||||
query.to_pandas()
|
||||
if blob_mode == "default"
|
||||
else query.to_pandas(blob_mode=blob_mode)
|
||||
)
|
||||
|
||||
assert df["id"].tolist() == [1]
|
||||
assert ("_rowid" in df.columns) == (with_row_id is True)
|
||||
if blob_mode == "bytes":
|
||||
assert df["blob"].tolist() == [b"one"]
|
||||
else:
|
||||
descriptor = df["blob"].iloc[0]
|
||||
assert descriptor != b"one"
|
||||
assert not hasattr(descriptor, "readall")
|
||||
assert "_lance_row_id" not in descriptor
|
||||
|
||||
|
||||
@pytest.mark.parametrize("blob_mode", ["lazy", "bytes", "descriptions"])
|
||||
def test_blob_v2_non_scan_query_shapes_to_pandas(tmp_db, blob_mode):
|
||||
pytest.importorskip("lance")
|
||||
table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_query_shapes")
|
||||
table.create_index("tag", config=FTS())
|
||||
|
||||
queries = [
|
||||
table.search("keep", query_type="fts").limit(2),
|
||||
table.search(query_type="hybrid").vector([1.0, 0.0]).text("keep").limit(2),
|
||||
table.take_offsets([1, 2]).select(["id", "blob"]),
|
||||
]
|
||||
for query in queries:
|
||||
df = query.to_pandas(blob_mode=blob_mode)
|
||||
assert len(df) == 2
|
||||
if blob_mode == "bytes":
|
||||
assert df["blob"].tolist() == [
|
||||
{1: b"one", 2: b"two", 3: b"three", 4: b"four"}[row_id]
|
||||
for row_id in df["id"]
|
||||
]
|
||||
else:
|
||||
assert all("_lance_row_id" not in blob for blob in df["blob"])
|
||||
assert "_rowid" not in df.columns
|
||||
|
||||
|
||||
@pytest.mark.parametrize("rename_parent", [False, True])
|
||||
def test_vector_query_to_pandas_nested_blob_bytes(tmp_db, rename_parent):
|
||||
pytest.importorskip("lance")
|
||||
table = tmp_db.create_table(
|
||||
f"test_vector_nested_blob_{rename_parent}",
|
||||
schema=_nested_blob_v2_query_schema(),
|
||||
)
|
||||
table.add(_nested_blob_v2_query_rows())
|
||||
|
||||
projection = {"id": "id", "renamed": "info"} if rename_parent else ["id", "info"]
|
||||
df = (
|
||||
table.search([1.0, 0.0])
|
||||
.select(projection)
|
||||
.limit(2)
|
||||
.to_pandas(blob_mode="bytes")
|
||||
)
|
||||
|
||||
assert df["id"].tolist() == [1, 2]
|
||||
parent = "renamed" if rename_parent else "info"
|
||||
assert df[parent].tolist() == [
|
||||
{"label": "one", "blob": b"one"},
|
||||
{"label": "two", "blob": b"two"},
|
||||
]
|
||||
for blob_mode in ("lazy", "descriptions"):
|
||||
descriptors = (
|
||||
table.search([1.0, 0.0])
|
||||
.select(projection)
|
||||
.limit(1)
|
||||
.to_pandas(blob_mode=blob_mode)
|
||||
)
|
||||
assert "_lance_row_id" not in descriptors[parent].iloc[0]["blob"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_vector_query_to_pandas_nested_blob_bytes(tmp_db_async):
|
||||
pytest.importorskip("lance")
|
||||
table = await tmp_db_async.create_table(
|
||||
"test_async_vector_nested_blob",
|
||||
schema=_nested_blob_v2_query_schema(),
|
||||
)
|
||||
await table.add(_nested_blob_v2_query_rows())
|
||||
|
||||
df = await (
|
||||
table.vector_search([1.0, 0.0])
|
||||
.select(["id", "info"])
|
||||
.limit(2)
|
||||
.to_pandas(blob_mode="bytes")
|
||||
)
|
||||
|
||||
assert df["id"].tolist() == [1, 2]
|
||||
assert df["info"].tolist() == [
|
||||
{"label": "one", "blob": b"one"},
|
||||
{"label": "two", "blob": b"two"},
|
||||
]
|
||||
|
||||
|
||||
def test_blob_v2_vector_projection_and_plain_scan_options(tmp_db, monkeypatch):
|
||||
pytest.importorskip("lance")
|
||||
table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_pandas_options")
|
||||
|
||||
df = (
|
||||
table.search([1.0, 0.0])
|
||||
.select({"payload": "blob", "id": "id"})
|
||||
.limit(1)
|
||||
.to_pandas(blob_mode="bytes")
|
||||
)
|
||||
assert df["payload"].tolist() == [b"one"]
|
||||
|
||||
query = table.search().where("id = 1")
|
||||
|
||||
def fail_to_arrow(*args, **kwargs):
|
||||
raise AssertionError("the native scanner should handle bytes with flatten")
|
||||
|
||||
monkeypatch.setattr(query, "to_arrow", fail_to_arrow)
|
||||
assert query.to_pandas(blob_mode="bytes", flatten=True)["blob"].tolist() == [b"one"]
|
||||
|
||||
flattened_lazy = table.search().where("id = 1").to_pandas(flatten=True)
|
||||
assert flattened_lazy["id"].tolist() == [1]
|
||||
assert not any("_lance_row_id" in col for col in flattened_lazy.columns)
|
||||
|
||||
timed = (
|
||||
table.search()
|
||||
.where("id = 1")
|
||||
.to_pandas(blob_mode="bytes", timeout=timedelta(seconds=10))
|
||||
)
|
||||
assert timed["blob"].tolist() == [b"one"]
|
||||
|
||||
timed_lazy = table.search().where("id = 1").to_pandas(timeout=timedelta(seconds=10))
|
||||
assert timed_lazy["id"].tolist() == [1]
|
||||
|
||||
|
||||
def test_blob_v2_pandas_preserves_fetch_error(tmp_db, monkeypatch):
|
||||
pytest.importorskip("lance")
|
||||
table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_fetch_error")
|
||||
|
||||
def missing_blob(*args, **kwargs):
|
||||
raise FileNotFoundError("missing obj.bin")
|
||||
|
||||
monkeypatch.setattr(table, "fetch_blobs", missing_blob)
|
||||
with pytest.raises(FileNotFoundError, match="missing obj.bin"):
|
||||
table.search().to_pandas(blob_mode="bytes")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("blob_mode", ["lazy", "bytes", "descriptions"])
|
||||
@pytest.mark.parametrize("with_row_id", [False, True])
|
||||
async def test_async_vector_query_to_pandas_blob_v2_modes(
|
||||
tmp_db_async, blob_mode, with_row_id
|
||||
):
|
||||
pytest.importorskip("lance")
|
||||
table = await _create_blob_v2_query_table_async(
|
||||
tmp_db_async, f"test_async_vector_blob_v2_{blob_mode}"
|
||||
)
|
||||
|
||||
query = table.vector_search([1.0, 0.0]).limit(1)
|
||||
if with_row_id:
|
||||
query = query.with_row_id()
|
||||
df = await query.to_pandas(blob_mode=blob_mode)
|
||||
assert df["id"].tolist() == [1]
|
||||
assert ("_rowid" in df.columns) == with_row_id
|
||||
if blob_mode == "bytes":
|
||||
assert df["blob"].tolist() == [b"one"]
|
||||
else:
|
||||
assert "_lance_row_id" not in df["blob"].iloc[0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_take_query_to_pandas_blob_v2_bytes(tmp_db_async):
|
||||
pytest.importorskip("lance")
|
||||
table = await _create_blob_v2_query_table_async(
|
||||
tmp_db_async, "test_async_take_blob_v2"
|
||||
)
|
||||
|
||||
df = await table.take_offsets([1, 2]).to_pandas(blob_mode="bytes")
|
||||
assert set(df["blob"]) == {b"two", b"three"}
|
||||
|
||||
|
||||
def test_order_by_plain_query(mem_db):
|
||||
|
||||
@@ -4,6 +4,7 @@ from concurrent.futures import ThreadPoolExecutor
|
||||
import contextlib
|
||||
from datetime import timedelta
|
||||
import http.server
|
||||
import io
|
||||
import json
|
||||
import multiprocessing as mp
|
||||
import pickle
|
||||
@@ -2728,6 +2729,11 @@ def test_create_client():
|
||||
assert isinstance(db.client_config, ClientConfig)
|
||||
assert db.client_config.retry_config.retries == 42
|
||||
|
||||
db = lancedb.connect(
|
||||
**mandatory_args, client_config={"blob_request_concurrency": 3}
|
||||
)
|
||||
assert db.client_config.blob_request_concurrency == 3
|
||||
|
||||
with pytest.warns(DeprecationWarning):
|
||||
db = lancedb.connect(**mandatory_args, connection_timeout=42)
|
||||
assert db.client_config.timeout_config.connect_timeout == timedelta(seconds=42)
|
||||
@@ -3463,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):
|
||||
@@ -3494,10 +3539,32 @@ def test_remote_blob_files_are_lazy_seekable_handles():
|
||||
assert empty.read_range(0, 0) == b""
|
||||
alpha.close()
|
||||
assert alpha.closed
|
||||
with pytest.raises(RuntimeError, match="already closed"):
|
||||
with pytest.raises(ValueError, match="read of closed file"):
|
||||
alpha.read_range(0, 1)
|
||||
|
||||
|
||||
def test_remote_blob_file_matches_file_error_and_range_contract():
|
||||
with blob_remote_table() as table:
|
||||
handle = table.fetch_blob_files("image", [10])[0]
|
||||
handle.seek(2)
|
||||
|
||||
with pytest.raises(ValueError, match="negative seek value -1"):
|
||||
handle.seek(-3, io.SEEK_CUR)
|
||||
assert handle.tell() == 2
|
||||
assert handle.read(None) == b"pha"
|
||||
|
||||
handle.seek(2)
|
||||
assert handle.read1(2) == b"ph"
|
||||
assert handle.read_ranges([(0, 2), (3, 2)]) == [b"al", b"ha"]
|
||||
assert handle.tell() == 4
|
||||
with pytest.raises(ValueError):
|
||||
handle.read_range(4, 2)
|
||||
|
||||
handle.close()
|
||||
with pytest.raises(ValueError, match="read of closed file"):
|
||||
handle.read(1)
|
||||
|
||||
|
||||
def test_remote_blob_fetch_accepts_query_table():
|
||||
hits = pa.table({"_rowid": pa.array([10, 20, 30], type=pa.uint64())})
|
||||
|
||||
@@ -3517,6 +3584,56 @@ def test_remote_blob_query_stashes_row_ids_for_fetch():
|
||||
assert blobs.to_pylist() == [b"alpha", None, b"gamma"]
|
||||
|
||||
|
||||
def test_remote_blob_query_to_pandas():
|
||||
with blob_remote_table() as table:
|
||||
query = table.search().select(["id", "image"]).limit(3)
|
||||
df = query.to_pandas(blob_mode="descriptions", split_blocks=True)
|
||||
|
||||
assert df["id"].tolist() == [1, 2, 3]
|
||||
assert "_rowid" not in df.columns
|
||||
assert df["image"].iloc[1] is None
|
||||
assert "_lance_row_id" not in df["image"].iloc[0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("blob_mode", ["default", "lazy", "bytes"])
|
||||
def test_remote_blob_query_to_pandas_rejects_unpinned_fetch(monkeypatch, blob_mode):
|
||||
with blob_remote_table() as table:
|
||||
|
||||
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.search().select(["id", "image"]).limit(3)
|
||||
pandas_kwargs = {} if blob_mode == "default" else {"blob_mode": blob_mode}
|
||||
with pytest.raises(NotImplementedError, match="stable table snapshot"):
|
||||
query.to_pandas(**pandas_kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("blob_mode", ["lazy", "bytes"])
|
||||
def test_remote_nonblob_projection_to_pandas(blob_mode):
|
||||
with blob_remote_table() as table:
|
||||
df = table.search().select(["id"]).to_pandas(blob_mode=blob_mode)
|
||||
|
||||
assert df["id"].tolist() == [1, 2, 3]
|
||||
assert list(df.columns) == ["id"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_remote_blob_query_to_pandas_flatten():
|
||||
with blob_remote_table() as table:
|
||||
df = (
|
||||
await table._table.query()
|
||||
.select(["id", "image"])
|
||||
.limit(3)
|
||||
.to_pandas(blob_mode="descriptions", flatten=True, split_blocks=True)
|
||||
)
|
||||
|
||||
assert df["id"].tolist() == [1, 2, 3]
|
||||
assert df["image.size"].iloc[0] == 5
|
||||
assert "image._lance_row_id" not in df.columns
|
||||
assert "_rowid" not in df.columns
|
||||
|
||||
|
||||
def test_remote_blob_query_survives_a_server_that_ignores_the_row_id_request():
|
||||
def handler(request):
|
||||
if request.path == "/v1/table/test/describe/":
|
||||
@@ -3546,6 +3663,8 @@ def test_remote_blob_query_survives_a_server_that_ignores_the_row_id_request():
|
||||
assert "_lance_row_id" not in hits.schema.field("image").type.names
|
||||
with pytest.raises(ValueError, match="pass a list of row ids"):
|
||||
table.fetch_blobs("image", hits)
|
||||
with pytest.raises(NotImplementedError, match="stable table snapshot"):
|
||||
table.search().select(["id", "image"]).to_pandas(blob_mode="bytes")
|
||||
|
||||
|
||||
def test_remote_blob_byte_apis_not_supported_on_old_server():
|
||||
|
||||
@@ -4016,6 +4016,26 @@ def test_search_nested_vector_no_candidates(mem_db: DBConnection):
|
||||
table.search([0.0, 1.0]).limit(1).to_arrow()
|
||||
|
||||
|
||||
def test_search_single_vector_column_wrong_query_dim(mem_db: DBConnection):
|
||||
"""A query vector whose dimension doesn't match the (only) vector column
|
||||
must still resolve that column, surfacing a dimension-mismatch error from
|
||||
query execution rather than a misleading "no vector column" error."""
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.uint32(), False),
|
||||
pa.field("vector_col", pa.list_(pa.float32(), 32), False),
|
||||
]
|
||||
)
|
||||
table = mem_db.create_table(
|
||||
"wrong_query_dim",
|
||||
data=[{"id": 1, "vector_col": [0.1] * 32}],
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="dim"):
|
||||
table.search([0.01] * 120).limit(1).to_arrow()
|
||||
|
||||
|
||||
def test_compact_cleanup(tmp_db: DBConnection):
|
||||
pytest.importorskip("lance")
|
||||
table = tmp_db.create_table(
|
||||
|
||||
@@ -28,7 +28,13 @@ import pandas as pd
|
||||
import polars as pl
|
||||
import pytest
|
||||
import lancedb
|
||||
from lancedb.util import flatten_columns, get_uri_scheme, join_uri, value_to_sql
|
||||
from lancedb.util import (
|
||||
flatten_columns,
|
||||
get_uri_scheme,
|
||||
inf_vector_column_query,
|
||||
join_uri,
|
||||
value_to_sql,
|
||||
)
|
||||
from utils import exception_output
|
||||
|
||||
|
||||
@@ -987,6 +993,39 @@ def test_cast_to_target_schema_coerces_binary_to_blob_v2():
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"string_type", [pa.string(), pa.large_string(), pa.string_view()]
|
||||
)
|
||||
def test_cast_to_target_schema_coerces_string_to_blob_v2_uri(string_type):
|
||||
data = pa.table({"image": pa.array(["file:///payload", None], type=string_type)})
|
||||
target = pa.schema([lancedb.blob("image")])
|
||||
|
||||
output = _cast_to_target_schema(data.to_reader(), target).read_all()
|
||||
|
||||
image = output["image"].chunk(0)
|
||||
assert type(image.type) is lancedb.BlobType
|
||||
assert image.storage.to_pylist() == [
|
||||
{"data": None, "uri": "file:///payload", "position": None, "size": None},
|
||||
None,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"source_type, value",
|
||||
[
|
||||
(pa.binary(4), b"abcd"),
|
||||
(pa.list_(pa.uint8()), [1, 2]),
|
||||
(pa.int64(), 42),
|
||||
],
|
||||
)
|
||||
def test_cast_to_target_schema_rejects_unsupported_blob_input(source_type, value):
|
||||
data = pa.table({"image": pa.array([value], type=source_type)})
|
||||
target = pa.schema([lancedb.blob("image")])
|
||||
|
||||
with pytest.raises(ValueError, match="cannot coerce column 'image' with type"):
|
||||
_cast_to_target_schema(data.to_reader(), target)
|
||||
|
||||
|
||||
def test_cast_to_target_schema_coerces_binary_to_metadata_blob_struct():
|
||||
storage = lancedb.blob("image").type.storage_type
|
||||
target = pa.schema(
|
||||
@@ -1218,3 +1257,41 @@ def test_infer_vector_column_raises_clear_error(tmp_path):
|
||||
# Plain vector search on a table with no vector column should raise
|
||||
# a clear ValueError, not a cryptic TypeError.
|
||||
table.search([1.0, 2.0]).to_list()
|
||||
|
||||
|
||||
def test_inf_vector_column_query_dim_mismatch_still_resolves():
|
||||
"""Regression: a query dimension that matches no column must not hide an
|
||||
otherwise-unambiguous vector column behind "no vector column in the data".
|
||||
|
||||
`dim` exists to disambiguate between *multiple* vector columns of
|
||||
different sizes. When the schema has exactly one vector column, that
|
||||
column is still the only sensible one to search; a dimension mismatch is
|
||||
a query-execution concern, not a missing-column one. This previously
|
||||
broke remote search (#4409 wired up client-side inference there), which
|
||||
expects the server to reject the mismatched dimension itself.
|
||||
"""
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.uint32(), False),
|
||||
pa.field("vector_col", pa.list_(pa.float32(), 32), False),
|
||||
pa.field("s", pa.string(), False),
|
||||
]
|
||||
)
|
||||
|
||||
assert inf_vector_column_query(schema, dim=32) == "vector_col"
|
||||
assert inf_vector_column_query(schema, dim=120) == "vector_col"
|
||||
assert inf_vector_column_query(schema) == "vector_col"
|
||||
|
||||
|
||||
def test_inf_vector_column_query_dim_mismatch_with_multiple_columns():
|
||||
"""With more than one vector column, a dimension that matches none of
|
||||
them is still ambiguous and must raise, not silently pick one."""
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("vector1", pa.list_(pa.float32(), 10), False),
|
||||
pa.field("vector2", pa.list_(pa.float32(), 20), False),
|
||||
]
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="more than one vector column"):
|
||||
inf_vector_column_query(schema, dim=99)
|
||||
|
||||
@@ -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`"
|
||||
|
||||
@@ -0,0 +1,280 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use crate::error::PythonErrorExt;
|
||||
use crate::runtime::future_into_py;
|
||||
use lancedb::authz::*;
|
||||
use pyo3::exceptions::{PyRuntimeError, PyValueError};
|
||||
use pyo3::types::{PyAnyMethods, PyBytes, PyString};
|
||||
use pyo3::{Bound, PyAny, PyRef, PyResult, pyclass, pymethods};
|
||||
use std::sync::Arc;
|
||||
|
||||
#[pyclass]
|
||||
pub struct AuthorizationClient {
|
||||
pub(crate) inner: Arc<dyn Authorization>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl AuthorizationClient {
|
||||
fn list_acls<'py>(self_: PyRef<'py, Self>, request: &str) -> PyResult<Bound<'py, PyAny>> {
|
||||
let request: ListAccessControlEntriesRequest = serde_json::from_str(request)
|
||||
.map_err(|_| PyValueError::new_err("Invalid authorization request"))?;
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let result = inner.list_acls(request).await.infer_error()?;
|
||||
serde_json::to_string(&result).map_err(|err| PyRuntimeError::new_err(err.to_string()))
|
||||
})
|
||||
}
|
||||
fn add_acl<'py>(self_: PyRef<'py, Self>, request: &str) -> PyResult<Bound<'py, PyAny>> {
|
||||
let request: AccessControlEntryRequest = serde_json::from_str(request)
|
||||
.map_err(|_| PyValueError::new_err("Invalid authorization request"))?;
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.add_acl(request).await.infer_error()
|
||||
})
|
||||
}
|
||||
fn delete_acl<'py>(self_: PyRef<'py, Self>, request: &str) -> PyResult<Bound<'py, PyAny>> {
|
||||
let request: AccessControlEntryRequest = serde_json::from_str(request)
|
||||
.map_err(|_| PyValueError::new_err("Invalid authorization request"))?;
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.delete_acl(request).await.infer_error()
|
||||
})
|
||||
}
|
||||
fn list_role_bindings<'py>(
|
||||
self_: PyRef<'py, Self>,
|
||||
request: &str,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let request: ListRoleBindingsRequest = serde_json::from_str(request)
|
||||
.map_err(|_| PyValueError::new_err("Invalid authorization request"))?;
|
||||
request.validate().infer_error()?;
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let result = inner.list_role_bindings(request).await.infer_error()?;
|
||||
serde_json::to_string(&result).map_err(|err| PyRuntimeError::new_err(err.to_string()))
|
||||
})
|
||||
}
|
||||
fn add_role_binding<'py>(
|
||||
self_: PyRef<'py, Self>,
|
||||
request: &str,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let request: RoleBindingRequest = serde_json::from_str(request)
|
||||
.map_err(|_| PyValueError::new_err("Invalid authorization request"))?;
|
||||
request.validate().infer_error()?;
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.add_role_binding(request).await.infer_error()
|
||||
})
|
||||
}
|
||||
fn delete_role_binding<'py>(
|
||||
self_: PyRef<'py, Self>,
|
||||
request: &str,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let request: RoleBindingRequest = serde_json::from_str(request)
|
||||
.map_err(|_| PyValueError::new_err("Invalid authorization request"))?;
|
||||
request.validate().infer_error()?;
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.delete_role_binding(request).await.infer_error()
|
||||
})
|
||||
}
|
||||
fn get_principal<'py>(self_: PyRef<'py, Self>, request: &str) -> PyResult<Bound<'py, PyAny>> {
|
||||
let request: GetPrincipalRequest = serde_json::from_str(request)
|
||||
.map_err(|_| PyValueError::new_err("Invalid authorization request"))?;
|
||||
request.validate().infer_error()?;
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let result = inner.get_principal(request).await.infer_error()?;
|
||||
serde_json::to_string(&result).map_err(|err| PyRuntimeError::new_err(err.to_string()))
|
||||
})
|
||||
}
|
||||
fn get_group<'py>(self_: PyRef<'py, Self>, request: &str) -> PyResult<Bound<'py, PyAny>> {
|
||||
let request: GetGroupRequest = serde_json::from_str(request)
|
||||
.map_err(|_| PyValueError::new_err("Invalid authorization request"))?;
|
||||
request.validate().infer_error()?;
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let result = inner.get_group(request).await.infer_error()?;
|
||||
serde_json::to_string(&result).map_err(|err| PyRuntimeError::new_err(err.to_string()))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// The shared privilege parser and canonical wire names.
|
||||
#[pyclass(frozen)]
|
||||
pub struct AuthzPrivilege {
|
||||
inner: Privilege,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl AuthzPrivilege {
|
||||
#[new]
|
||||
fn new(value: &str) -> PyResult<Self> {
|
||||
Ok(Self {
|
||||
inner: value.parse().map_err(PyValueError::new_err)?,
|
||||
})
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn known() -> Vec<String> {
|
||||
Privilege::KNOWN.iter().map(ToString::to_string).collect()
|
||||
}
|
||||
|
||||
fn __str__(&self) -> String {
|
||||
self.inner.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(frozen)]
|
||||
pub struct AuthzSubject {
|
||||
inner: Subject,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl AuthzSubject {
|
||||
#[new]
|
||||
fn new(kind: &str, value: &Bound<'_, PyAny>) -> PyResult<Self> {
|
||||
let value: String = value
|
||||
.extract()
|
||||
.map_err(|_| PyValueError::new_err("Subject value must be a nonempty string"))?;
|
||||
Ok(Self {
|
||||
inner: Subject::new(kind.parse().infer_error()?, value).infer_error()?,
|
||||
})
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn from_canonical(value: &str) -> PyResult<Self> {
|
||||
Ok(Self {
|
||||
inner: Subject::from_canonical(value).infer_error()?,
|
||||
})
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn kind(&self) -> &str {
|
||||
self.inner.kind().as_str()
|
||||
}
|
||||
#[getter]
|
||||
fn value(&self) -> &str {
|
||||
self.inner.value()
|
||||
}
|
||||
#[getter]
|
||||
fn display_value(&self) -> &str {
|
||||
self.inner.display_value()
|
||||
}
|
||||
|
||||
#[pyo3(signature = (context=None))]
|
||||
fn to_wire(&self, context: Option<&str>) -> PyResult<String> {
|
||||
match context {
|
||||
None => {}
|
||||
Some("principal") => self.inner.validate_principal().infer_error()?,
|
||||
Some("group") => self.inner.validate_group().infer_error()?,
|
||||
Some("role_member") => self.inner.validate_role_member().infer_error()?,
|
||||
Some(_) => return Err(PyValueError::new_err("Invalid subject context")),
|
||||
}
|
||||
Ok(self.inner.to_wire())
|
||||
}
|
||||
|
||||
fn __repr__(&self) -> String {
|
||||
format!("{:?}", self.inner)
|
||||
}
|
||||
fn __str__(&self) -> String {
|
||||
self.inner.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(frozen)]
|
||||
pub struct AuthzObject {
|
||||
inner: Object,
|
||||
}
|
||||
|
||||
impl AuthzObject {
|
||||
fn validated(inner: Object) -> PyResult<Self> {
|
||||
inner.validate().map_err(PyValueError::new_err)?;
|
||||
Ok(Self { inner })
|
||||
}
|
||||
}
|
||||
|
||||
fn namespace_path(value: &Bound<'_, PyAny>) -> PyResult<NamespacePath> {
|
||||
if value.is_instance_of::<PyString>() || value.is_instance_of::<PyBytes>() {
|
||||
return Err(PyValueError::new_err(
|
||||
"namespace must be a sequence of components",
|
||||
));
|
||||
}
|
||||
let components = value
|
||||
.extract()
|
||||
.map_err(|_| PyValueError::new_err("namespace must be a sequence of string components"))?;
|
||||
NamespacePath::new(components).map_err(PyValueError::new_err)
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl AuthzObject {
|
||||
#[new]
|
||||
fn new(value: &str) -> PyResult<Self> {
|
||||
Ok(Self {
|
||||
inner: value.parse().infer_error()?,
|
||||
})
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn system() -> Self {
|
||||
Self {
|
||||
inner: Object::System,
|
||||
}
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn database(name: &str) -> PyResult<Self> {
|
||||
Self::validated(Object::Database(DatabaseObject::new(name.into())))
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn namespace(database: &str, namespace: &Bound<'_, PyAny>) -> PyResult<Self> {
|
||||
Self::validated(Object::Namespace(NamespaceObject {
|
||||
database: database.into(),
|
||||
namespace: namespace_path(namespace)?,
|
||||
}))
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn table(database: &str, namespace: &Bound<'_, PyAny>, name: &str) -> PyResult<Self> {
|
||||
Self::validated(Object::Table(TableObject {
|
||||
database: database.into(),
|
||||
namespace: namespace_path(namespace)?,
|
||||
table: name.into(),
|
||||
}))
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn view(database: &str, namespace: &Bound<'_, PyAny>, name: &str) -> PyResult<Self> {
|
||||
Self::validated(Object::View(ViewObject {
|
||||
database: database.into(),
|
||||
namespace: namespace_path(namespace)?,
|
||||
view: name.into(),
|
||||
}))
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn secret(database: &str, namespace: &Bound<'_, PyAny>, name: &str) -> PyResult<Self> {
|
||||
Self::validated(Object::Secret(SecretObject {
|
||||
database: database.into(),
|
||||
namespace: namespace_path(namespace)?,
|
||||
secret: name.into(),
|
||||
}))
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn function(database: &str, namespace: &Bound<'_, PyAny>, name: &str) -> PyResult<Self> {
|
||||
Self::validated(Object::Function(FunctionObject {
|
||||
database: database.into(),
|
||||
namespace: namespace_path(namespace)?,
|
||||
function: name.into(),
|
||||
}))
|
||||
}
|
||||
|
||||
fn __str__(&self) -> String {
|
||||
self.inner.to_string()
|
||||
}
|
||||
fn __repr__(&self) -> String {
|
||||
format!("{:?}", self.inner)
|
||||
}
|
||||
}
|
||||
@@ -21,6 +21,13 @@ pub struct Catalog {
|
||||
|
||||
#[pymethods]
|
||||
impl Catalog {
|
||||
#[getter]
|
||||
fn authz(&self) -> PyResult<crate::authz::AuthorizationClient> {
|
||||
Ok(crate::authz::AuthorizationClient {
|
||||
inner: self.inner.authz().infer_error()?,
|
||||
})
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn uri(&self) -> &str {
|
||||
self.inner.uri()
|
||||
|
||||
@@ -1220,6 +1220,7 @@ pub struct PyClientConfig {
|
||||
user_agent: String,
|
||||
retry_config: Option<PyClientRetryConfig>,
|
||||
timeout_config: Option<PyClientTimeoutConfig>,
|
||||
blob_request_concurrency: Option<usize>,
|
||||
extra_headers: Option<HashMap<String, String>>,
|
||||
id_delimiter: Option<String>,
|
||||
tls_config: Option<PyClientTlsConfig>,
|
||||
@@ -1305,6 +1306,7 @@ impl From<PyClientConfig> for lancedb::remote::ClientConfig {
|
||||
user_agent: value.user_agent,
|
||||
retry_config: value.retry_config.map(Into::into).unwrap_or_default(),
|
||||
timeout_config: value.timeout_config.map(Into::into).unwrap_or_default(),
|
||||
blob_request_concurrency: value.blob_request_concurrency,
|
||||
extra_headers: value.extra_headers.unwrap_or_default(),
|
||||
id_delimiter: value.id_delimiter,
|
||||
tls_config: value.tls_config.map(Into::into),
|
||||
|
||||
@@ -22,6 +22,7 @@ use table::{
|
||||
};
|
||||
|
||||
pub mod arrow;
|
||||
pub mod authz;
|
||||
pub mod catalog;
|
||||
pub mod connection;
|
||||
pub mod error;
|
||||
@@ -70,6 +71,10 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<listing::FunctionListing>()?;
|
||||
m.add_class::<listing::JobListing>()?;
|
||||
|
||||
m.add_class::<authz::AuthorizationClient>()?;
|
||||
m.add_class::<authz::AuthzSubject>()?;
|
||||
m.add_class::<authz::AuthzObject>()?;
|
||||
m.add_class::<authz::AuthzPrivilege>()?;
|
||||
m.add_function(wrap_pyfunction!(catalog::connect_catalog, m)?)?;
|
||||
m.add_class::<Session>()?;
|
||||
m.add_class::<Table>()?;
|
||||
|
||||
+58
-15
@@ -697,6 +697,14 @@ impl PyBlobFile {
|
||||
}
|
||||
}
|
||||
|
||||
fn blob_file_error(operation: &str, error: lance_core::Error) -> pyo3::PyErr {
|
||||
let message = format!("blob {operation} failed: {error}");
|
||||
match error {
|
||||
lance_core::Error::InvalidInput { .. } => PyValueError::new_err(message),
|
||||
_ => PyRuntimeError::new_err(message),
|
||||
}
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyBlobFile {
|
||||
fn read_bytes(self_: PyRef<'_, Self>) -> PyResult<Py<PyBytes>> {
|
||||
@@ -705,7 +713,7 @@ impl PyBlobFile {
|
||||
let py = self_.py();
|
||||
let bytes = py
|
||||
.detach(move || block_on(async move { inner.read().await }))
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
|
||||
.map_err(|e| blob_file_error("read", e))?;
|
||||
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
|
||||
}
|
||||
|
||||
@@ -713,10 +721,7 @@ impl PyBlobFile {
|
||||
self_.ensure_open()?;
|
||||
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}")))?;
|
||||
let bytes = inner.read().await.map_err(|e| blob_file_error("read", e))?;
|
||||
Python::attach(|py| Ok(PyBytes::new(py, bytes.as_ref()).unbind()))
|
||||
})
|
||||
}
|
||||
@@ -739,7 +744,7 @@ impl PyBlobFile {
|
||||
self_
|
||||
.py()
|
||||
.detach(move || block_on(async move { inner.close().await }))
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob close failed: {e}")))
|
||||
.map_err(|e| blob_file_error("close", e))
|
||||
}
|
||||
|
||||
fn is_closed(self_: PyRef<'_, Self>) -> bool {
|
||||
@@ -752,7 +757,7 @@ impl PyBlobFile {
|
||||
self_
|
||||
.py()
|
||||
.detach(move || block_on(async move { inner.seek(position).await }))
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob seek failed: {e}")))
|
||||
.map_err(|e| blob_file_error("seek", e))
|
||||
}
|
||||
|
||||
fn tell(self_: PyRef<'_, Self>) -> PyResult<u64> {
|
||||
@@ -761,7 +766,7 @@ impl PyBlobFile {
|
||||
self_
|
||||
.py()
|
||||
.detach(move || block_on(async move { inner.tell().await }))
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob tell failed: {e}")))
|
||||
.map_err(|e| blob_file_error("tell", e))
|
||||
}
|
||||
|
||||
fn size(self_: PyRef<'_, Self>) -> u64 {
|
||||
@@ -778,17 +783,39 @@ impl PyBlobFile {
|
||||
let py = self_.py();
|
||||
let bytes = py
|
||||
.detach(move || block_on(async move { inner.read_range(offset..end).await }))
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read_range failed: {e}")))?;
|
||||
.map_err(|e| blob_file_error("read_range", e))?;
|
||||
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
|
||||
}
|
||||
|
||||
/// Read blob-local byte ranges without moving the cursor.
|
||||
fn read_ranges(self_: PyRef<'_, Self>, ranges: Vec<(u64, u64)>) -> PyResult<Vec<Py<PyBytes>>> {
|
||||
let ranges = ranges
|
||||
.into_iter()
|
||||
.map(|(offset, length)| {
|
||||
offset
|
||||
.checked_add(length)
|
||||
.map(|end| offset..end)
|
||||
.ok_or_else(|| PyValueError::new_err("offset + length overflowed"))
|
||||
})
|
||||
.collect::<PyResult<Vec<_>>>()?;
|
||||
let inner = self_.inner.clone();
|
||||
let py = self_.py();
|
||||
let bytes = py
|
||||
.detach(move || block_on(async move { inner.read_ranges(&ranges).await }))
|
||||
.map_err(|e| blob_file_error("read_ranges", e))?;
|
||||
Ok(bytes
|
||||
.iter()
|
||||
.map(|bytes| PyBytes::new(py, bytes.as_ref()).unbind())
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn read_up_to(self_: PyRef<'_, Self>, length: usize) -> PyResult<Py<PyBytes>> {
|
||||
self_.ensure_open()?;
|
||||
let inner = self_.inner.clone();
|
||||
let py = self_.py();
|
||||
let bytes = py
|
||||
.detach(move || block_on(async move { inner.read_up_to(length).await }))
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read_up_to failed: {e}")))?;
|
||||
.map_err(|e| blob_file_error("read_up_to", e))?;
|
||||
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
|
||||
}
|
||||
}
|
||||
@@ -1462,14 +1489,30 @@ impl Table {
|
||||
pub fn fetch_blob_ranges(
|
||||
self_: PyRef<'_, Self>,
|
||||
column: String,
|
||||
requests: Vec<(u64, u64, u64)>,
|
||||
requests: Vec<Vec<i128>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
let requests = requests
|
||||
.into_iter()
|
||||
.map(|values| {
|
||||
let [row_id, offset, length] = <[i128; 3]>::try_from(values).map_err(|_| {
|
||||
PyValueError::new_err("each blob range request must have three values")
|
||||
})?;
|
||||
let to_u64 = |value, name| {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
PyValueError::new_err(format!(
|
||||
"blob range {name} must be a non-negative 64-bit integer"
|
||||
))
|
||||
})
|
||||
};
|
||||
Ok(BlobRangeRequest::new(
|
||||
to_u64(row_id, "row id")?,
|
||||
to_u64(offset, "offset")?,
|
||||
to_u64(length, "length")?,
|
||||
))
|
||||
})
|
||||
.collect::<PyResult<Vec<_>>>()?;
|
||||
future_into_py(self_.py(), async move {
|
||||
let requests = requests
|
||||
.into_iter()
|
||||
.map(|(row_id, offset, length)| BlobRangeRequest::new(row_id, offset, length))
|
||||
.collect::<Vec<_>>();
|
||||
let blobs: LargeBinaryArray = inner
|
||||
.fetch_blob_ranges(column, requests)
|
||||
.await
|
||||
|
||||
Vendored
+1
-1
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lance-arrow"
|
||||
version = "14.0.0-beta.15"
|
||||
version = "14.0.0-beta.18"
|
||||
edition = "2024"
|
||||
authors = ["Lance Devs <dev@lance.org>"]
|
||||
license = "Apache-2.0"
|
||||
|
||||
+2
-2
@@ -1,5 +1,5 @@
|
||||
This directory contains `rust/lance-arrow` from Lance tag `v14.0.0-beta.15`
|
||||
(commit `765641dc0d7100f032fdebc00d73eef49711c4d4`), licensed under Apache-2.0.
|
||||
This directory contains `rust/lance-arrow` from Lance tag `v14.0.0-beta.18`
|
||||
(commit `f80bf1d1e59f40d249bce2539b78e41d93dee9ec`), licensed under Apache-2.0.
|
||||
|
||||
LanceDB carries a local fix in `src/lib.rs` for sliced child validity bitmaps
|
||||
until a Lance release includes the fix. The patch is selected in the root
|
||||
|
||||
Reference in New Issue
Block a user