Merge branch 'merge-insert-source-stats' into merge-insert-source-stats-bindings

This commit is contained in:
Will Jones
2026-10-16 21:18:09 +00:00
312 changed files with 21498 additions and 7261 deletions
+1 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "lancedb-python"
version = "0.41.0-beta.3"
version = "0.41.0-beta.7"
publish = false
edition.workspace = true
description = "Python bindings for LanceDB"
+2 -1
View File
@@ -63,7 +63,7 @@ tests = [
"polars>=0.19, <=1.32.3",
"pyarrow<25",
"pyarrow-stubs>=16.0",
"pylance==14.0.0b10",
"pylance==14.0.0b12",
"requests>=2.31.0",
"datafusion>=54,<55",
"opentelemetry-sdk>=1.30.0",
@@ -73,6 +73,7 @@ dev = [
"pre-commit>=3.5.0",
"pyright>=1.1.350",
'typing-extensions>=4.0.0; python_version < "3.11"',
"maturin>=1.10",
]
docs = ["mkdocs", "mkdocs-jupyter", "mkdocs-material", "mkdocstrings-python"]
clip = ["torch", "pillow>=12.1.1", "open-clip-torch"]
+2 -2
View File
@@ -15,7 +15,7 @@ from ._lancedb import FtsToken
from ._lancedb import FunctionErrorFragment as FunctionErrorFragment
from ._lancedb import FunctionErrorRecord as FunctionErrorRecord
from ._lancedb import FunctionErrors as FunctionErrors
from ._lancedb import LsmWriteSpec
from ._lancedb import MemWalWriteSpec
from ._lancedb import tokenize as _tokenize
from .common import URI, sanitize_uri
from urllib.parse import urlparse
@@ -629,7 +629,7 @@ __all__ = [
"SqlQuery",
"LanceDBConnection",
"LanceNamespaceDBConnection",
"LsmWriteSpec",
"MemWalWriteSpec",
"RemoteDBConnection",
"Session",
"Table",
+250 -53
View File
@@ -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,10 +136,12 @@ 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 source in columns:
if output in columns or any(
column == source or column.startswith(f"{source}.") for column in columns
):
return True
return False
@@ -119,17 +150,69 @@ 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 {}
columns = set(blob_columns)
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}
return {
output: source
for output, source in _iter_projection_pairs(projection)
if source in columns
}
sources: dict[str, str] = {}
entries: dict[str, str] = {}
for output, source in _iter_projection_pairs(projection):
for column in blob_columns:
if column == source:
path = output
elif column.startswith(f"{source}."):
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
def v2_projection_needs_row_id(
@@ -179,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 output_name not in tbl.column_names:
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
@@ -192,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 output_name not in tbl.column_names:
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")
@@ -224,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
@@ -287,9 +402,37 @@ def _iter_projection_pairs(
yield name, source if source is not None else expr.to_sql()
def _set_blob_column(tbl: pa.Table, output_name: str, blobs: pa.Array) -> pa.Table:
index = tbl.schema.get_field_index(output_name)
return tbl.set_column(index, pa.field(output_name, blobs.type), [blobs])
def _has_column_path(schema: pa.Schema, path: str) -> bool:
top_name, rest = _split_column_path(schema, path)
index = schema.get_field_index(top_name)
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:
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])
leaf_name = rest[-1]
if isinstance(blobs, pa.ChunkedArray):
blobs = blobs.combine_chunks()
def replace_blob(children: list, child_fields: list) -> None:
for i, field in enumerate(child_fields):
if field.name == leaf_name:
children[i] = blobs
child_fields[i] = pa.field(
leaf_name, blobs.type, nullable=field.nullable
)
return
raise KeyError(f"Blob field '{leaf_name}' not found in '{output_name}'")
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:
@@ -297,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:
@@ -314,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)
@@ -381,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]]:
+89 -22
View File
@@ -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:
@@ -537,17 +538,17 @@ class Table:
async def initial_storage_options(self) -> Optional[Dict[str, str]]: ...
async def latest_storage_options(self) -> Optional[Dict[str, str]]: ...
async def set_unenforced_primary_key(self, columns: List[str]) -> None: ...
async def set_lsm_write_spec(self, spec: LsmWriteSpec) -> None: ...
async def unset_lsm_write_spec(self) -> None: ...
async def get_lsm_write_spec(self) -> Optional[LsmWriteSpec]: ...
async def enable_memwal(self, spec: MemWalWriteSpec) -> None: ...
async def disable_memwal(self) -> None: ...
async def get_memwal_spec(self) -> Optional[MemWalWriteSpec]: ...
# WAL-PK-FUSION: delete both.
def hybrid_pk_fusion_learned(self) -> bool: ...
def note_hybrid_pk_fusion(self) -> None: ...
async def checkpoint_lsm(self) -> None: ...
async def flush_lsm(self) -> None: ...
async def compact_lsm(self) -> None: ...
async def get_lsm_stats(self, include_generation_rows: bool) -> Optional[dict]: ...
async def close_lsm_writers(self) -> None: ...
async def checkpoint_memwal(self) -> None: ...
async def flush_memwal(self) -> None: ...
async def compact_memwal(self) -> None: ...
async def get_memwal_stats(self, include_sstable_rows: bool) -> Optional[dict]: ...
async def close_memwal_writers(self) -> None: ...
@property
def tags(self) -> Tags: ...
@property
@@ -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]
@@ -661,7 +662,7 @@ class Query:
def fast_search(self): ...
def with_row_id(self): ...
def postfilter(self): ...
def use_lsm(self, enable: bool): ...
def use_memwal(self, enable: bool): ...
def nearest_to(self, query_vec: pa.Array) -> VectorQuery: ...
def nearest_to_text(self, query: dict) -> FTSQuery: ...
def order_by(self, ordering: Optional[List[ColumnOrdering]]): ...
@@ -678,7 +679,7 @@ class Query:
class TakeQuery:
def select(self, columns: List[str]): ...
def with_row_id(self): ...
def use_lsm(self, enable: bool): ...
def use_memwal(self, enable: bool): ...
async def output_schema(self) -> pa.Schema: ...
async def execute(self) -> RecordBatchStream: ...
async def explain_plan(self, verbose: Optional[bool]) -> str: ...
@@ -697,7 +698,7 @@ class FTSQuery:
def fast_search(self): ...
def with_row_id(self): ...
def postfilter(self): ...
def use_lsm(self, enable: bool): ...
def use_memwal(self, enable: bool): ...
def get_query(self) -> str: ...
def add_query_vector(self, query_vec: pa.Array) -> None: ...
def nearest_to(self, query_vec: pa.Array) -> HybridQuery: ...
@@ -725,9 +726,10 @@ class VectorQuery:
def column(self, column: str): ...
def distance_type(self, distance_type: str): ...
def postfilter(self): ...
def use_lsm(self, enable: bool): ...
def use_memwal(self, enable: bool): ...
def refine_factor(self, refine_factor: int): ...
def nprobes(self, nprobes: int): ...
def search_effort(self, search_effort: Optional[float]): ...
def minimum_nprobes(self, minimum_nprobes: int): ...
def maximum_nprobes(self, maximum_nprobes: int): ...
def bypass_vector_index(self): ...
@@ -749,10 +751,11 @@ class HybridQuery:
def fast_search(self): ...
def with_row_id(self): ...
def postfilter(self): ...
def use_lsm(self, enable: bool): ...
def use_memwal(self, enable: bool): ...
def distance_type(self, distance_type: str): ...
def refine_factor(self, refine_factor: int): ...
def nprobes(self, nprobes: int): ...
def search_effort(self, search_effort: Optional[float]): ...
def minimum_nprobes(self, minimum_nprobes: int): ...
def maximum_nprobes(self, maximum_nprobes: int): ...
def bypass_vector_index(self): ...
@@ -776,12 +779,13 @@ class PyQueryRequest:
select_source_columns: Optional[Dict[str, str]]
fast_search: Optional[bool]
with_row_id: Optional[bool]
use_lsm: Optional[bool]
use_memwal: Optional[bool]
column: Optional[str]
query_vector: Optional[List[pa.Array]]
nprobes: Optional[int]
minimum_nprobes: Optional[int]
maximum_nprobes: Optional[int]
search_effort: Optional[float]
lower_bound: Optional[float]
upper_bound: Optional[float]
ef: Optional[int]
@@ -836,22 +840,26 @@ class FtsToken:
text: str
position: int
class LsmWriteSpec:
class MemWalWriteSpec:
"""Specification selecting Lance's MemWAL LSM-style write path for
`merge_insert`."""
@staticmethod
def bucket(column: str, num_buckets: int) -> "LsmWriteSpec": ...
def bucket(column: str, num_buckets: int) -> "MemWalWriteSpec": ...
@staticmethod
def identity(column: str) -> "LsmWriteSpec": ...
def identity(column: str) -> "MemWalWriteSpec": ...
@staticmethod
def unsharded() -> "LsmWriteSpec": ...
def with_maintained_indexes(self, indexes: Optional[List[str]]) -> "LsmWriteSpec":
def unsharded() -> "MemWalWriteSpec": ...
def with_maintained_indexes(
self, indexes: Optional[List[str]]
) -> "MemWalWriteSpec":
"""Set which indexes the MemWAL keeps up to date. None resolves every
index on the table at install, failing if one cannot be maintained;
a list is verbatim, empty means none."""
...
def with_writer_config_defaults(self, defaults: Dict[str, str]) -> "LsmWriteSpec":
def with_writer_config_defaults(
self, defaults: Dict[str, str]
) -> "MemWalWriteSpec":
"""Return a copy of this spec recording the given default
`ShardWriter` configuration in the MemWAL index."""
...
@@ -953,6 +961,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(
@@ -999,3 +1009,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: ...
+4 -4
View File
@@ -76,7 +76,7 @@ _WAL_ROW_ID_REFUSAL = "does not support with_row_id"
WAL_ROW_ID_UNSUPPORTED = (
"hybrid search on a MemWAL table cannot return _rowid: the fresh tier has "
"no stable row id, and the ids the fusion joins on are synthesized from "
"the primary key. Set use_lsm(False) to read the base table only (results "
"the primary key. Set use_memwal(False) to read the base table only (results "
"will exclude un-compacted MemWAL data)"
)
@@ -151,13 +151,13 @@ def _stamp_surrogate_row_ids(
return stamped
def pk_fusion_learned(table, use_lsm: Optional[bool]) -> bool:
def pk_fusion_learned(table, use_memwal: Optional[bool]) -> bool:
"""Whether a previous hybrid query on `table` was refused `_rowid`, so this
one should go straight to the primary key.
Learned from a refusal, never probed, so asking is free.
"""
if table is None or use_lsm is False:
if table is None or use_memwal is False:
return False
# Only an unambiguous True forks to the key: a table type that answers the
# question loosely keeps the `_rowid` path.
@@ -192,7 +192,7 @@ class PkFusion:
raise NotImplementedError(
"hybrid search cannot project a blob column on a MemWAL table: "
"fetching blobs needs a real _rowid, and the fresh tier has no "
"stable one. Set use_lsm(False) to read the base table only "
"stable one. Set use_memwal(False) to read the base table only "
"(results will exclude un-compacted MemWAL data)"
)
return cls(_pk_columns(schema))
+588
View File
@@ -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))
+12
View File
@@ -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."""
+7 -7
View File
@@ -37,7 +37,7 @@ class LanceMergeInsertBuilder(object):
self._when_not_matched_by_source_condition_expr = None
self._timeout = None
self._use_index = True
self._use_lsm = None
self._use_memwal = None
self._validate_single_shard = None
self._source_collect_threshold_bytes = None
@@ -116,22 +116,22 @@ class LanceMergeInsertBuilder(object):
self._use_index = use_index
return self
def use_lsm(self, enable: bool) -> LanceMergeInsertBuilder:
def use_memwal(self, enable: bool) -> LanceMergeInsertBuilder:
"""
Control MemWAL routing for this merge.
By default (unset), a `merge_insert` on a table with an LSM write spec is
By default (unset), a `merge_insert` on a table with a MemWAL write spec is
routed through Lance's MemWAL shard writer, and a table without one uses
the standard path.
Parameters
----------
enable: bool
``True`` forces MemWAL routing and errors if the table has no LSM
``True`` forces MemWAL routing and errors if the table has no MemWAL
write spec. ``False`` forces the standard write path even when a spec
is set.
"""
self._use_lsm = enable
self._use_memwal = enable
return self
def validate_single_shard(
@@ -140,13 +140,13 @@ class LanceMergeInsertBuilder(object):
"""
Controls how an LSM merge checks that its input targets a single shard.
When a table has an LSM write spec, every row in a `merge_insert` call
When a table has a MemWAL write spec, every row in a `merge_insert` call
must route to the same shard. When `True` (the default), every row is
inspected to verify this. When `False`, only the first row is inspected
and the shard it routes to is used for the whole input — a faster path
for callers that have already pre-sharded their input.
Has no effect on tables without an LSM write spec.
Has no effect on tables without a MemWAL write spec.
Parameters
----------
+5
View File
@@ -105,6 +105,11 @@ def _query_to_namespace_request(
if not isinstance(query, Query):
raise TypeError(f"Expected Query, got {type(query)}")
if query.search_effort is not None:
raise ValueError(
"namespace QueryTable does not support search_effort; "
"execute the query locally"
)
vector_request = None
if query.vector is not None:
+1 -1
View File
@@ -48,7 +48,7 @@ class PermutationBuilder:
By default, the permutation builder will create a single split that contains all
rows in the same order as the base table.
Tables with an LSM write spec are rejected: unflushed rows have no row id.
Tables with a MemWAL write spec are rejected: unflushed rows have no row id.
"""
self._async = async_permutation_builder(table)
File diff suppressed because it is too large Load Diff
+5
View File
@@ -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
+1
View File
@@ -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),
+35 -28
View File
@@ -31,7 +31,7 @@ from lancedb._lancedb import (
DropColumnsResult,
FtsToken,
IndexConfig,
LsmWriteSpec,
MemWalWriteSpec,
MergeResult,
UpdateResult,
)
@@ -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]):
@@ -1031,20 +1034,20 @@ class RemoteTable(Table):
return LOOP.run(self._table.drop_columns(columns))
def set_unenforced_primary_key(self, columns: Union[str, Iterable[str]]) -> None:
"""Not supported for remote tables."""
"""Set the unenforced primary key for this table to a single column."""
return LOOP.run(self._table.set_unenforced_primary_key(columns))
def set_lsm_write_spec(self, spec: "LsmWriteSpec") -> None:
"""Install an LsmWriteSpec."""
return LOOP.run(self._table.set_lsm_write_spec(spec))
def enable_memwal(self, spec: "MemWalWriteSpec") -> None:
"""Install an MemWalWriteSpec."""
return LOOP.run(self._table.enable_memwal(spec))
def unset_lsm_write_spec(self) -> None:
"""Remove the LsmWriteSpec."""
return LOOP.run(self._table.unset_lsm_write_spec())
def disable_memwal(self) -> None:
"""Remove the MemWalWriteSpec."""
return LOOP.run(self._table.disable_memwal())
def get_lsm_write_spec(self) -> Optional["LsmWriteSpec"]:
"""Read the installed LsmWriteSpec, or ``None``."""
return LOOP.run(self._table.get_lsm_write_spec())
def get_memwal_spec(self) -> Optional["MemWalWriteSpec"]:
"""Read the installed MemWalWriteSpec, or ``None``."""
return LOOP.run(self._table.get_memwal_spec())
# WAL-PK-FUSION: delete both hooks.
def _hybrid_pk_fusion_learned(self) -> bool:
@@ -1054,31 +1057,31 @@ class RemoteTable(Table):
def _note_hybrid_pk_fusion(self) -> None:
self._table._note_hybrid_pk_fusion()
def checkpoint_lsm(self) -> None:
def checkpoint_memwal(self) -> None:
"""Synchronous version of
[`AsyncTable.checkpoint_lsm`][lancedb.AsyncTable.checkpoint_lsm]."""
return LOOP.run(self._table.checkpoint_lsm())
[`AsyncTable.checkpoint_memwal`][lancedb.AsyncTable.checkpoint_memwal]."""
return LOOP.run(self._table.checkpoint_memwal())
def flush_lsm(self) -> None:
def flush_memwal(self) -> None:
"""Synchronous version of
[`AsyncTable.flush_lsm`][lancedb.AsyncTable.flush_lsm]."""
return LOOP.run(self._table.flush_lsm())
[`AsyncTable.flush_memwal`][lancedb.AsyncTable.flush_memwal]."""
return LOOP.run(self._table.flush_memwal())
def compact_lsm(self) -> None:
def compact_memwal(self) -> None:
"""Synchronous version of
[`AsyncTable.compact_lsm`][lancedb.AsyncTable.compact_lsm]."""
return LOOP.run(self._table.compact_lsm())
[`AsyncTable.compact_memwal`][lancedb.AsyncTable.compact_memwal]."""
return LOOP.run(self._table.compact_memwal())
def get_lsm_stats(self, *, include_generation_rows: bool = False) -> Optional[dict]:
def get_memwal_stats(self, *, include_sstable_rows: bool = False) -> Optional[dict]:
"""Synchronous version of
[`AsyncTable.get_lsm_stats`][lancedb.AsyncTable.get_lsm_stats]."""
[`AsyncTable.get_memwal_stats`][lancedb.AsyncTable.get_memwal_stats]."""
return LOOP.run(
self._table.get_lsm_stats(include_generation_rows=include_generation_rows)
self._table.get_memwal_stats(include_sstable_rows=include_sstable_rows)
)
def close_lsm_writers(self) -> None:
def close_memwal_writers(self) -> None:
"""No-op for remote tables (no local shard writers)."""
return LOOP.run(self._table.close_lsm_writers())
return LOOP.run(self._table.close_memwal_writers())
def drop_index(self, index_name: str):
return LOOP.run(self._table.drop_index(index_name))
@@ -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))
+51 -4
View File
@@ -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:
+269 -118
View File
@@ -223,7 +223,7 @@ if TYPE_CHECKING:
DeleteResult,
DropColumnsResult,
FtsToken,
LsmWriteSpec,
MemWalWriteSpec,
MergeResult,
UpdateResult,
)
@@ -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".
@@ -4440,6 +4562,7 @@ class LanceTable(Table):
and not self._route_pushdown_to_rust
and self.current_branch() is None
and query.take_offsets is None
and query.search_effort is None
):
from lancedb.namespace import _execute_server_side_query
@@ -4704,47 +4827,47 @@ class LanceTable(Table):
[`AsyncTable.set_unenforced_primary_key`][lancedb.AsyncTable.set_unenforced_primary_key]."""
return LOOP.run(self._table.set_unenforced_primary_key(columns))
def set_lsm_write_spec(self, spec: "LsmWriteSpec") -> None:
"""Install an LsmWriteSpec. See
[`AsyncTable.set_lsm_write_spec`][lancedb.AsyncTable.set_lsm_write_spec]."""
return LOOP.run(self._table.set_lsm_write_spec(spec))
def enable_memwal(self, spec: "MemWalWriteSpec") -> None:
"""Install an MemWalWriteSpec. See
[`AsyncTable.enable_memwal`][lancedb.AsyncTable.enable_memwal]."""
return LOOP.run(self._table.enable_memwal(spec))
def unset_lsm_write_spec(self) -> None:
"""Remove the LsmWriteSpec. See
[`AsyncTable.unset_lsm_write_spec`][lancedb.AsyncTable.unset_lsm_write_spec]."""
return LOOP.run(self._table.unset_lsm_write_spec())
def disable_memwal(self) -> None:
"""Remove the MemWalWriteSpec. See
[`AsyncTable.disable_memwal`][lancedb.AsyncTable.disable_memwal]."""
return LOOP.run(self._table.disable_memwal())
def get_lsm_write_spec(self) -> Optional["LsmWriteSpec"]:
"""Read the installed LsmWriteSpec, or ``None``. See
[`AsyncTable.get_lsm_write_spec`][lancedb.AsyncTable.get_lsm_write_spec]."""
return LOOP.run(self._table.get_lsm_write_spec())
def get_memwal_spec(self) -> Optional["MemWalWriteSpec"]:
"""Read the installed MemWalWriteSpec, or ``None``. See
[`AsyncTable.get_memwal_spec`][lancedb.AsyncTable.get_memwal_spec]."""
return LOOP.run(self._table.get_memwal_spec())
def checkpoint_lsm(self) -> None:
def checkpoint_memwal(self) -> None:
"""Synchronous version of
[`AsyncTable.checkpoint_lsm`][lancedb.AsyncTable.checkpoint_lsm]."""
return LOOP.run(self._table.checkpoint_lsm())
[`AsyncTable.checkpoint_memwal`][lancedb.AsyncTable.checkpoint_memwal]."""
return LOOP.run(self._table.checkpoint_memwal())
def flush_lsm(self) -> None:
def flush_memwal(self) -> None:
"""Synchronous version of
[`AsyncTable.flush_lsm`][lancedb.AsyncTable.flush_lsm]."""
return LOOP.run(self._table.flush_lsm())
[`AsyncTable.flush_memwal`][lancedb.AsyncTable.flush_memwal]."""
return LOOP.run(self._table.flush_memwal())
def compact_lsm(self) -> None:
def compact_memwal(self) -> None:
"""Synchronous version of
[`AsyncTable.compact_lsm`][lancedb.AsyncTable.compact_lsm]."""
return LOOP.run(self._table.compact_lsm())
[`AsyncTable.compact_memwal`][lancedb.AsyncTable.compact_memwal]."""
return LOOP.run(self._table.compact_memwal())
def get_lsm_stats(self, *, include_generation_rows: bool = False) -> Optional[dict]:
def get_memwal_stats(self, *, include_sstable_rows: bool = False) -> Optional[dict]:
"""Synchronous version of
[`AsyncTable.get_lsm_stats`][lancedb.AsyncTable.get_lsm_stats]."""
[`AsyncTable.get_memwal_stats`][lancedb.AsyncTable.get_memwal_stats]."""
return LOOP.run(
self._table.get_lsm_stats(include_generation_rows=include_generation_rows)
self._table.get_memwal_stats(include_sstable_rows=include_sstable_rows)
)
def close_lsm_writers(self) -> None:
def close_memwal_writers(self) -> None:
"""Close cached MemWAL shard writers. See
[`AsyncTable.close_lsm_writers`][lancedb.AsyncTable.close_lsm_writers]."""
return LOOP.run(self._table.close_lsm_writers())
[`AsyncTable.close_memwal_writers`][lancedb.AsyncTable.close_memwal_writers]."""
return LOOP.run(self._table.close_memwal_writers())
def uses_v2_manifest_paths(self) -> bool:
"""
@@ -5402,18 +5525,18 @@ class AsyncTable:
columns = list(columns)
await self._inner.set_unenforced_primary_key(columns)
async def set_lsm_write_spec(self, spec: "LsmWriteSpec") -> None:
"""Install an LsmWriteSpec on this table.
async def enable_memwal(self, spec: "MemWalWriteSpec") -> None:
"""Install an MemWalWriteSpec on this table.
The spec selects Lance's MemWAL LSM-style write path for future
`merge_insert` calls. ``LsmWriteSpec`` chooses one of three sharding
`merge_insert` calls. ``MemWalWriteSpec`` chooses one of three sharding
strategies:
- ``LsmWriteSpec.bucket(column, num_buckets)`` — hash-bucket writes by
- ``MemWalWriteSpec.bucket(column, num_buckets)`` — hash-bucket writes by
the single-column unenforced primary key.
- ``LsmWriteSpec.identity(column)`` — shard by the raw value of a
- ``MemWalWriteSpec.identity(column)`` — shard by the raw value of a
scalar column.
- ``LsmWriteSpec.unsharded()`` — route every write to a single shard.
- ``MemWalWriteSpec.unsharded()`` — route every write to a single shard.
All variants require the table to have an unenforced primary key set
via [`set_unenforced_primary_key`]; bucket sharding additionally
@@ -5428,30 +5551,30 @@ class AsyncTable:
Parameters
----------
spec : LsmWriteSpec
spec : MemWalWriteSpec
The sharding spec to install.
Examples
--------
>>> from lancedb import LsmWriteSpec
>>> from lancedb import MemWalWriteSpec
>>> # table.set_unenforced_primary_key("id")
>>> # table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 16))
>>> # table.enable_memwal(MemWalWriteSpec.bucket("id", 16))
"""
await self._inner.set_lsm_write_spec(spec)
await self._inner.enable_memwal(spec)
async def unset_lsm_write_spec(self) -> None:
"""Remove the LsmWriteSpec from this table.
async def disable_memwal(self) -> None:
"""Remove the MemWalWriteSpec from this table.
Reverts to the standard `merge_insert` write path. Errors if no spec
is currently set.
"""
await self._inner.unset_lsm_write_spec()
await self._inner.disable_memwal()
async def get_lsm_write_spec(self) -> Optional["LsmWriteSpec"]:
"""Read the LsmWriteSpec currently installed on this table.
async def get_memwal_spec(self) -> Optional["MemWalWriteSpec"]:
"""Read the MemWalWriteSpec currently installed on this table.
Returns ``None`` when the LSM write path is not enabled at all — no
spec has been set, or one was removed with `unset_lsm_write_spec`.
Returns ``None`` when the MemWAL write path is not enabled at all — no
spec has been set, or one was removed with `disable_memwal`.
That is a different answer from a spec whose ``maintained_indexes``
is ``None``, which is an installed spec selecting indexes
automatically.
@@ -5461,7 +5584,7 @@ class AsyncTable:
or gains later, ``[]`` maintains none, and a non-empty list maintains
exactly those. All three round-trip.
"""
return await self._inner.get_lsm_write_spec()
return await self._inner.get_memwal_spec()
# WAL-PK-FUSION: delete both hooks.
def _hybrid_pk_fusion_learned(self) -> bool:
@@ -5472,19 +5595,19 @@ class AsyncTable:
"""See [`Table._note_hybrid_pk_fusion`][lancedb.table.Table]."""
self._inner.note_hybrid_pk_fusion()
async def checkpoint_lsm(self) -> None:
"""Converge this table's LSM write path into its base table.
async def checkpoint_memwal(self) -> None:
"""Converge this table's MemWAL write path into its base table.
One flush, sealing every memtable into L0, then compaction triggers
until every generation that existed at that moment has reached base.
The loop runs client-side, reading progress from ``get_lsm_stats``.
One flush, freezing every memtable into an SSTable, then compaction triggers
until every SSTable that existed at that moment has reached base.
The loop runs client-side, reading progress from ``get_memwal_stats``.
Best-effort: generations created *while* it runs are deliberately not
Best-effort: SSTables created *while* it runs are deliberately not
waited on, which is what lets it terminate on a table taking writes.
Idempotent and safe on a cadence.
There is no deadline, and the caller owns that. It returns when the
target generations are gone, raises on a terminal server fault, and
target SSTables are gone, raises on a terminal server fault, and
otherwise waits however long the server takes. A slow table and a
stuck one are the same picture from the client: the compactor pool is
shared across every table on the node, so a checkpoint queued behind
@@ -5492,56 +5615,56 @@ class AsyncTable:
``asyncio.wait_for`` for a wall-clock bound; abandoning it partway
costs nothing.
"""
await self._inner.checkpoint_lsm()
await self._inner.checkpoint_memwal()
async def flush_lsm(self) -> None:
"""Seal every bucket's active memtable into L0.
async def flush_memwal(self) -> None:
"""Freeze every table shard's active memtable into an SSTable.
Does not touch the base table — moving L0 into base is
`compact_lsm`. On a node that has not claimed this table, this claims
Does not touch the base table — compacting SSTables into base is
`compact_memwal`. On a node that has not claimed this table, this claims
it and replays its WAL log first.
"""
await self._inner.flush_lsm()
await self._inner.flush_memwal()
async def compact_lsm(self) -> None:
"""Trigger a background L0 to base compaction pass per bucket.
async def compact_memwal(self) -> None:
"""Trigger a background SSTable compaction pass per table shard.
Returns once the passes are dispatched, not once they finish: watch
``get_lsm_stats`` for progress, or use ``checkpoint_lsm`` to loop
until the current L0 has reached base.
``get_memwal_stats`` for progress, or use ``checkpoint_memwal`` to loop
until the current SSTables have reached base.
"""
await self._inner.compact_lsm()
await self._inner.compact_memwal()
async def get_lsm_stats(
self, *, include_generation_rows: bool = False
async def get_memwal_stats(
self, *, include_sstable_rows: bool = False
) -> Optional[dict]:
"""Read live per-bucket LSM state.
"""Read live per-table-shard MemWAL state.
Answers "how far behind is my fresh tier", "which bucket is hot", and
Answers "how far behind is my fresh tier", "which table shard is hot", and
"why is my fresh-tier vector search brute-force". Mutates no table
state, though on a node that has not claimed this table it claims it,
exactly as a read would.
Returns ``None`` only when the LSM write path is not enabled.
Returns ``None`` only when the MemWAL write path is not enabled.
Parameters
----------
include_generation_rows
Report a row count per L0 generation. Off by default: each count
opens an uncached Lance dataset, and ``checkpoint_lsm`` polls this
include_sstable_rows
Report a row count per SSTable. Off by default: each count
opens an uncached Lance dataset, and ``checkpoint_memwal`` polls this
needing only generation numbers.
"""
return await self._inner.get_lsm_stats(include_generation_rows)
return await self._inner.get_memwal_stats(include_sstable_rows)
async def close_lsm_writers(self) -> None:
async def close_memwal_writers(self) -> None:
"""Drain and close any cached MemWAL shard writers for this table.
When an LSM write spec is installed, `merge_insert` opens MemWAL shard
When a MemWAL write spec is installed, `merge_insert` opens MemWAL shard
writers and caches them for reuse across calls. This closes them,
flushing pending data; writers reopen lazily on the next
`merge_insert`. It is a no-op when no writers are cached.
"""
await self._inner.close_lsm_writers()
await self._inner.close_memwal_writers()
@property
def name(self) -> str:
@@ -5686,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,
@@ -5997,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".
@@ -6035,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
):
@@ -6409,8 +6553,8 @@ class AsyncTable:
take_query = self.take_offsets(query.take_offsets)
if query.columns:
take_query = take_query.select(query.columns)
if query.use_lsm is not None:
take_query = take_query.use_lsm(query.use_lsm)
if query.use_memwal is not None:
take_query = take_query.use_memwal(query.use_memwal)
if query.with_row_id:
take_query = take_query.with_row_id()
return take_query
@@ -6426,8 +6570,8 @@ class AsyncTable:
async_query = async_query.where(query.filter)
if query.fast_search:
async_query = async_query.fast_search()
if query.use_lsm is not None:
async_query = async_query.use_lsm(query.use_lsm)
if query.use_memwal is not None:
async_query = async_query.use_memwal(query.use_memwal)
if query.with_row_id:
async_query = async_query.with_row_id()
if query.order_by:
@@ -6441,6 +6585,8 @@ class AsyncTable:
async_query = async_query.distance_type(query.distance_type)
if query.nprobes is not None:
async_query = async_query.nprobes(query.nprobes)
if query.search_effort is not None:
async_query = async_query.search_effort(query.search_effort)
if query.minimum_nprobes is not None and query.maximum_nprobes is not None:
async_query = async_query.minimum_nprobes(
query.minimum_nprobes
@@ -6481,6 +6627,7 @@ class AsyncTable:
)
and not self._route_pushdown_to_rust
and query.take_offsets is None
and query.search_effort is None
):
from lancedb.namespace import _execute_server_side_query
@@ -6550,7 +6697,7 @@ class AsyncTable:
when_not_matched_by_source_condition_expr=merge._when_not_matched_by_source_condition_expr,
timeout=merge._timeout,
use_index=merge._use_index,
use_lsm=merge._use_lsm,
use_memwal=merge._use_memwal,
validate_single_shard=merge._validate_single_shard,
source_collect_threshold_bytes=merge._source_collect_threshold_bytes,
),
@@ -7076,7 +7223,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)
@@ -7085,12 +7234,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)
+13
View File
@@ -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. "
+821
View File
@@ -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
)
File diff suppressed because it is too large Load Diff
+42
View File
@@ -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 == []
@@ -1,8 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Tests for installing and clearing an LsmWriteSpec via
`Table.set_lsm_write_spec` / `Table.unset_lsm_write_spec`.
"""Tests for installing and clearing an MemWalWriteSpec via
`Table.enable_memwal` / `Table.disable_memwal`.
"""
from datetime import timedelta
@@ -10,7 +10,7 @@ from datetime import timedelta
import lancedb
import pyarrow as pa
import pytest
from lancedb._lancedb import LsmWriteSpec
from lancedb._lancedb import MemWalWriteSpec
from lancedb.index import BTree
SCHEMA = pa.schema(
@@ -21,9 +21,9 @@ SCHEMA = pa.schema(
)
def test_lsm_write_spec_module_metadata():
assert lancedb.LsmWriteSpec is LsmWriteSpec
assert LsmWriteSpec.__module__ == "lancedb._lancedb"
def test_memwal_spec_module_metadata():
assert lancedb.MemWalWriteSpec is MemWalWriteSpec
assert MemWalWriteSpec.__module__ == "lancedb._lancedb"
def _batch(ids, vs):
@@ -43,35 +43,49 @@ def _make_table(tmp_path):
return db, table
def test_set_lsm_write_spec_validates(tmp_path):
def test_enable_memwal_validates(tmp_path):
_db, table = _make_table(tmp_path)
table.set_unenforced_primary_key("id")
# Out-of-range num_buckets.
with pytest.raises(Exception, match="num_buckets"):
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 0))
table.enable_memwal(MemWalWriteSpec.bucket("id", 0))
with pytest.raises(Exception, match="num_buckets"):
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 1025))
table.enable_memwal(MemWalWriteSpec.bucket("id", 1025))
# Happy path, then a set over the installed spec is refused.
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 4))
table.enable_memwal(MemWalWriteSpec.bucket("id", 4))
with pytest.raises(Exception, match="already set"):
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 8))
table.enable_memwal(MemWalWriteSpec.bucket("id", 8))
def test_unset_lsm_write_spec(tmp_path):
def test_disable_memwal(tmp_path):
_db, table = _make_table(tmp_path)
table.set_unenforced_primary_key("id")
# unset errors when no spec is set.
with pytest.raises(Exception, match="no LSM write spec"):
table.unset_lsm_write_spec()
with pytest.raises(Exception, match="no MemWAL write spec"):
table.disable_memwal()
# Install a spec, then remove it; afterwards a fresh spec can be set.
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 4))
table.unset_lsm_write_spec()
table.enable_memwal(MemWalWriteSpec.bucket("id", 4))
table.disable_memwal()
# A second unset errors — there is no spec left to remove.
with pytest.raises(Exception, match="no LSM write spec"):
table.unset_lsm_write_spec()
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 8))
with pytest.raises(Exception, match="no MemWAL write spec"):
table.disable_memwal()
table.enable_memwal(MemWalWriteSpec.bucket("id", 8))
def test_enable_memwal_requires_primary_key(tmp_path):
_db, table = _make_table(tmp_path)
for spec in (
MemWalWriteSpec.unsharded(),
MemWalWriteSpec.bucket("id", 4),
MemWalWriteSpec.identity("id"),
):
with pytest.raises(Exception, match="primary key"):
table.enable_memwal(spec)
assert table.get_memwal_spec() is None
def test_set_unsharded_spec(tmp_path):
@@ -79,12 +93,12 @@ def test_set_unsharded_spec(tmp_path):
# Lance MemWAL still requires a primary key on the dataset; Unsharded
# just skips per-row hashing.
table.set_unenforced_primary_key("id")
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
table.unset_lsm_write_spec()
table.enable_memwal(MemWalWriteSpec.unsharded())
table.disable_memwal()
def test_lsm_write_spec_repr():
s = LsmWriteSpec.bucket("id", 4)
def test_memwal_spec_repr():
s = MemWalWriteSpec.bucket("id", 4)
assert s.spec_type == "bucket"
assert s.column == "id"
assert s.num_buckets == 4
@@ -95,20 +109,20 @@ def test_lsm_write_spec_repr():
assert "id" in repr(s)
assert "4" in repr(s)
u = LsmWriteSpec.unsharded()
u = MemWalWriteSpec.unsharded()
assert u.spec_type == "unsharded"
assert u.column is None
assert u.num_buckets is None
assert "unsharded" in repr(u)
def test_lsm_write_spec_with_maintained_indexes():
s = LsmWriteSpec.bucket("id", 4).with_maintained_indexes(["idx_a", "idx_b"])
def test_memwal_spec_with_maintained_indexes():
s = MemWalWriteSpec.bucket("id", 4).with_maintained_indexes(["idx_a", "idx_b"])
assert s.maintained_indexes == ["idx_a", "idx_b"]
@pytest.mark.asyncio
async def test_async_set_unset_lsm_write_spec(tmp_path):
async def test_async_set_disable_memwal(tmp_path):
db = await lancedb.connect_async(
tmp_path, read_consistency_interval=timedelta(seconds=0)
)
@@ -118,11 +132,11 @@ async def test_async_set_unset_lsm_write_spec(tmp_path):
)
await table.set_unenforced_primary_key("id")
await table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 4))
await table.unset_lsm_write_spec()
await table.enable_memwal(MemWalWriteSpec.bucket("id", 4))
await table.disable_memwal()
# A second unset errors.
with pytest.raises(Exception, match="no LSM write spec"):
await table.unset_lsm_write_spec()
with pytest.raises(Exception, match="no MemWAL write spec"):
await table.disable_memwal()
def test_set_identity_spec(tmp_path):
@@ -130,12 +144,12 @@ def test_set_identity_spec(tmp_path):
# Identity sharding still requires an unenforced primary key on the
# table; it shards by the raw value of the given column.
table.set_unenforced_primary_key("id")
table.set_lsm_write_spec(LsmWriteSpec.identity("v"))
table.unset_lsm_write_spec()
table.enable_memwal(MemWalWriteSpec.identity("v"))
table.disable_memwal()
def test_lsm_write_spec_identity_and_writer_config_defaults():
s = LsmWriteSpec.identity("v")
def test_memwal_spec_identity_and_writer_config_defaults():
s = MemWalWriteSpec.identity("v")
assert s.spec_type == "identity"
assert s.column == "v"
assert s.num_buckets is None
@@ -146,12 +160,12 @@ def test_lsm_write_spec_identity_and_writer_config_defaults():
assert "durable_write" in repr(s)
def test_get_lsm_write_spec(tmp_path):
def test_get_memwal_spec(tmp_path):
_db, table = _make_table(tmp_path)
table.set_unenforced_primary_key("id")
# None when nothing is installed.
assert table.get_lsm_write_spec() is None
assert table.get_memwal_spec() is None
# A real scalar index is needed to name it as a maintained index.
table.create_index("id", config=BTree())
@@ -159,12 +173,12 @@ def test_get_lsm_write_spec(tmp_path):
# Bucket spec round-trips, including maintained indexes and writer config
# defaults.
table.set_lsm_write_spec(
LsmWriteSpec.bucket("id", 4)
table.enable_memwal(
MemWalWriteSpec.bucket("id", 4)
.with_maintained_indexes([idx_name])
.with_writer_config_defaults({"durable_write": "false"})
)
spec = table.get_lsm_write_spec()
spec = table.get_memwal_spec()
assert spec is not None
assert spec.spec_type == "bucket"
assert spec.column == "id"
@@ -173,30 +187,30 @@ def test_get_lsm_write_spec(tmp_path):
assert spec.writer_config_defaults == {"durable_write": "false"}
# After unset, None again.
table.unset_lsm_write_spec()
assert table.get_lsm_write_spec() is None
table.disable_memwal()
assert table.get_memwal_spec() is None
# Identity round-trips (column recovered from the schema). An unnamed
# maintained set reads back as None: what is stored is the intent to
# maintain every index, not the set it resolves to now.
table.set_lsm_write_spec(LsmWriteSpec.identity("id"))
spec = table.get_lsm_write_spec()
table.enable_memwal(MemWalWriteSpec.identity("id"))
spec = table.get_memwal_spec()
assert spec.spec_type == "identity"
assert spec.column == "id"
assert spec.maintained_indexes is None
table.unset_lsm_write_spec()
table.disable_memwal()
# Unsharded round-trips (no routing column). Opting out is distinct from
# the unnamed default.
table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([]))
spec = table.get_lsm_write_spec()
table.enable_memwal(MemWalWriteSpec.unsharded().with_maintained_indexes([]))
spec = table.get_memwal_spec()
assert spec.spec_type == "unsharded"
assert spec.column is None
assert spec.maintained_indexes == []
@pytest.mark.asyncio
async def test_async_get_lsm_write_spec(tmp_path):
async def test_async_get_memwal_spec(tmp_path):
db = await lancedb.connect_async(
tmp_path, read_consistency_interval=timedelta(seconds=0)
)
@@ -205,20 +219,21 @@ async def test_async_get_lsm_write_spec(tmp_path):
pa.RecordBatchReader.from_batches(SCHEMA, [_batch(["seed"], [0])]),
)
assert await table.get_lsm_write_spec() is None
assert await table.get_memwal_spec() is None
await table.set_unenforced_primary_key("id")
# A real scalar index is needed to name it as a maintained index.
await table.create_index("id", config=BTree())
idx_name = (await table.list_indices())[0].name
await table.set_lsm_write_spec(
LsmWriteSpec.bucket("id", 8).with_maintained_indexes([idx_name])
await table.enable_memwal(
MemWalWriteSpec.bucket("id", 8).with_maintained_indexes([idx_name])
)
spec = await table.get_lsm_write_spec()
spec = await table.get_memwal_spec()
assert spec is not None
assert spec.spec_type == "bucket"
assert spec.column == "id"
assert spec.num_buckets == 8
assert spec.maintained_indexes == [idx_name]
await table.unset_lsm_write_spec()
assert await table.get_lsm_write_spec() is None
await table.disable_memwal()
assert await table.get_memwal_spec() is None
+50 -52
View File
@@ -8,7 +8,7 @@ from datetime import timedelta
import lancedb
import pyarrow as pa
import pytest
from lancedb._lancedb import LsmWriteSpec
from lancedb._lancedb import MemWalWriteSpec
from lancedb.index import FTS, IvfPq
SCHEMA = pa.schema(
@@ -49,12 +49,12 @@ def _region_reader(rows):
def _bucket_table(tmp_path):
"""A table with ``id`` as the primary key and a single-bucket LSM spec."""
"""A table with ``id`` as the primary key and a single-bucket MemWAL spec."""
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
table = db.create_table("t", _reader([1, 2, 3]))
table.set_unenforced_primary_key("id")
# num_buckets = 1: every row routes to the single bucket.
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 1))
table.enable_memwal(MemWalWriteSpec.bucket("id", 1))
return table
@@ -78,7 +78,7 @@ def test_lsm_merge_insert_unsharded(tmp_path):
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
table = db.create_table("t", _reader([1, 2, 3]))
table.set_unenforced_primary_key("id")
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
table.enable_memwal(MemWalWriteSpec.unsharded())
result = (
table.merge_insert("id")
.when_matched_update_all()
@@ -92,7 +92,7 @@ def test_lsm_merge_insert_identity(tmp_path):
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
table = db.create_table("t", _region_reader([(1, "us"), (2, "us")]))
table.set_unenforced_primary_key("id")
table.set_lsm_write_spec(LsmWriteSpec.identity("region"))
table.enable_memwal(MemWalWriteSpec.identity("region"))
# All rows share one identity value, so they route to one shard.
result = (
table.merge_insert([])
@@ -103,31 +103,31 @@ def test_lsm_merge_insert_identity(tmp_path):
assert result.num_rows == 2
def test_lsm_merge_insert_use_lsm_false(tmp_path):
def test_lsm_merge_insert_use_memwal_false(tmp_path):
table = _bucket_table(tmp_path) # rows id = 1, 2, 3
# use_lsm(False) opts out: the standard path runs and commits even with a spec.
# use_memwal(False) opts out: the standard path runs and commits even with a spec.
result = (
table.merge_insert("id")
.when_not_matched_insert_all()
.use_lsm(False)
.use_memwal(False)
.execute(_reader([3, 4, 5]))
)
assert result.num_inserted_rows == 2
assert table.count_rows() == 5
def test_lsm_merge_insert_use_lsm_true_without_spec_errors(tmp_path):
# A table with a primary key but no LSM write spec installed.
def test_lsm_merge_insert_use_memwal_true_without_spec_errors(tmp_path):
# A table with a primary key but no MemWAL write spec installed.
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
table = db.create_table("t", _reader([1, 2, 3]))
table.set_unenforced_primary_key("id")
# use_lsm(True) demands MemWAL routing; without a spec it errors.
with pytest.raises(Exception, match="use_lsm"):
# use_memwal(True) demands MemWAL routing; without a spec it errors.
with pytest.raises(Exception, match="use_memwal"):
(
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.use_lsm(True)
.use_memwal(True)
.execute(_reader([3, 4, 5]))
)
@@ -145,7 +145,7 @@ def test_lsm_merge_insert_validate_single_shard_off(tmp_path):
def test_lsm_merge_insert_no_spec_uses_standard_path(tmp_path):
# A table with a primary key but no LSM write spec installed.
# A table with a primary key but no MemWAL write spec installed.
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
table = db.create_table("t", _reader([1, 2, 3]))
table.set_unenforced_primary_key("id")
@@ -186,7 +186,7 @@ def test_lsm_close_writers(tmp_path):
.when_not_matched_insert_all()
.execute(_reader([7, 8]))
)
table.close_lsm_writers()
table.close_memwal_writers()
# The writer reopens lazily on the next merge_insert.
result = (
table.merge_insert([])
@@ -204,14 +204,14 @@ async def test_async_lsm_merge_insert(tmp_path):
)
table = await db.create_table("t", _reader([1, 2, 3]))
await table.set_unenforced_primary_key("id")
await table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 1))
await table.enable_memwal(MemWalWriteSpec.bucket("id", 1))
builder = (
table.merge_insert([]).when_matched_update_all().when_not_matched_insert_all()
)
result = await builder.execute(_reader([3, 4, 5]))
assert result.num_rows == 3
await table.close_lsm_writers()
await table.close_memwal_writers()
def _lsm_upsert(table, ids):
@@ -228,7 +228,7 @@ def test_lsm_read_sees_active_memtable(tmp_path):
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
table = db.create_table("t", _reader([1, 2, 3])) # base ids 1,2,3
table.set_unenforced_primary_key("id")
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
table.enable_memwal(MemWalWriteSpec.unsharded())
_lsm_upsert(table, [4, 5]) # active memtable only, not committed to base
@@ -236,8 +236,8 @@ def test_lsm_read_sees_active_memtable(tmp_path):
lsm = table.search().to_arrow()
assert sorted(lsm["id"].to_pylist()) == [1, 2, 3, 4, 5]
# use_lsm(False) bypasses the MemWAL and reads the base table only.
base_only = table.search().use_lsm(False).to_arrow()
# use_memwal(False) bypasses the MemWAL and reads the base table only.
base_only = table.search().use_memwal(False).to_arrow()
assert sorted(base_only["id"].to_pylist()) == [1, 2, 3]
@@ -245,7 +245,7 @@ def test_lsm_read_dedup_newest_wins(tmp_path):
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
table = db.create_table("t", _reader([1, 2, 3])) # id 2 -> value 1
table.set_unenforced_primary_key("id")
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
table.enable_memwal(MemWalWriteSpec.unsharded())
_lsm_upsert(table, [2, 3, 4]) # ids 2,3,4 -> values 0,1,2
@@ -258,22 +258,22 @@ def test_lsm_read_dedup_newest_wins(tmp_path):
def test_lsm_read_without_spec_reads_base(tmp_path):
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
table = db.create_table("t", _reader([1, 2, 3]))
table.set_unenforced_primary_key("id") # no LSM write spec
table.set_unenforced_primary_key("id") # no MemWAL write spec
# No spec: default read and use_lsm(False) both read the base table, no error.
# No spec: default read and use_memwal(False) both read the base table, no error.
assert sorted(table.search().to_arrow()["id"].to_pylist()) == [1, 2, 3]
assert sorted(table.search().use_lsm(False).to_arrow()["id"].to_pylist()) == [
assert sorted(table.search().use_memwal(False).to_arrow()["id"].to_pylist()) == [
1,
2,
3,
]
def test_lsm_read_unsupported_shape_errors_without_use_lsm_false(tmp_path):
def test_lsm_read_unsupported_shape_errors_without_use_memwal_false(tmp_path):
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
table = db.create_table("t", _reader([1, 2, 3]))
table.set_unenforced_primary_key("id")
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
table.enable_memwal(MemWalWriteSpec.unsharded())
_lsm_upsert(table, [4])
# with_row_id is unsupported by the LSM scanner; on a MemWAL table the default
@@ -281,8 +281,8 @@ def test_lsm_read_unsupported_shape_errors_without_use_lsm_false(tmp_path):
with pytest.raises(Exception):
table.search().with_row_id(True).to_arrow()
# use_lsm(False) is the escape hatch: it reads the base table only.
base = table.search().with_row_id(True).use_lsm(False).to_arrow()
# use_memwal(False) is the escape hatch: it reads the base table only.
base = table.search().with_row_id(True).use_memwal(False).to_arrow()
assert sorted(base["id"].to_pylist()) == [1, 2, 3]
@@ -293,7 +293,7 @@ async def test_async_lsm_read(tmp_path):
)
table = await db.create_table("t", _reader([1, 2, 3]))
await table.set_unenforced_primary_key("id")
await table.set_lsm_write_spec(LsmWriteSpec.unsharded())
await table.enable_memwal(MemWalWriteSpec.unsharded())
builder = (
table.merge_insert([]).when_matched_update_all().when_not_matched_insert_all()
@@ -331,7 +331,7 @@ def _vector_reader(rows):
def _vector_table(tmp_path):
"""Base table whose vector column is indexed so its rows are visible to the LSM
vector scanner (the base arm uses ``fast_search`` — indexed data only), plus an
unsharded LSM spec that maintains that index for the memtable.
unsharded MemWAL spec that maintains that index for the memtable.
Rows 1,2 are category ``a``, row 3 is ``b``, and 4..60 are filler ``c`` that
give the tiny IVF index enough data to train.
@@ -353,8 +353,8 @@ def _vector_table(tmp_path):
"vector", config=IvfPq(num_partitions=1, num_sub_vectors=2, num_bits=4)
)
index_name = table.list_indices()[0].name
table.set_lsm_write_spec(
LsmWriteSpec.unsharded().with_maintained_indexes([index_name])
table.enable_memwal(
MemWalWriteSpec.unsharded().with_maintained_indexes([index_name])
)
return table
@@ -379,9 +379,9 @@ def test_lsm_read_vector_sees_memtable(tmp_path):
assert {1, 2, 3} <= ids # indexed base rows
assert 1000 in ids # in-flight memtable row
# use_lsm(False) bypasses the MemWAL, so the in-flight row is not visible.
# use_memwal(False) bypasses the MemWAL, so the in-flight row is not visible.
base_ids = set(
table.search(query).use_lsm(False).limit(100).to_arrow()["id"].to_pylist()
table.search(query).use_memwal(False).limit(100).to_arrow()["id"].to_pylist()
)
assert {1, 2, 3} <= base_ids
assert 1000 not in base_ids
@@ -453,8 +453,8 @@ def test_lsm_read_fts_sees_memtable(tmp_path):
# Native FTS index (tantivy is not compatible with the LSM memtable index).
table.create_index("text", config=FTS())
index_name = table.list_indices()[0].name
table.set_lsm_write_spec(
LsmWriteSpec.unsharded().with_maintained_indexes([index_name])
table.enable_memwal(
MemWalWriteSpec.unsharded().with_maintained_indexes([index_name])
)
# in-flight doc 4 lands in the memtable's maintained FTS index.
@@ -498,8 +498,8 @@ def test_lsm_read_vector_unsupported_knobs_error(tmp_path):
with pytest.raises(Exception, match="use_index"):
table.search(query).bypass_vector_index().to_arrow()
# use_lsm(False) is the escape hatch: the base-only standard path honors them.
base = table.search(query).distance_range(0.0, 100.0).use_lsm(False).to_arrow()
# use_memwal(False) is the escape hatch: the base-only standard path honors them.
base = table.search(query).distance_range(0.0, 100.0).use_memwal(False).to_arrow()
assert 1000 not in set(base["id"].to_pylist())
@@ -544,7 +544,7 @@ def test_lsm_read_fts_unmaintained_index_errors(tmp_path):
table.create_index("text", config=FTS())
# No maintained indexes: the active memtable FTS arm cannot serve un-compacted
# docs, so the search would silently omit them — reject instead.
table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([]))
table.enable_memwal(MemWalWriteSpec.unsharded().with_maintained_indexes([]))
with pytest.raises(Exception, match="maintained"):
table.search("fox", query_type="fts", fts_columns="text").to_arrow()
@@ -553,7 +553,7 @@ def test_lsm_read_time_travel_errors(tmp_path):
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
table = db.create_table("t", _reader([1, 2, 3]))
table.set_unenforced_primary_key("id")
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
table.enable_memwal(MemWalWriteSpec.unsharded())
pinned = table.version
table.add(_reader([4, 5])) # standard add commits a newer version
table.checkout(pinned) # detached head at the historical version
@@ -562,8 +562,8 @@ def test_lsm_read_time_travel_errors(tmp_path):
# historical version is rejected.
with pytest.raises(Exception, match="time-travel"):
table.search().to_arrow()
# use_lsm(False) reads the base table at the pinned version.
base = table.search().use_lsm(False).to_arrow()
# use_memwal(False) reads the base table at the pinned version.
base = table.search().use_memwal(False).to_arrow()
assert sorted(base["id"].to_pylist()) == [1, 2, 3]
@@ -571,14 +571,14 @@ def test_lsm_read_take_row_ids_errors(tmp_path):
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
table = db.create_table("t", _reader([1, 2, 3]))
table.set_unenforced_primary_key("id")
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
table.enable_memwal(MemWalWriteSpec.unsharded())
_lsm_upsert(table, [4])
# take-by-row-id auto-routes through the LSM scanner, which has no stable _rowid,
# so it hard-errors instead of failing with an opaque column-not-found error.
with pytest.raises(Exception, match="row id"):
table.take_row_ids([0, 1]).to_arrow()
# use_lsm(False) is the escape hatch: it reads the base table.
base = table.take_row_ids([0, 1]).use_lsm(False).to_arrow()
# use_memwal(False) is the escape hatch: it reads the base table.
base = table.take_row_ids([0, 1]).use_memwal(False).to_arrow()
assert base.num_rows == 2
@@ -588,8 +588,8 @@ def test_lsm_read_fts_postfilter_errors(tmp_path):
table.set_unenforced_primary_key("id")
table.create_index("text", config=FTS())
index_name = table.list_indices()[0].name
table.set_lsm_write_spec(
LsmWriteSpec.unsharded().with_maintained_indexes([index_name])
table.enable_memwal(
MemWalWriteSpec.unsharded().with_maintained_indexes([index_name])
)
# The LSM scanner always prefilters; postfilter on FTS changes result semantics,
# so it hard-errors (previously only the vector arm rejected it).
@@ -608,9 +608,7 @@ def test_lsm_read_fts_multiple_same_type_indexes_errors(tmp_path):
table.set_unenforced_primary_key("id")
table.create_index("text", config=FTS(), name="fts_a")
table.create_index("text", config=FTS(), name="fts_b", replace=False)
table.set_lsm_write_spec(
LsmWriteSpec.unsharded().with_maintained_indexes(["fts_a"])
)
table.enable_memwal(MemWalWriteSpec.unsharded().with_maintained_indexes(["fts_a"]))
# Two FTS indexes on the column: the base planner's chosen index is ambiguous, so
# the scanner can't pick a catch-up watermark and rejects rather than risk
# dropping rows the actually-used index has not caught up to.
@@ -631,7 +629,7 @@ def test_lsm_read_vector_unmaintained_index_errors(tmp_path):
)
# Spec with NO maintained indexes: the base vector index's catch-up is untracked,
# so the scanner rejects rather than risk dropping compacted-but-unindexed rows.
table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([]))
table.enable_memwal(MemWalWriteSpec.unsharded().with_maintained_indexes([]))
with pytest.raises(Exception, match="maintained"):
table.search([1.0] * VECTOR_DIM).to_arrow()
@@ -644,7 +642,7 @@ def test_lsm_read_fts_optimized_index_not_rejected(tmp_path):
table.add(_fts_reader([(i, "lazy fox") for i in range(6, 11)]))
table.optimize() # may split the FTS index into multiple physical segments
name = table.list_indices()[0].name
table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([name]))
table.enable_memwal(MemWalWriteSpec.unsharded().with_maintained_indexes([name]))
# Multiple physical segments of one logical index must not be miscounted as
# multiple indexes and rejected.
ids = set(
+12
View File
@@ -1209,3 +1209,15 @@ def test_local_table_to_arrow_and_to_pandas_are_unchanged(tmp_path):
assert table.to_arrow().column("id").to_pylist() == [1, 2]
assert table.to_pandas()["id"].tolist() == [1, 2]
@pytest.mark.parametrize("effort", [0.0, 0.5, 1.0])
def test_namespace_request_rejects_search_effort(effort):
from lancedb.query import Query
with pytest.raises(
ValueError, match="namespace QueryTable does not support search_effort"
):
_query_to_namespace_request(
["test"], Query(vector=[1.0, 0.0], search_effort=effort)
)
+348 -5
View File
@@ -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):
@@ -2265,3 +2484,127 @@ def test_blob_v2_to_batches_row_id(tmp_db):
assert "_rowid" in hits.column_names
blobs = table.fetch_blobs("blob", hits)
assert [blobs[i].as_py() for i in range(len(blobs))] == [b"one", b"two"]
@pytest.mark.parametrize("metric", ["l2", "cosine", "dot"])
@pytest.mark.parametrize(
"effort,probes", [(None, 4), (0, 1), (0.25, 2), (0.5, 4), (0.75, 8), (1, 16)]
)
def test_search_effort_probe_budget(tmp_path, monkeypatch, metric, effort, probes):
import re
from lancedb.index import IvfFlat
monkeypatch.setenv("LANCE_AUTO_PROBE_MARGIN", "0")
monkeypatch.setenv("LANCE_AUTO_MIN_INITIAL_NPROBES", "4")
monkeypatch.setenv("LANCE_AUTO_MAX_INITIAL_NPROBES", "4")
rng = np.random.default_rng(42)
vectors = np.repeat(np.eye(16, dtype=np.float32), 32, axis=0)
vectors += rng.normal(0, 0.001, vectors.shape).astype(np.float32)
data = pa.table(
{
"id": np.arange(512),
"text": ["dog"] * 512,
"vector": pa.FixedSizeListArray.from_arrays(pa.array(vectors.ravel()), 16),
}
)
table = lancedb.connect(tmp_path).create_table("effort", data.slice(0, 256))
table.add(data.slice(256))
table.create_index(
"vector", config=IvfFlat(distance_type=metric, num_partitions=16)
)
vector = np.eye(16, dtype=np.float32)[0]
def search(value=effort):
return table.search(vector).distance_type(metric).search_effort(value).limit(8)
def partitions(query):
vector_plan = query.analyze_plan().split("FTS Search Plan:", 1)[0]
values = re.findall(r"partitions_searched=(\d+)", vector_plan)
assert values
return sum(map(int, values))
assert partitions(search()) == probes
exact = set(
table.search(vector)
.distance_type(metric)
.bypass_vector_index()
.limit(8)
.to_arrow()["id"]
.to_pylist()
)
actual = set(search().to_arrow()["id"].to_pylist())
assert len(actual & exact) / 8 >= 0.5
expected = 1 if effort is None else probes
assert partitions(search().nprobes(1)) == expected
assert (
partitions(
table.search(vector)
.distance_type(metric)
.nprobes(1)
.search_effort(effort)
.limit(8)
)
== expected
)
assert partitions(search().nprobes(1).search_effort(None)) == 1
assert partitions(search().maximum_nprobes(2)) <= 2
assert partitions(search().minimum_nprobes(3).maximum_nprobes(3)) == 3
assert partitions(search(0).nprobes(16).maximum_nprobes(0)) == 1
if effort is None:
assert search().to_arrow() == search(0.5).to_arrow()
filtered = search(0).maximum_nprobes(2).where("id % 8 = 0", prefilter=True)
assert partitions(filtered) <= 2
assert all(row["id"] % 8 == 0 for row in filtered.to_list())
if metric == "l2" and effort == 0.5:
table.create_index("text", config=FTS())
hybrid = (
table.search(query_type="hybrid")
.vector(vector)
.text("dog")
.nprobes(1)
.search_effort(0.5)
.limit(8)
)
assert partitions(hybrid) == 4
assert partitions(hybrid.search_effort(None)) == 1
@pytest.mark.parametrize(
"effort", [-0.1, 1.1, float("nan"), float("inf"), -float("inf")]
)
def test_search_effort_invalid_sync(table, effort):
with pytest.raises(ValueError, match="search_effort must be finite and in"):
table.search([0, 0]).search_effort(effort).to_list()
@pytest.mark.asyncio
@pytest.mark.parametrize("hybrid", [False, True])
@pytest.mark.parametrize("effort", [None, 0, 0.25, 0.5, 0.75, 1])
async def test_search_effort_async_roundtrip(table_async, hybrid, effort):
query = table_async.vector_search([0, 0])
if hybrid:
query = query.nearest_to_text("dog")
query = query.nprobes(5).search_effort(effort).minimum_nprobes(2).maximum_nprobes(3)
request = query.to_query_object()
assert request.search_effort == effort
assert request.nprobes == 5
assert request.minimum_nprobes == 2
assert request.maximum_nprobes == 3
assert query.search_effort(None).to_query_object().search_effort is None
query.search_effort(effort)
if not hybrid:
assert len(await query.limit(1).to_list()) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("hybrid", [False, True])
@pytest.mark.parametrize(
"effort", [-0.1, 1.1, float("nan"), float("inf"), -float("inf")]
)
async def test_search_effort_invalid_async(table_async, hybrid, effort):
query = table_async.vector_search([0, 0])
if hybrid:
query = query.nearest_to_text("dog")
with pytest.raises(ValueError, match="search_effort must be finite and in"):
query.search_effort(effort)
+147 -29
View File
@@ -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
@@ -1622,7 +1623,6 @@ def remote_unsupported_operations():
return [
("uses_v2_manifest_paths", (), {}, "uses_v2_manifest_paths"),
("migrate_v2_manifest_paths", (), {}, "migrate_manifest_paths_v2"),
("set_unenforced_primary_key", (["id"],), {}, "set_unenforced_primary_key"),
("fetch_blob_ranges", ("image", [(1, 0, 1)]), {}, "fetch_blob_ranges"),
("optimize", (), {}, "optimize"),
("create_index", ("vector",), {"config": HnswPq()}, "IVF_HNSW_PQ"),
@@ -1845,10 +1845,10 @@ def lsm_test_table(lsm_handler):
"""A remote table whose LSM routes are served by ``lsm_handler``.
``lsm_handler(request, route)`` is called for ``/v1/table/test/<route>/``
where route is one of flush_lsm, compact_lsm, get_lsm_stats, and is
where route is one of flush_memwal, compact_memwal, get_memwal_stats, and is
responsible for writing the response.
"""
routes = ("flush_lsm", "compact_lsm", "get_lsm_stats")
routes = ("flush_memwal", "compact_memwal", "get_memwal_stats")
def handler(request):
match = re.fullmatch(r"/v1/table/test/(\w+)/", request.path)
@@ -1880,9 +1880,9 @@ def send_json(request, payload, status=200):
request.wfile.write(json.dumps(payload).encode())
def test_get_lsm_stats_sync():
def test_get_memwal_stats_sync():
"""The sync wrapper round-trips the server payload into a dict."""
bucket = {
table_shard = {
"shard_id": "b0",
"status": "Active",
"writer_epoch": 3,
@@ -1890,7 +1890,7 @@ def test_get_lsm_stats_sync():
"current_generation": 6,
"replay_after_wal_entry_position": 40,
"wal_entry_position_last_seen": 42,
"generations": [{"generation": 5, "bytes": 1024, "rows": 7}],
"sstables": [{"generation": 5, "bytes": 1024, "rows": 7}],
"compacting": False,
"memtables": [
{
@@ -1905,29 +1905,29 @@ def test_get_lsm_stats_sync():
seen_bodies = []
def lsm_handler(request, route):
assert route == "get_lsm_stats"
assert route == "get_memwal_stats"
seen_bodies.append(read_json_body(request))
send_json(request, {"lsm_stats": {"buckets": [bucket]}})
send_json(request, {"memwal_stats": {"table_shards": [table_shard]}})
with lsm_test_table(lsm_handler) as table:
assert table.get_lsm_stats() == {"buckets": [bucket]}
assert table.get_memwal_stats() == {"table_shards": [table_shard]}
# Off by default, and forwarded when asked for.
assert seen_bodies == [{"include_generation_rows": False}]
table.get_lsm_stats(include_generation_rows=True)
assert seen_bodies[-1] == {"include_generation_rows": True}
assert seen_bodies == [{"include_sstable_rows": False}]
table.get_memwal_stats(include_sstable_rows=True)
assert seen_bodies[-1] == {"include_sstable_rows": True}
def test_get_lsm_stats_sync_returns_none_when_lsm_disabled():
"""A null envelope means the LSM write path is not enabled, not an error."""
def test_get_memwal_stats_sync_returns_none_when_memwal_disabled():
"""A null envelope means the MemWAL write path is not enabled, not an error."""
def lsm_handler(request, route):
send_json(request, {"lsm_stats": None})
send_json(request, {"memwal_stats": None})
with lsm_test_table(lsm_handler) as table:
assert table.get_lsm_stats() is None
assert table.get_memwal_stats() is None
def test_flush_and_compact_lsm_sync():
def test_flush_and_compact_memwal_sync():
"""Both are one-shot POSTs answered 202 with no body."""
called = []
@@ -1937,13 +1937,13 @@ def test_flush_and_compact_lsm_sync():
request.end_headers()
with lsm_test_table(lsm_handler) as table:
assert table.flush_lsm() is None
assert table.compact_lsm() is None
assert called == ["flush_lsm", "compact_lsm"]
assert table.flush_memwal() is None
assert table.compact_memwal() is None
assert called == ["flush_memwal", "compact_memwal"]
def test_checkpoint_lsm_sync():
"""Seal, read the watermark, and return once L0 holds nothing.
def test_checkpoint_memwal_sync():
"""Freeze, read the watermark, and return once no SSTables remain.
The convergence loop itself is covered in Rust; this pins the sync
binding to the endpoints it drives.
@@ -1952,17 +1952,17 @@ def test_checkpoint_lsm_sync():
def lsm_handler(request, route):
called.append(route)
if route == "get_lsm_stats":
# An empty L0 yields no target watermark, so the loop is done
# after the seal without ever polling compaction.
send_json(request, {"lsm_stats": {"buckets": []}})
if route == "get_memwal_stats":
# An empty SSTable tier yields no target watermark, so the loop is done
# after the freeze without ever polling compaction.
send_json(request, {"memwal_stats": {"table_shards": []}})
else:
request.send_response(202)
request.end_headers()
with lsm_test_table(lsm_handler) as table:
assert table.checkpoint_lsm() is None
assert called == ["flush_lsm", "get_lsm_stats"]
assert table.checkpoint_memwal() is None
assert called == ["flush_memwal", "get_memwal_stats"]
@contextlib.contextmanager
@@ -2728,6 +2728,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 +3468,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 +3538,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 +3583,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 +3662,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():
+20
View File
@@ -4083,6 +4083,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(
+78 -1
View File
@@ -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)
+8 -6
View File
@@ -81,15 +81,15 @@ def test_base_table_hybrid_still_joins_on_row_ids():
assert builder._vector_query._columns is None
def test_wal_hybrid_with_use_lsm_false_joins_on_row_ids():
builder = _wal_hybrid().use_lsm(False)
def test_wal_hybrid_with_use_memwal_false_joins_on_row_ids():
builder = _wal_hybrid().use_memwal(False)
builder._create_query_builders()
assert builder._pk_fusion is None
def test_wal_hybrid_refuses_an_explicit_with_row_id():
builder = _wal_hybrid().with_row_id(True)
with pytest.raises(NotImplementedError, match="use_lsm"):
with pytest.raises(NotImplementedError, match="use_memwal"):
builder._create_query_builders()
@@ -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`"
@@ -139,7 +141,7 @@ def test_wal_hybrid_refuses_row_id_even_with_return_score_all(row_id_first: bool
else:
builder = builder.rerank(reranker).with_row_id(True)
with pytest.raises(NotImplementedError, match="use_lsm"):
with pytest.raises(NotImplementedError, match="use_memwal"):
builder._create_query_builders()
@@ -197,7 +199,7 @@ def test_fallback_explains_the_refusal_to_a_caller_who_asked_for_row_ids():
def run():
raise _REFUSAL
with pytest.raises(NotImplementedError, match="use_lsm") as info:
with pytest.raises(NotImplementedError, match="use_memwal") as info:
with_pk_fallback(
run, mock.Mock(), fused=lambda: False, caller_requested_row_id=True
)
+3 -3
View File
@@ -73,10 +73,10 @@ class MockPermutationServer:
},
},
)
if path == f"/v1/table/{self.name}/get_lsm_write_spec/":
if path == f"/v1/table/{self.name}/get_memwal_spec/":
self._read_body(request)
# Null spec: this table has no LSM write path.
return self._json(request, {"lsm_write_spec": None})
# Null spec: this table has no MemWAL write path.
return self._json(request, {"memwal_spec": None})
if path == f"/v1/table/{self.name}/count_rows/":
self._read_body(request)
return self._json(request, self.num_rows)
+280
View File
@@ -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)
}
}
+7
View File
@@ -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()
+2
View File
@@ -1224,6 +1224,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>,
@@ -1309,6 +1310,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),
+7 -2
View File
@@ -16,12 +16,13 @@ use query::{FTSQuery, HybridQuery, Query, VectorQuery};
use session::Session;
use table::{
AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, FtsToken,
FunctionErrorFragment, FunctionErrorRecord, FunctionErrors, LsmWriteSpec, MergeResult,
FunctionErrorFragment, FunctionErrorRecord, FunctionErrors, MemWalWriteSpec, MergeResult,
PyBlobFile, RefreshColumnResult, RefreshMaterializedViewResult, Table,
UpdateFieldMetadataResult, UpdateResult,
};
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>()?;
@@ -99,7 +104,7 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<UpdateFieldMetadataResult>()?;
m.add_class::<AddResult>()?;
m.add_class::<MergeResult>()?;
m.add_class::<LsmWriteSpec>()?;
m.add_class::<MemWalWriteSpec>()?;
m.add_class::<DeleteResult>()?;
m.add_class::<DropColumnsResult>()?;
m.add_class::<UpdateResult>()?;
+45 -17
View File
@@ -71,7 +71,7 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB<FtsQuery> {
match class_name(&ob)?.as_str() {
"MatchQuery" => {
let query = ob.getattr("query")?.extract()?;
let column = ob.getattr("column")?.extract()?;
let column = ob.getattr("column")?.extract::<Option<String>>()?;
let boost = ob.getattr("boost")?.extract()?;
let fuzziness = ob.getattr("fuzziness")?.extract()?;
let max_expansions = ob.getattr("max_expansions")?.extract()?;
@@ -88,7 +88,7 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB<FtsQuery> {
let mut query =
MatchQuery::new(query)
.with_column(Some(column))
.with_column(column)
.with_boost(boost)
.with_fuzziness(fuzziness)
.with_max_expansions(max_expansions)
@@ -199,7 +199,7 @@ impl<'py> IntoPyObject<'py> for PyLanceDB<FtsQuery> {
}
namespace
.getattr(intern!(py, "MatchQuery"))?
.call((query.terms, query.column.unwrap()), Some(&kwargs))
.call((query.terms, query.column), Some(&kwargs))
}
FtsQuery::Phrase(query) => {
let kwargs = PyDict::new(py);
@@ -333,12 +333,13 @@ pub struct PyQueryRequest {
pub select_source_columns: Option<HashMap<String, String>>,
pub fast_search: Option<bool>,
pub with_row_id: Option<bool>,
pub use_lsm: Option<bool>,
pub use_memwal: Option<bool>,
pub column: Option<String>,
pub query_vector: Option<PyQueryVectors>,
pub nprobes: Option<usize>,
pub minimum_nprobes: Option<usize>,
pub maximum_nprobes: Option<usize>,
pub search_effort: Option<f64>,
pub lower_bound: Option<f32>,
pub upper_bound: Option<f32>,
pub ef: Option<usize>,
@@ -365,12 +366,13 @@ impl From<AnyQuery> for PyQueryRequest {
select: PySelect(query_request.select),
fast_search: Some(query_request.fast_search),
with_row_id: Some(query_request.with_row_id),
use_lsm: query_request.use_lsm,
use_memwal: query_request.use_memwal,
column: None,
query_vector: None,
nprobes: None,
minimum_nprobes: None,
maximum_nprobes: None,
search_effort: None,
lower_bound: None,
upper_bound: None,
ef: None,
@@ -393,12 +395,13 @@ impl From<AnyQuery> for PyQueryRequest {
select: PySelect(vector_query.base.select),
fast_search: Some(vector_query.base.fast_search),
with_row_id: Some(vector_query.base.with_row_id),
use_lsm: vector_query.base.use_lsm,
use_memwal: vector_query.base.use_memwal,
column: vector_query.column,
query_vector: Some(PyQueryVectors(vector_query.query_vector)),
nprobes: vector_query.nprobes,
minimum_nprobes: vector_query.minimum_nprobes,
maximum_nprobes: vector_query.maximum_nprobes,
search_effort: vector_query.search_effort,
lower_bound: vector_query.lower_bound,
upper_bound: vector_query.upper_bound,
ef: vector_query.ef,
@@ -537,8 +540,8 @@ impl Query {
self.inner = self.inner.clone().fast_search();
}
pub fn use_lsm(&mut self, enable: bool) {
self.inner = self.inner.clone().use_lsm(enable);
pub fn use_memwal(&mut self, enable: bool) {
self.inner = self.inner.clone().use_memwal(enable);
}
pub fn with_row_id(&mut self) {
@@ -703,8 +706,8 @@ impl TakeQuery {
self.inner = self.inner.clone().with_row_id();
}
pub fn use_lsm(&mut self, enable: bool) {
self.inner = self.inner.clone().use_lsm(enable);
pub fn use_memwal(&mut self, enable: bool) {
self.inner = self.inner.clone().use_memwal(enable);
}
#[pyo3(signature = ())]
@@ -816,8 +819,8 @@ impl FTSQuery {
self.inner = self.inner.clone().fast_search();
}
pub fn use_lsm(&mut self, enable: bool) {
self.inner = self.inner.clone().use_lsm(enable);
pub fn use_memwal(&mut self, enable: bool) {
self.inner = self.inner.clone().use_memwal(enable);
}
pub fn with_row_id(&mut self) {
@@ -967,8 +970,8 @@ impl VectorQuery {
self.inner = self.inner.clone().fast_search();
}
pub fn use_lsm(&mut self, enable: bool) {
self.inner = self.inner.clone().use_lsm(enable);
pub fn use_memwal(&mut self, enable: bool) {
self.inner = self.inner.clone().use_memwal(enable);
}
pub fn with_row_id(&mut self) {
@@ -998,6 +1001,15 @@ impl VectorQuery {
Ok(())
}
pub fn search_effort(&mut self, search_effort: Option<f64>) -> PyResult<()> {
self.inner = self
.inner
.clone()
.search_effort(search_effort)
.infer_error()?;
Ok(())
}
pub fn minimum_nprobes(&mut self, minimum_nprobes: u32) -> PyResult<()> {
self.inner = self
.inner
@@ -1166,9 +1178,9 @@ impl HybridQuery {
self.inner_fts.postfilter();
}
pub fn use_lsm(&mut self, enable: bool) {
self.inner_vec.use_lsm(enable);
self.inner_fts.use_lsm(enable);
pub fn use_memwal(&mut self, enable: bool) {
self.inner_vec.use_memwal(enable);
self.inner_fts.use_memwal(enable);
}
pub fn add_query_vector(&mut self, vector: Bound<'_, PyAny>) -> PyResult<()> {
@@ -1191,6 +1203,18 @@ impl HybridQuery {
self.inner_vec.nprobes(nprobe)
}
pub fn search_effort(&mut self, search_effort: Option<f64>) -> PyResult<()> {
self.inner_vec.search_effort(search_effort)
}
pub fn minimum_nprobes(&mut self, minimum_nprobes: u32) -> PyResult<()> {
self.inner_vec.minimum_nprobes(minimum_nprobes)
}
pub fn maximum_nprobes(&mut self, maximum_nprobes: u32) -> PyResult<()> {
self.inner_vec.maximum_nprobes(maximum_nprobes)
}
pub fn ef(&mut self, ef: u32) {
self.inner_vec.ef(ef);
}
@@ -1235,6 +1259,10 @@ impl HybridQuery {
req.query_vector = vec_req.query_vector;
req.column = vec_req.column;
req.distance_type = vec_req.distance_type;
req.nprobes = vec_req.nprobes;
req.minimum_nprobes = vec_req.minimum_nprobes;
req.maximum_nprobes = vec_req.maximum_nprobes;
req.search_effort = vec_req.search_effort;
req.ef = vec_req.ef;
req.refine_factor = vec_req.refine_factor;
req.lower_bound = vec_req.lower_bound;
+128 -86
View File
@@ -38,16 +38,16 @@ use pyo3::{
mod scannable;
/// Convert `LsmStats` to a Python dict, preserving the per-bucket list.
/// Convert `MemWalStats` to a Python dict, preserving the per-table-shard list.
///
/// Deliberately not flattened to a table-level summary: a table is N
/// buckets on one node, and the per-bucket detail is the reason the
/// endpoint exists — flattening hides the single hot bucket someone opened
/// table shards on one node, and the per-shard detail is the reason the
/// endpoint exists — flattening hides the single hot table shard someone opened
/// it to find.
fn lsm_stats_to_py(py: Python<'_>, stats: &lancedb::table::LsmStats) -> PyResult<Py<PyDict>> {
fn memwal_stats_to_py(py: Python<'_>, stats: &lancedb::table::MemWalStats) -> PyResult<Py<PyDict>> {
let out = PyDict::new(py);
let buckets = PyList::empty(py);
for b in &stats.buckets {
let table_shards = PyList::empty(py);
for b in &stats.table_shards {
let e = PyDict::new(py);
e.set_item("shard_id", &b.shard_id)?;
e.set_item("status", &b.status)?;
@@ -63,15 +63,15 @@ fn lsm_stats_to_py(py: Python<'_>, stats: &lancedb::table::LsmStats) -> PyResult
b.wal_entry_position_last_seen,
)?;
let generations = PyList::empty(py);
for g in &b.generations {
let sstables = PyList::empty(py);
for g in &b.sstables {
let ge = PyDict::new(py);
ge.set_item("generation", g.generation)?;
ge.set_item("bytes", g.bytes)?;
ge.set_item("rows", g.rows)?;
generations.append(ge)?;
sstables.append(ge)?;
}
e.set_item("generations", generations)?;
e.set_item("sstables", sstables)?;
e.set_item("compacting", b.compacting)?;
e.set_item(
@@ -93,9 +93,9 @@ fn lsm_stats_to_py(py: Python<'_>, stats: &lancedb::table::LsmStats) -> PyResult
})
.transpose()?,
)?;
buckets.append(e)?;
table_shards.append(e)?;
}
out.set_item("buckets", buckets)?;
out.set_item("table_shards", table_shards)?;
Ok(out.unbind())
}
@@ -269,17 +269,17 @@ fn fmt_maintained(maintained: &Option<Vec<String>>) -> String {
/// table has, including ones created later.
#[pyclass(module = "lancedb._lancedb", from_py_object)]
#[derive(Clone, Debug)]
pub struct LsmWriteSpec {
inner: lancedb::table::LsmWriteSpec,
pub struct MemWalWriteSpec {
inner: lancedb::table::MemWalWriteSpec,
}
#[pymethods]
impl LsmWriteSpec {
impl MemWalWriteSpec {
/// Hash-bucket sharding by the unenforced primary key column.
#[staticmethod]
pub fn bucket(column: String, num_buckets: u32) -> Self {
Self {
inner: lancedb::table::LsmWriteSpec::bucket(column, num_buckets),
inner: lancedb::table::MemWalWriteSpec::bucket(column, num_buckets),
}
}
@@ -293,7 +293,7 @@ impl LsmWriteSpec {
#[staticmethod]
pub fn identity(column: String) -> Self {
Self {
inner: lancedb::table::LsmWriteSpec::identity(column),
inner: lancedb::table::MemWalWriteSpec::identity(column),
}
}
@@ -302,7 +302,7 @@ impl LsmWriteSpec {
#[staticmethod]
pub fn unsharded() -> Self {
Self {
inner: lancedb::table::LsmWriteSpec::unsharded(),
inner: lancedb::table::MemWalWriteSpec::unsharded(),
}
}
@@ -326,33 +326,33 @@ impl LsmWriteSpec {
pub fn __repr__(&self) -> String {
match &self.inner {
lancedb::table::LsmWriteSpec::Bucket {
lancedb::table::MemWalWriteSpec::Bucket {
column,
num_buckets,
maintained_indexes,
writer_config_defaults,
} => format!(
"LsmWriteSpec.bucket(column={:?}, num_buckets={}, maintained_indexes={}, writer_config_defaults={:?})",
"MemWalWriteSpec.bucket(column={:?}, num_buckets={}, maintained_indexes={}, writer_config_defaults={:?})",
column,
num_buckets,
fmt_maintained(maintained_indexes),
writer_config_defaults,
),
lancedb::table::LsmWriteSpec::Identity {
lancedb::table::MemWalWriteSpec::Identity {
column,
maintained_indexes,
writer_config_defaults,
} => format!(
"LsmWriteSpec.identity(column={:?}, maintained_indexes={}, writer_config_defaults={:?})",
"MemWalWriteSpec.identity(column={:?}, maintained_indexes={}, writer_config_defaults={:?})",
column,
fmt_maintained(maintained_indexes),
writer_config_defaults,
),
lancedb::table::LsmWriteSpec::Unsharded {
lancedb::table::MemWalWriteSpec::Unsharded {
maintained_indexes,
writer_config_defaults,
} => format!(
"LsmWriteSpec.unsharded(maintained_indexes={}, writer_config_defaults={:?})",
"MemWalWriteSpec.unsharded(maintained_indexes={}, writer_config_defaults={:?})",
fmt_maintained(maintained_indexes),
writer_config_defaults,
),
@@ -364,9 +364,9 @@ impl LsmWriteSpec {
#[getter]
pub fn spec_type(&self) -> &'static str {
match &self.inner {
lancedb::table::LsmWriteSpec::Bucket { .. } => "bucket",
lancedb::table::LsmWriteSpec::Identity { .. } => "identity",
lancedb::table::LsmWriteSpec::Unsharded { .. } => "unsharded",
lancedb::table::MemWalWriteSpec::Bucket { .. } => "bucket",
lancedb::table::MemWalWriteSpec::Identity { .. } => "identity",
lancedb::table::MemWalWriteSpec::Unsharded { .. } => "unsharded",
}
}
@@ -374,9 +374,9 @@ impl LsmWriteSpec {
#[getter]
pub fn column(&self) -> Option<String> {
match &self.inner {
lancedb::table::LsmWriteSpec::Bucket { column, .. }
| lancedb::table::LsmWriteSpec::Identity { column, .. } => Some(column.clone()),
lancedb::table::LsmWriteSpec::Unsharded { .. } => None,
lancedb::table::MemWalWriteSpec::Bucket { column, .. }
| lancedb::table::MemWalWriteSpec::Identity { column, .. } => Some(column.clone()),
lancedb::table::MemWalWriteSpec::Unsharded { .. } => None,
}
}
@@ -384,7 +384,7 @@ impl LsmWriteSpec {
#[getter]
pub fn num_buckets(&self) -> Option<u32> {
match &self.inner {
lancedb::table::LsmWriteSpec::Bucket { num_buckets, .. } => Some(*num_buckets),
lancedb::table::MemWalWriteSpec::Bucket { num_buckets, .. } => Some(*num_buckets),
_ => None,
}
}
@@ -402,14 +402,14 @@ impl LsmWriteSpec {
}
}
impl From<LsmWriteSpec> for lancedb::table::LsmWriteSpec {
fn from(spec: LsmWriteSpec) -> Self {
impl From<MemWalWriteSpec> for lancedb::table::MemWalWriteSpec {
fn from(spec: MemWalWriteSpec) -> Self {
spec.inner
}
}
impl From<lancedb::table::LsmWriteSpec> for LsmWriteSpec {
fn from(inner: lancedb::table::LsmWriteSpec) -> Self {
impl From<lancedb::table::MemWalWriteSpec> for MemWalWriteSpec {
fn from(inner: lancedb::table::MemWalWriteSpec) -> Self {
Self { inner }
}
}
@@ -696,6 +696,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>> {
@@ -704,7 +712,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())
}
@@ -712,10 +720,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()))
})
}
@@ -738,7 +743,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 {
@@ -751,7 +756,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> {
@@ -760,7 +765,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 {
@@ -777,17 +782,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())
}
}
@@ -1461,14 +1488,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
@@ -1591,8 +1634,8 @@ impl Table {
if let Some(use_index) = parameters.use_index {
builder.use_index(use_index);
}
if let Some(use_lsm) = parameters.use_lsm {
builder.use_lsm(use_lsm);
if let Some(use_memwal) = parameters.use_memwal {
builder.use_memwal(use_memwal);
}
if let Some(validate_single_shard) = parameters.validate_single_shard {
builder.validate_single_shard(validate_single_shard);
@@ -1620,29 +1663,29 @@ impl Table {
})
}
pub fn set_lsm_write_spec<'a>(
pub fn enable_memwal<'a>(
self_: PyRef<'a, Self>,
spec: LsmWriteSpec,
spec: MemWalWriteSpec,
) -> PyResult<Bound<'a, PyAny>> {
let inner = self_.inner_ref()?.clone();
let native_spec = lancedb::table::LsmWriteSpec::from(spec);
let native_spec = lancedb::table::MemWalWriteSpec::from(spec);
future_into_py(self_.py(), async move {
inner.set_lsm_write_spec(native_spec).await.infer_error()
inner.enable_memwal(native_spec).await.infer_error()
})
}
pub fn unset_lsm_write_spec(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
pub fn disable_memwal(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
inner.unset_lsm_write_spec().await.infer_error()
inner.disable_memwal().await.infer_error()
})
}
pub fn get_lsm_write_spec(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
pub fn get_memwal_spec(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let spec = inner.get_lsm_write_spec().await.infer_error()?;
Ok(spec.map(LsmWriteSpec::from))
let spec = inner.get_memwal_spec().await.infer_error()?;
Ok(spec.map(MemWalWriteSpec::from))
})
}
@@ -1661,55 +1704,54 @@ impl Table {
Ok(())
}
/// Converge the table's LSM write path into its base table.
/// Converge the table's MemWAL write path into its base table.
///
/// Best-effort: with writes flowing, new rows may land after the last
/// pass. Errors if the table stops making progress.
pub fn checkpoint_lsm(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
pub fn checkpoint_memwal(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
inner.checkpoint_lsm().await.infer_error()
inner.checkpoint_memwal().await.infer_error()
})
}
/// Seal every bucket's active memtable into L0.
pub fn flush_lsm(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(
self_.py(),
async move { inner.flush_lsm().await.infer_error() },
)
}
/// Trigger a background L0 → base pass per bucket. Returns once the
/// passes are dispatched, not once they finish — watch `get_lsm_stats`.
pub fn compact_lsm(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
/// Freeze every table shard's active memtable into an SSTable.
pub fn flush_memwal(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
inner.compact_lsm().await.infer_error()
inner.flush_memwal().await.infer_error()
})
}
/// Live LSM state, or `None` when the LSM write path is not enabled.
#[pyo3(signature = (include_generation_rows=false))]
pub fn get_lsm_stats(
/// Trigger a background SSTable compaction pass per table shard. Returns once the
/// passes are dispatched, not once they finish — watch `get_memwal_stats`.
pub fn compact_memwal(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
inner.compact_memwal().await.infer_error()
})
}
/// Live MemWAL state, or `None` when the MemWAL write path is not enabled.
#[pyo3(signature = (include_sstable_rows=false))]
pub fn get_memwal_stats(
self_: PyRef<'_, Self>,
include_generation_rows: bool,
include_sstable_rows: bool,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let stats = inner
.get_lsm_stats(include_generation_rows)
.get_memwal_stats(include_sstable_rows)
.await
.infer_error()?;
Python::attach(|py| stats.map(|s| lsm_stats_to_py(py, &s)).transpose())
Python::attach(|py| stats.map(|s| memwal_stats_to_py(py, &s)).transpose())
})
}
pub fn close_lsm_writers(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
pub fn close_memwal_writers(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
inner.close_lsm_writers().await.infer_error()
inner.close_memwal_writers().await.infer_error()
})
}
@@ -2041,7 +2083,7 @@ pub struct MergeInsertParams {
when_not_matched_by_source_condition_expr: Option<PyExpr>,
timeout: Option<std::time::Duration>,
use_index: Option<bool>,
use_lsm: Option<bool>,
use_memwal: Option<bool>,
validate_single_shard: Option<bool>,
source_collect_threshold_bytes: Option<usize>,
}
+34 -8
View File
@@ -1900,6 +1900,7 @@ clip = [
{ name = "torch" },
]
dev = [
{ name = "maturin" },
{ name = "pre-commit" },
{ name = "pyright" },
{ name = "ruff" },
@@ -1981,6 +1982,7 @@ requires-dist = [
{ name = "ibm-watsonx-ai", marker = "python_full_version >= '3.10' and extra == 'embeddings'", specifier = ">=1.1.2" },
{ name = "instructorembedding", marker = "extra == 'embeddings'", specifier = ">=1.0.1" },
{ name = "lance-namespace", specifier = ">=0.13.0" },
{ name = "maturin", marker = "extra == 'dev'", specifier = ">=1.10" },
{ name = "mkdocs", marker = "extra == 'docs'" },
{ name = "mkdocs-jupyter", marker = "extra == 'docs'" },
{ name = "mkdocs-material", marker = "extra == 'docs'" },
@@ -2005,7 +2007,7 @@ requires-dist = [
{ name = "pyarrow-stubs", marker = "extra == 'tests'", specifier = ">=16.0" },
{ name = "pydantic", specifier = ">=2.7.4,<3" },
{ name = "pylance", marker = "extra == 'pylance'", specifier = ">=5.0.0b5" },
{ name = "pylance", marker = "extra == 'tests'", specifier = "==14.0.0b10" },
{ name = "pylance", marker = "extra == 'tests'", specifier = "==14.0.0b12" },
{ name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.350" },
{ name = "pytest", marker = "extra == 'tests'", specifier = ">=7.0" },
{ name = "pytest-asyncio", marker = "extra == 'tests'", specifier = ">=0.21" },
@@ -2156,6 +2158,30 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/41/09/5b161152e2d90f7b87f781c2e1267494aef9c32498df793f73ad0a0a494a/matplotlib_inline-0.2.2-py3-none-any.whl", hash = "sha256:3c821cf1c209f59fb2d2d64abbf5b23b67bcb2210d663f9918dd851c6da1fcf6", size = 9534, upload-time = "2026-05-08T17:33:32.055Z" },
]
[[package]]
name = "maturin"
version = "1.15.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "tomli", marker = "python_full_version < '3.11'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/b9/c8/22e5e21b2679c9bce6415ca578034ca2cc9316be0642ae21e051a2d5198c/maturin-1.15.0.tar.gz", hash = "sha256:94b26cc8e8aba61a5f2099715fe640e18c5f678e9a500408b38761263954228a", size = 385504, upload-time = "2026-08-24T12:11:22.665Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/14/69/5c01b461044eb1f45ddcce006706eb88110c793cdb11c7ae0b5e08492e94/maturin-1.15.0-py3-none-linux_armv6l.whl", hash = "sha256:6bf6dc62e22d4dcfd5a51244ff0d58975fa4979c48209fe84159617648956d82", size = 10206220, upload-time = "2026-08-24T12:10:53.327Z" },
{ url = "https://files.pythonhosted.org/packages/eb/1f/2b431554e11687cdb1077e0cdadcc118c53f611086b3af00c8545a67c6a5/maturin-1.15.0-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:cd35772633f489841132bc8e71d6fc7f842df30b9c05cd5cdf1ee1ddcb744cc7", size = 19416513, upload-time = "2026-08-24T12:10:56.126Z" },
{ url = "https://files.pythonhosted.org/packages/51/36/e23a21cb34a648b711036b9b2fe1d4f3f4ee24f8db54215d73f1a9a3a3ec/maturin-1.15.0-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:c40b4eae7bf5ef1f4b1af8d623fe4105016f93578fb15b764e741d08ec3b92dd", size = 10014962, upload-time = "2026-08-24T12:10:58.486Z" },
{ url = "https://files.pythonhosted.org/packages/8c/80/33b15cb2d8f30f12c807955e8f2fd775027692904e30ec0784744ce8cd83/maturin-1.15.0-py3-none-manylinux_2_12_i686.manylinux2010_i686.musllinux_1_1_i686.whl", hash = "sha256:7eb066372f541f8eb4909c79c5d9bd0b9e8125980bdf1ec9e8aba23c6c8d6c55", size = 10196223, upload-time = "2026-08-24T12:11:00.696Z" },
{ url = "https://files.pythonhosted.org/packages/fe/91/b495e19e2f5c503b540452b2039115e7b2363867e8c5ad4179eb752fa92c/maturin-1.15.0-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.musllinux_1_1_x86_64.whl", hash = "sha256:653020a63525bb224e5ab0adf02e17a2e08bc86dbea7fc1399c9a56d7529b99e", size = 10541186, upload-time = "2026-08-24T12:11:02.857Z" },
{ url = "https://files.pythonhosted.org/packages/5d/2b/2abff58037188d852b124871b1f0d720e1c2bfb3d4f1b03d87c52cd66488/maturin-1.15.0-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:0ebf9767892725083138e671c34482c660317a2f3d6a29fc0e0f34e9d8c99136", size = 10083468, upload-time = "2026-08-24T12:11:05.012Z" },
{ url = "https://files.pythonhosted.org/packages/3f/07/b7e9f8be99a6627849e81ac7b6694876bce8f50a92995fe17e3cf2610f0a/maturin-1.15.0-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.musllinux_1_1_armv7l.whl", hash = "sha256:7ab7eebffd7b8debca2265985de4eaeb332141276d24b9560b5ad484d4b3add1", size = 10047786, upload-time = "2026-08-24T12:11:07.042Z" },
{ url = "https://files.pythonhosted.org/packages/ca/ab/167e3cb7accee11b507dbe53e0e87aeccb376d44ae66284c96ee4df3a9fd/maturin-1.15.0-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.musllinux_1_1_ppc64le.whl", hash = "sha256:126e12e618b4db42f68c779a56d41f82a390145ba36ac3f621d057eb34f5ad9d", size = 13315332, upload-time = "2026-08-24T12:11:09.433Z" },
{ url = "https://files.pythonhosted.org/packages/14/4d/801379f646cbc6b00998e5289b0630a886be3a4ee4c75b6bc9b87478a7f1/maturin-1.15.0-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:4f9d33e6c3f9615c8caceecbbbd440f8eb25a3ddeb687077682cd5eca2e9ae15", size = 10807183, upload-time = "2026-08-24T12:11:11.73Z" },
{ url = "https://files.pythonhosted.org/packages/89/27/2e612e1cbd1580e9e94d4722c227b5180dca27b32b955a34b79918aa1292/maturin-1.15.0-py3-none-manylinux_2_31_riscv64.musllinux_1_1_riscv64.whl", hash = "sha256:bf29beddd0c6708f112db51d5275fc28b28b9e9c9c5faae387eaef662918b176", size = 10413274, upload-time = "2026-08-24T12:11:14.04Z" },
{ url = "https://files.pythonhosted.org/packages/70/d8/202a7b4d75a51f20f84ec9ce3b7345b12b822207072164e2b1c6ef665125/maturin-1.15.0-py3-none-win32.whl", hash = "sha256:da649988be98e87e009e51b1bf0d301b6a301bc0cecbdd60d40d8ba60748d1ca", size = 8928744, upload-time = "2026-08-24T12:11:16.269Z" },
{ url = "https://files.pythonhosted.org/packages/40/dc/4e90da594986ba78dd3bc8a5921ecdcdb11085b22b02a412caab3b225601/maturin-1.15.0-py3-none-win_amd64.whl", hash = "sha256:552c2be4afd43fe8d5c9f3ec8d4c4756d973b8dcbe94c14084390301f50243e1", size = 10335085, upload-time = "2026-08-24T12:11:18.326Z" },
{ url = "https://files.pythonhosted.org/packages/8b/10/15d4314edf130955edf2dc237aa393a8a7c10f2b9b57b89fa2f61f915659/maturin-1.15.0-py3-none-win_arm64.whl", hash = "sha256:c7dc0c66c78d3debdd9c5aa807e861fbcbf07f3505d34b125df74c03986b0f48", size = 9713795, upload-time = "2026-08-24T12:11:20.83Z" },
]
[[package]]
name = "mdit-py-plugins"
version = "0.6.1"
@@ -3912,7 +3938,7 @@ crypto = [
[[package]]
name = "pylance"
version = "14.0.0b10"
version = "14.0.0b12"
source = { registry = "https://pypi.fury.io/lance-format/" }
dependencies = [
{ name = "lance-namespace" },
@@ -3922,12 +3948,12 @@ dependencies = [
{ name = "pyarrow" },
]
wheels = [
{ url = "https://pypi.fury.io/lance-format/-/ver_4hvE/pylance-14.0.0b10-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:415b6b1b9fe735fbae7c755eb9774ae2fe0feddc90bf80bee24087cbfbfdf8c1" },
{ url = "https://pypi.fury.io/lance-format/-/ver_pM1km/pylance-14.0.0b10-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:87317cad22f4d35fea11eb70aa5c0d71a0e1ea3b96de38f3deca933daa749fff" },
{ url = "https://pypi.fury.io/lance-format/-/ver_24jdyj/pylance-14.0.0b10-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6da554a2fc2197ba254676b12824d8120b48d65937697e0661a61302458a4740" },
{ url = "https://pypi.fury.io/lance-format/-/ver_XclaZ/pylance-14.0.0b10-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:36c80a801675bb9dcbd681fa99d0cdd5ca642dfb3ca08d71ac980278a2f7748b" },
{ url = "https://pypi.fury.io/lance-format/-/ver_L0yh/pylance-14.0.0b10-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:e50ad0b866faefb2fa3ae87328b93fa0d06b2bdd5bb5d0c1a6c74b907cc2fd93" },
{ url = "https://pypi.fury.io/lance-format/-/ver_2gbNr8/pylance-14.0.0b10-cp310-abi3-win_amd64.whl", hash = "sha256:3d86451fe00f71510005bb700737fb1a4f9361554541cbf802f3c5072554d014" },
{ url = "https://pypi.fury.io/lance-format/-/ver_1Ixasc/pylance-14.0.0b12-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:84350a3c5c1eea3398f4c4fad7cfe85b831d831a5e43b5ca2a783dbdb0e11969" },
{ url = "https://pypi.fury.io/lance-format/-/ver_sCEkG/pylance-14.0.0b12-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:76b386c182a6a8e7b86fbef28f84d1dd39e4626ac8681a0cf36f40a45dd7cb87" },
{ url = "https://pypi.fury.io/lance-format/-/ver_1vdT1C/pylance-14.0.0b12-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6aa548c7937fcdfa8f1943a2d59041a42f8b500199880dbb509c7195fa253381" },
{ url = "https://pypi.fury.io/lance-format/-/ver_2fXu1V/pylance-14.0.0b12-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:79f82d6a6c50f37aab338470e3733436e6a27fbb5972d29fcf453ece56d3b3c6" },
{ url = "https://pypi.fury.io/lance-format/-/ver_24UavN/pylance-14.0.0b12-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:de0305999257002785f60169e6926896584d66f1da050bb5afd1d48c6e3d9bb9" },
{ url = "https://pypi.fury.io/lance-format/-/ver_XNynX/pylance-14.0.0b12-cp310-abi3-win_amd64.whl", hash = "sha256:4d0cdb689e898ce71e8e434d00a3b97639cea8fb97da21f1c921543be363268f" },
]
[[package]]