rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-24 04:32:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge main into gatekeeper/fix-4435-1

Browse Source
This commit is contained in:
Gatefixer
2026-10-22 20:40:56 +00:00
parent 4503f0c2bc 88203d2515
commit f9d6589a93
453 changed files with 42568 additions and 8581 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/Cargo.toml
+3 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "lancedb-python"
version = "0.41.0-beta.3"
version = "0.41.0-beta.13"
publish = false
edition.workspace = true
description = "Python bindings for LanceDB"
@@ -22,6 +22,8 @@ bytes.workspace = true
lancedb.workspace = true
datafusion-common.workspace = true
lance-core.workspace = true
# The workspace path dependency includes the patched crate in Maturin sources.
lance-arrow.workspace = true
lance-namespace.workspace = true
lance-namespace-impls.workspace = true
lance-io.workspace = true
python/README.md
+2
View File
@@ -8,6 +8,8 @@ A Python library for [LanceDB](https://github.com/lancedb/lancedb).
pip install lancedb
```
If no wheel is published for your platform, pip builds `lancedb` from the source distribution. That build needs a [Rust toolchain](https://rustup.rs/) and [protoc](https://grpc.io/docs/protoc-installation/) (the Protocol Buffers compiler) on your `PATH`, and can take a while.
### Pre-Haswell x86_64 hosts: `lancedb-compat`
The default `lancedb` wheel targets `x86-64-haswell` (AVX2 + FMA + F16C) for full performance on modern hardware. Pre-Haswell hosts — Intel Sandy Bridge / Ivy Bridge / Westmere; AMD Bulldozer / Piledriver / Steamroller — don't have AVX2 and crash with `Illegal instruction` at `import lancedb`.
python/pyproject.toml
+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"]
python/python/lancedb/__init__.py
+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 RTEWriteSpec
from ._lancedb import tokenize as _tokenize
from .common import URI, sanitize_uri
from urllib.parse import urlparse
@@ -610,7 +610,7 @@ __all__ = [
"SqlQuery",
"LanceDBConnection",
"LanceNamespaceDBConnection",
"LsmWriteSpec",
"RTEWriteSpec",
"RemoteDBConnection",
"Session",
"Table",
python/python/lancedb/_blob.py
+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]]:
python/python/lancedb/_lancedb.pyi
+113 -30
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
@@ -270,6 +270,8 @@ class Connection(object):
filter: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[float] = None,
) -> Table: ...
async def create_materialized_view_async(
self,
@@ -279,6 +281,8 @@ class Connection(object):
filter: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[float] = None,
) -> Job: ...
def list_materialized_views(
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
@@ -311,6 +315,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:
@@ -500,10 +505,17 @@ class Table:
async def delete(self, filter: Union[str, PyExpr]) -> DeleteResult: ...
async def add_columns(self, columns: list[tuple[str, str]]) -> AddColumnsResult: ...
async def add_computed_columns(
self, columns: list[tuple[str, str]]
self,
columns: list[tuple[str, str]],
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[int] = None,
) -> AddColumnsResult: ...
async def add_function_columns(
self, application_json: str, output_name: Optional[str]
self,
application_json: str,
output_name: Optional[str],
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[int] = None,
) -> AddColumnsResult: ...
async def refresh_column(self, column: str) -> RefreshColumnResult: ...
async def refresh_column_async(self, column: str) -> Job: ...
@@ -537,17 +549,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_rte(self, spec: RTEWriteSpec) -> None: ...
async def disable_rte(self) -> None: ...
async def get_rte_spec(self) -> Optional[RTEWriteSpec]: ...
# 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_rte(self) -> None: ...
async def flush_rte(self) -> None: ...
async def compact_rte(self) -> None: ...
async def get_rte_stats(self, include_sstable_rows: bool) -> Optional[dict]: ...
async def close_rte_writers(self) -> None: ...
@property
def tags(self) -> Tags: ...
@property
@@ -563,7 +575,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]
@@ -660,7 +672,7 @@ class Query:
def fast_search(self): ...
def with_row_id(self): ...
def postfilter(self): ...
def use_lsm(self, enable: bool): ...
def use_rte(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]]): ...
@@ -677,7 +689,7 @@ class Query:
class TakeQuery:
def select(self, columns: List[str]): ...
def with_row_id(self): ...
def use_lsm(self, enable: bool): ...
def use_rte(self, enable: bool): ...
async def output_schema(self) -> pa.Schema: ...
async def execute(self) -> RecordBatchStream: ...
async def explain_plan(self, verbose: Optional[bool]) -> str: ...
@@ -696,7 +708,7 @@ class FTSQuery:
def fast_search(self): ...
def with_row_id(self): ...
def postfilter(self): ...
def use_lsm(self, enable: bool): ...
def use_rte(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: ...
@@ -724,9 +736,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_rte(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): ...
@@ -748,10 +761,11 @@ class HybridQuery:
def fast_search(self): ...
def with_row_id(self): ...
def postfilter(self): ...
def use_lsm(self, enable: bool): ...
def use_rte(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): ...
@@ -775,12 +789,13 @@ class PyQueryRequest:
select_source_columns: Optional[Dict[str, str]]
fast_search: Optional[bool]
with_row_id: Optional[bool]
use_lsm: Optional[bool]
use_rte: 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]
@@ -835,22 +850,22 @@ class FtsToken:
text: str
position: int
class LsmWriteSpec:
"""Specification selecting Lance's MemWAL LSM-style write path for
`merge_insert`."""
class RTEWriteSpec:
"""Specification enabling the realtime engine (RTE), which routes
`merge_insert` through Lance's MemWAL LSM-style write path."""
@staticmethod
def bucket(column: str, num_buckets: int) -> "LsmWriteSpec": ...
def bucket(column: str, num_buckets: int) -> "RTEWriteSpec": ...
@staticmethod
def identity(column: str) -> "LsmWriteSpec": ...
def identity(column: str) -> "RTEWriteSpec": ...
@staticmethod
def unsharded() -> "LsmWriteSpec": ...
def with_maintained_indexes(self, indexes: Optional[List[str]]) -> "LsmWriteSpec":
"""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 unsharded() -> "RTEWriteSpec": ...
def with_maintained_indexes(self, indexes: Optional[List[str]]) -> "RTEWriteSpec":
"""Set which indexes the RTE keeps up to date. None is every index on
the table, now and later, skipping any it cannot maintain; 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]) -> "RTEWriteSpec":
"""Return a copy of this spec recording the given default
`ShardWriter` configuration in the MemWAL index."""
...
@@ -864,7 +879,7 @@ class LsmWriteSpec:
def num_buckets(self) -> Optional[int]: ...
@property
def maintained_indexes(self) -> Optional[List[str]]:
"""Indexes the MemWAL keeps up to date, or None for every supported one."""
"""Indexes the RTE keeps up to date, or None for every supported one."""
...
@property
def writer_config_defaults(self) -> Dict[str, str]: ...
@@ -936,6 +951,15 @@ class AsyncPermutationBuilder:
counts: Optional[List[int]] = None,
fixed: Optional[int] = None,
) -> "AsyncPermutationBuilder": ...
def bootstrap_sample(
self,
num_samples: int,
*,
ratios: Optional[List[float]] = None,
counts: Optional[List[int]] = None,
fixed: Optional[int] = None,
seed: Optional[int] = None,
) -> "AsyncPermutationBuilder": ...
def split_calculated(self, calculation: str) -> "AsyncPermutationBuilder": ...
def shuffle(
self, seed: Optional[int], clump_size: Optional[int]
@@ -952,6 +976,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(
@@ -998,3 +1024,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: ...
python/python/lancedb/_wal_hybrid.py
+19 -19
View File
@@ -1,14 +1,14 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""WAL-PK-FUSION: hybrid search's primary-key fallback for MemWAL tables.
"""WAL-PK-FUSION: hybrid search's primary-key fallback for RTE tables.
Hybrid search joins its vector and full-text legs on `_rowid`, which a MemWAL
Hybrid search joins its vector and full-text legs on `_rowid`, which an RTE
table cannot supply: the fresh tier has no stable row id, so the server rejects
`with_row_id` before it plans. When that happens the legs are joined on a
surrogate `_rowid` built from the table's primary key instead.
This is a stopgap until MemWAL supports `_rowid`, and everything it needs lives
This is a stopgap until RTE supports `_rowid`, and everything it needs lives
here so it can be removed in one pass:
1. Delete this file and `tests/test_wal_hybrid.py`.
@@ -69,21 +69,21 @@ def _pk_columns(schema: pa.Schema) -> List[str]:
# builds it from the same words -- so this substring survives either path and
# the wrapping each adds. It is a coupling to a message rather than a code,
# because the refusal arrives as a generic 400; if the server ever rewords it
# the fallback stops firing and hybrid on a MemWAL table errors visibly again,
# the fallback stops firing and hybrid on an RTE table errors visibly again,
# which is loud rather than silently wrong.
_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 "
"hybrid search on an RTE 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 "
"will exclude un-compacted MemWAL data)"
"the primary key. Set use_rte(False) to read the base table only (results "
"will exclude un-compacted RTE data)"
)
def _is_wal_row_id_refusal(error: BaseException) -> bool:
"""Whether `error` is the server declining `_rowid` because the table is
MemWAL-backed, as opposed to any other failure the query might hit."""
RTE-backed, as opposed to any other failure the query might hit."""
return _WAL_ROW_ID_REFUSAL in str(error)
@@ -112,7 +112,7 @@ def _stamp_surrogate_row_ids(
"""Give each table a `_rowid` column derived from its primary key.
The fusion only ever compares that column for equality — to deduplicate, to
group by in a reranker, and to restore the original scores — so on a MemWAL
group by in a reranker, and to restore the original scores — so on an RTE
table, where the server has no stable row id to hand out, any injective
mapping of the primary key does the same job. Ids are assigned in first-seen
order across `tables`, so passing them in fusion order (vector, then FTS)
@@ -124,7 +124,7 @@ def _stamp_surrogate_row_ids(
"""
if not pk_columns:
raise ValueError(
"hybrid search on a MemWAL table needs an unenforced primary key to "
"hybrid search on an RTE table needs an unenforced primary key to "
"deduplicate on, and this table declares none"
)
@@ -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_rte: 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_rte is False:
return False
# Only an unambiguous True forks to the key: a table type that answers the
# question loosely keeps the `_rowid` path.
@@ -165,7 +165,7 @@ def pk_fusion_learned(table, use_lsm: Optional[bool]) -> bool:
class PkFusion:
"""How the legs of a hybrid query on a MemWAL table are joined: on a
"""How the legs of a hybrid query on an RTE table are joined: on a
surrogate `_rowid` built from the primary key."""
def __init__(self, pk_columns: List[str]):
@@ -190,16 +190,16 @@ class PkFusion:
# has turned `_with_row_id` on — including `return_score="all"`.
if blob_v2_projection_sources(schema, projection):
raise NotImplementedError(
"hybrid search cannot project a blob column on a MemWAL table: "
"hybrid search cannot project a blob column on an RTE 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 "
"(results will exclude un-compacted MemWAL data)"
"stable one. Set use_rte(False) to read the base table only "
"(results will exclude un-compacted RTE data)"
)
return cls(_pk_columns(schema))
def inject(self, columns):
"""`columns` with the key columns it is missing, to project in place of
the `_rowid` a MemWAL table would refuse."""
the `_rowid` an RTE table would refuse."""
columns, self.injected = _with_pk_columns(columns, self.pk_columns)
return columns
@@ -226,7 +226,7 @@ class PkFusion:
def _check_refusal(error: Exception, *, fused: bool, caller_requested_row_id: bool):
"""Re-raise `error` unless it is a MemWAL `_rowid` refusal worth retrying on
"""Re-raise `error` unless it is an RTE `_rowid` refusal worth retrying on
the primary key."""
if fused or not _is_wal_row_id_refusal(error):
raise error
@@ -248,7 +248,7 @@ def with_pk_fallback(
`_rowid`.
`_rowid` is asked for first rather than paying a round trip up front to
find out: only a MemWAL table refuses, and the refusal has to be handled
find out: only an RTE table refuses, and the refusal has to be handled
anyway -- a write spec installed elsewhere can arrive between any two
queries. It costs nothing server-side: it is raised before the query is
planned. `fused` reports whether the attempt already used the key.
python/python/lancedb/authz.py
+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))
python/python/lancedb/catalog.py
+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."""
python/python/lancedb/db.py
+48 -3
View File
@@ -5,6 +5,7 @@
from __future__ import annotations
from abc import abstractmethod
import asyncio
from datetime import timedelta
from pathlib import Path
import sys
@@ -607,6 +608,8 @@ class DBConnection(EnforceOverrides):
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[float] = None,
) -> MaterializedView:
"""Define a materialized view named ``name`` over the table ``source``.
@@ -631,6 +634,13 @@ class DBConnection(EnforceOverrides):
Cap the view at this many rows, in materialization order.
with_no_data: bool, default False
Skip the initial refresh and leave the backing table empty.
auto_refresh_enabled: bool, optional
Enable or disable automatic refresh on Cloud and Enterprise.
When omitted, the server determines whether automatic refresh is enabled.
Embedded databases do not support automatic refresh.
auto_refresh_interval_seconds: float, optional
Positive, finite refresh interval in seconds on Cloud and Enterprise.
When omitted, the server determines when to refresh.
Returns
-------
@@ -649,6 +659,8 @@ class DBConnection(EnforceOverrides):
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[float] = None,
) -> Job[None]:
"""Submit materialized-view creation and return its job.
@@ -1288,6 +1300,10 @@ class LanceDBConnection(DBConnection):
self.storage_options = storage_options
self._manifest_enabled = manifest_enabled
self._namespace_client_properties = namespace_client_properties
# The native connection does not expose the Python Session wrapper.
# Keep the object supplied by the caller so this property stays
# synchronous for debugger inspection.
self._session = session
if _inner is not None:
self._conn = _inner
# Native-derived wrappers resolve this in their async reconstruction
@@ -1354,7 +1370,12 @@ class LanceDBConnection(DBConnection):
@property
def session(self) -> Optional[Session]:
return self._conn.session
"""Return the session supplied when this connection was opened.
None when the caller did not pass a session, including connections
reconstructed with :meth:`from_inner`.
"""
return self._session
@property
def uri(self) -> str:
@@ -1706,6 +1727,8 @@ class LanceDBConnection(DBConnection):
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[float] = None,
) -> MaterializedView:
"""Define a materialized view named ``name`` over the table ``source``.
See
@@ -1734,6 +1757,8 @@ class LanceDBConnection(DBConnection):
where=where,
limit=limit,
with_no_data=with_no_data,
auto_refresh_enabled=auto_refresh_enabled,
auto_refresh_interval_seconds=auto_refresh_interval_seconds,
)
)
return MaterializedView(self.open_table(name))
@@ -1748,6 +1773,8 @@ class LanceDBConnection(DBConnection):
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[float] = None,
) -> Job[None]:
job = LOOP.run(
self._conn.create_materialized_view_async(
@@ -1757,6 +1784,8 @@ class LanceDBConnection(DBConnection):
where=where,
limit=limit,
with_no_data=with_no_data,
auto_refresh_enabled=auto_refresh_enabled,
auto_refresh_interval_seconds=auto_refresh_interval_seconds,
)
)
return Job(job)
@@ -2633,8 +2662,16 @@ class AsyncConnection(object):
if fill_value is None:
fill_value = 0.0
data, schema = sanitize_create_table(
data, schema, metadata, on_bad_vectors, fill_value
# Input preparation may advance a user-provided iterator. Keep that work
# off the background event loop so an iterator can use the synchronous
# LanceDB API without blocking the loop that API needs to make progress.
data, schema = await asyncio.to_thread(
sanitize_create_table,
data,
schema,
metadata,
on_bad_vectors,
fill_value,
)
validate_schema(schema)
@@ -2767,6 +2804,8 @@ class AsyncConnection(object):
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[float] = None,
) -> AsyncMaterializedView:
"""Define a materialized view named ``name`` over the table ``source``.
See
@@ -2782,6 +2821,8 @@ class AsyncConnection(object):
filter=where,
limit=limit,
with_no_data=with_no_data,
auto_refresh_enabled=auto_refresh_enabled,
auto_refresh_interval_seconds=auto_refresh_interval_seconds,
**native_options,
)
return AsyncMaterializedView(AsyncTable(inner))
@@ -2795,6 +2836,8 @@ class AsyncConnection(object):
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[float] = None,
) -> AsyncJob[None]:
"""Submit materialized-view creation and return its job.
@@ -2810,6 +2853,8 @@ class AsyncConnection(object):
filter=where,
limit=limit,
with_no_data=with_no_data,
auto_refresh_enabled=auto_refresh_enabled,
auto_refresh_interval_seconds=auto_refresh_interval_seconds,
**native_options,
)
return AsyncJob(inner)
python/python/lancedb/embeddings/__init__.py
+7 -1
View File
@@ -3,7 +3,12 @@
# ruff: noqa: F401
from .base import EmbeddingFunction, EmbeddingFunctionConfig, TextEmbeddingFunction
from .base import (
EmbeddingFunction,
EmbeddingFunctionConfig,
TextEmbeddingFunction,
UntrustedEmbeddingConfigError,
)
from .bedrock import BedRockText
from .cohere import CohereEmbeddingFunction
from .gemini_text import GeminiText
@@ -29,6 +34,7 @@ __all__ = [
"EmbeddingFunction",
"EmbeddingFunctionConfig",
"TextEmbeddingFunction",
"UntrustedEmbeddingConfigError",
"EmbeddingFunctionRegistry",
"get_registry",
"register",
python/python/lancedb/embeddings/base.py
+107 -7
View File
@@ -3,7 +3,7 @@
from abc import ABC, abstractmethod
import copy
from typing import List, Union
from typing import List, Optional, Tuple, Union
from lancedb.util import add_note
import numpy as np
@@ -13,6 +13,18 @@ from pydantic import BaseModel, Field, PrivateAttr
from .utils import TEXT, retry_with_exponential_backoff
class UntrustedEmbeddingConfigError(PermissionError):
"""
Raised when using an embedding function restored from table metadata whose
settings this process has not approved.
See
[allow_remote_code][lancedb.embeddings.EmbeddingFunctionRegistry.allow_remote_code]
and
[allow_endpoint][lancedb.embeddings.EmbeddingFunctionRegistry.allow_endpoint].
"""
class EmbeddingFunction(BaseModel, ABC):
"""
An ABC for embedding functions.
@@ -26,9 +38,11 @@ class EmbeddingFunction(BaseModel, ABC):
3. ndims() which returns the number of dimensions of the vector column
"""
max_retries: int = (
7 # Setting 0 disables retires. Maybe this should not be enabled by default,
)
# Maximum number of retries for transient errors (rate limits, 5xx responses,
# connection errors, ...), with exponential backoff capped at 60 seconds between
# attempts. Permanent errors (HTTP 4xx other than 408/409/429) are never retried.
# Setting 0 disables retries.
max_retries: int = 7
_ndims: int = PrivateAttr()
_original_args: dict = PrivateAttr()
@@ -37,11 +51,29 @@ class EmbeddingFunction(BaseModel, ABC):
"""
Create an instance of the embedding function
"""
from .registry import EmbeddingFunctionRegistry
resolved_kwargs = cls.__resolveVariables(kwargs)
instance = cls(**resolved_kwargs)
instance._original_args = kwargs
instance._original_args = dict(kwargs)
if (
instance.remote_code_model() is not None
and "trust_remote_code" in cls.model_fields
):
# Record the flag even when left at its default, so a missing flag
# in table metadata only ever means a table written before this.
instance._original_args.setdefault(
"trust_remote_code", instance.trust_remote_code
)
# The caller wrote this config, so this process trusts it. Without this,
# a table created with it could not be reopened in the same process.
EmbeddingFunctionRegistry.get_instance()._trust(instance)
return instance
@classmethod
def _resolve_variables(cls, args: dict) -> dict:
return cls.__resolveVariables(args)
@classmethod
def __resolveVariables(cls, args: dict) -> dict:
"""
@@ -92,6 +124,38 @@ class EmbeddingFunction(BaseModel, ABC):
"""
return []
def remote_code_model(self) -> Optional[Tuple[str, Optional[str]]]:
"""
Return the model whose repository code this function runs when its
`trust_remote_code` field is enabled, as a `(name, revision)` pair.
Return None if the function never runs code from a model repository.
When the function is restored from table metadata, `trust_remote_code`
is only enabled for models allowed with
[allow_remote_code][lancedb.embeddings.EmbeddingFunctionRegistry.allow_remote_code].
"""
return None
@classmethod
def _legacy_trust_remote_code_default(cls) -> bool:
"""
The `trust_remote_code` to use for table metadata that omits it. Only
tables written before the flag was always recorded omit it, so this is
the default from that time.
"""
field = cls.model_fields.get("trust_remote_code")
return bool(field.default) if field is not None else False
def endpoints(self) -> List[str]:
"""
Return the non-default URLs this function sends data or credentials to.
When the function is restored from table metadata, each of these must be
allowed with
[allow_endpoint][lancedb.embeddings.EmbeddingFunctionRegistry.allow_endpoint].
"""
return []
@abstractmethod
def compute_query_embeddings(self, *args, **kwargs) -> list[Union[np.array, None]]:
"""
@@ -177,7 +241,9 @@ class EmbeddingFunction(BaseModel, ABC):
Creates a pydantic Field that can automatically annotate
the source column for this embedding function
"""
return Field(json_schema_extra={"source_column_for": self}, **kwargs)
json_schema_extra = dict(kwargs.pop("json_schema_extra", None) or {})
json_schema_extra["source_column_for"] = self
return Field(json_schema_extra=json_schema_extra, **kwargs)
def VectorField(self, **kwargs):
"""
@@ -188,7 +254,9 @@ class EmbeddingFunction(BaseModel, ABC):
static type checkers. LanceDB will infer the fixed vector dimension from
this embedding function.
"""
return Field(json_schema_extra={"vector_column_for": self}, **kwargs)
json_schema_extra = dict(kwargs.pop("json_schema_extra", None) or {})
json_schema_extra["vector_column_for"] = self
return Field(json_schema_extra=json_schema_extra, **kwargs)
def __eq__(self, __value: object) -> bool:
if not hasattr(__value, "__dict__"):
@@ -233,3 +301,35 @@ class TextEmbeddingFunction(EmbeddingFunction):
) -> list[Union[np.array, None]]:
"""Generate the embeddings for the given texts"""
pass
class UntrustedEmbeddingFunction(EmbeddingFunction):
"""
Stands in for an embedding function restored from table metadata whose
settings this process has not approved.
The table stays usable for anything that doesn't compute embeddings, while
computing embeddings raises
[UntrustedEmbeddingConfigError][lancedb.embeddings.UntrustedEmbeddingConfigError].
The real function is never constructed, since some load their model in
`__init__`.
"""
function_name: str
reasons: Tuple[str, ...]
def _error(self) -> UntrustedEmbeddingConfigError:
details = "\n".join(f" - {reason}" for reason in self.reasons)
return UntrustedEmbeddingConfigError(
f"Embedding function '{self.function_name}' was restored from table "
"metadata with settings this process has not approved:\n" + details
)
def compute_query_embeddings(self, *args, **kwargs):
raise self._error()
def compute_source_embeddings(self, *args, **kwargs):
raise self._error()
def ndims(self) -> int:
raise self._error()
python/python/lancedb/embeddings/jinaai.py
+13 -13
View File
@@ -5,9 +5,10 @@
import os
import io
import base64
from functools import cached_property
from urllib.parse import urlparse
from pathlib import Path
from typing import TYPE_CHECKING, ClassVar, List, Union, Optional, Any, Dict
from typing import TYPE_CHECKING, List, Union, Optional, Any, Dict
import numpy as np
import pyarrow as pa
@@ -51,7 +52,6 @@ class JinaEmbeddings(EmbeddingFunction):
name: str = "jina-clip-v1"
api_key: Optional[str] = None
_session: ClassVar = None
def ndims(self):
# TODO: fix hardcoding
@@ -206,8 +206,7 @@ class JinaEmbeddings(EmbeddingFunction):
texts: list[str] or np.ndarray (of str)
The texts to embed
"""
self._init_client()
resp = JinaEmbeddings._session.post( # type: ignore
resp = self._session.post(
API_URL, json={"input": input, "model": self.name}
).json()
if "data" not in resp:
@@ -220,14 +219,15 @@ class JinaEmbeddings(EmbeddingFunction):
return [result["embedding"] for result in sorted_embeddings]
def _init_client(self):
@cached_property
def _session(self):
import requests
if JinaEmbeddings._session is None:
if self.api_key is None and os.environ.get("JINA_API_KEY") is None:
api_key_not_found_help("jina")
api_key = self.api_key or os.environ.get("JINA_API_KEY")
JinaEmbeddings._session = requests.Session()
JinaEmbeddings._session.headers.update(
{"Authorization": f"Bearer {api_key}", "Accept-Encoding": "identity"}
)
if self.api_key is None and os.environ.get("JINA_API_KEY") is None:
api_key_not_found_help("jina")
api_key = self.api_key or os.environ.get("JINA_API_KEY")
session = requests.Session()
session.headers.update(
{"Authorization": f"Bearer {api_key}", "Accept-Encoding": "identity"}
)
return session
python/python/lancedb/embeddings/ollama.py
+20 -1
View File
@@ -14,6 +14,9 @@ if TYPE_CHECKING:
import ollama
DEFAULT_OLLAMA_HOST = "http://localhost:11434"
@register("ollama")
class OllamaEmbeddings(TextEmbeddingFunction):
"""
@@ -24,11 +27,25 @@ class OllamaEmbeddings(TextEmbeddingFunction):
"""
name: str = "nomic-embed-text"
host: str = "http://localhost:11434"
host: str = DEFAULT_OLLAMA_HOST
options: Optional[dict] = None # type = ollama.Options
keep_alive: Optional[Union[float, str]] = None
ollama_client_kwargs: Optional[dict] = {}
def endpoints(self) -> List[str]:
urls = []
if self.host.rstrip("/") != DEFAULT_OLLAMA_HOST:
urls.append(self.host)
# These kwargs go to the underlying httpx client, where a proxy would
# see every request.
for key in ("proxy", "proxies"):
proxy = (self.ollama_client_kwargs or {}).get(key)
if isinstance(proxy, str):
urls.append(proxy)
elif isinstance(proxy, dict):
urls.extend(v for v in proxy.values() if isinstance(v, str))
return urls
def ndims(self) -> int:
return len(self.generate_embeddings(["foo"])[0])
@@ -53,6 +70,8 @@ class OllamaEmbeddings(TextEmbeddingFunction):
The texts to embed
"""
# TODO retry, rate limit, token limit
if isinstance(texts, np.ndarray):
texts = texts.tolist()
embeddings = self._compute_embedding(texts)
return list(embeddings)
python/python/lancedb/embeddings/openai.py
+21 -7
View File
@@ -58,6 +58,9 @@ class OpenAIEmbeddings(TextEmbeddingFunction):
def sensitive_keys():
return ["api_key"]
def endpoints(self) -> List[str]:
return [self.base_url] if self.base_url else []
@staticmethod
def model_names():
return [
@@ -98,24 +101,35 @@ class OpenAIEmbeddings(TextEmbeddingFunction):
valid_indices.append(idx)
# TODO retry, rate limit, token limit
try:
def embed_batch(texts_batch, indices_batch):
kwargs = {
"input": valid_texts,
"input": texts_batch,
"model": self.name,
}
if self.name != "text-embedding-ada-002":
kwargs["dimensions"] = self.dim
rs = self._openai_client.embeddings.create(**kwargs)
valid_embeddings = {
idx: v.embedding for v, idx in zip(rs.data, valid_indices)
}
return {idx: v.embedding for v, idx in zip(rs.data, indices_batch)}
def embed_one_or_log(idx, text):
# Retry a single rejected input
try:
return embed_batch([text], [idx])
except openai.BadRequestError:
logging.exception("Bad request for input at index %d: %r", idx, text)
return {}
try:
valid_embeddings = embed_batch(valid_texts, valid_indices)
except openai.AuthenticationError:
logging.error("Authentication failed: Invalid API key provided")
raise
except openai.BadRequestError:
logging.exception("Bad request: %s", texts)
return [None] * len(texts)
# Retry each input individually
valid_embeddings = {}
for idx, text in zip(valid_indices, valid_texts):
valid_embeddings.update(embed_one_or_log(idx, text))
except Exception:
logging.exception("OpenAI embeddings error")
raise
python/python/lancedb/embeddings/registry.py
+138 -9
View File
@@ -2,9 +2,15 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import json
from typing import Dict, Optional, Type
from typing import Dict, Optional, Set, Tuple, Type
from .base import EmbeddingFunction, EmbeddingFunctionConfig
from pydantic import TypeAdapter
from .base import (
EmbeddingFunction,
EmbeddingFunctionConfig,
UntrustedEmbeddingFunction,
)
class EmbeddingFunctionRegistry:
@@ -42,6 +48,9 @@ class EmbeddingFunctionRegistry:
def __init__(self):
self._functions = {}
self._variables = {}
# Revision is a str, None for the default branch, or _ANY_REVISION.
self._allowed_remote_code: Set[Tuple[str, object]] = set()
self._allowed_endpoints: Set[str] = set()
def register(self, alias: Optional[str] = None):
"""
@@ -60,11 +69,15 @@ class EmbeddingFunctionRegistry:
def decorator(cls):
if not issubclass(cls, EmbeddingFunction):
raise TypeError("Must be a subclass of EmbeddingFunction")
if cls.__name__ in self._functions:
raise KeyError(f"{cls.__name__} was already registered")
key = alias or cls.__name__
if key in self._functions and (
not alias or self._functions[key] is not cls
):
raise KeyError(f"{key} was already registered")
self._functions[key] = cls
cls.__embedding_function_registry_alias__ = alias
# Store the name the class is registered under, so the table
# metadata written for it can be looked up again when reading.
cls.__embedding_function_registry_alias__ = key
return cls
return decorator
@@ -122,20 +135,129 @@ class EmbeddingFunctionRegistry:
obj["vector_column"]: EmbeddingFunctionConfig(
vector_column=obj["vector_column"],
source_column=obj["source_column"],
function=self.get(obj["name"]).create(**obj["model"]),
function=self._restore(obj["name"], obj["model"]),
)
for obj in raw_list
}
def _restore(self, name: str, args: dict) -> EmbeddingFunction:
"""
Create an embedding function from table metadata. Anyone who can write
the table controls `args`, so settings that run model code or send data
elsewhere must already be trusted by this process.
"""
cls = self.get(name)
resolved = cls._resolve_variables(args)
# Inspect the config without running __init__, which may load a model.
probe = cls.model_construct(**resolved)
reasons = []
model = probe.remote_code_model()
if model is not None:
allowed = self._remote_code_allowed(*model)
explicit = "trust_remote_code" in resolved
if explicit:
requested = TypeAdapter(bool).validate_python(
resolved["trust_remote_code"]
)
else:
requested = cls._legacy_trust_remote_code_default()
# Always set the flag, since a subclass may default it to True.
resolved["trust_remote_code"] = requested and allowed
# A legacy table that omits the flag is downgraded rather than
# blocked: most sentence-transformers models don't need remote code,
# and those tables relied on the old True default without asking.
if requested and not allowed and explicit:
model_name, revision = model
reasons.append(
f"runs code from model '{model_name}' at revision "
f"{revision!r}. If you trust it, call "
f"get_registry().allow_remote_code({model_name!r}, "
f"revision={revision!r})"
)
reasons.extend(
f"sends data to endpoint '{url}'. If you trust it, call "
f"get_registry().allow_endpoint({url!r})"
for url in probe.endpoints()
if _normalize_url(url) not in self._allowed_endpoints
)
if reasons:
func = UntrustedEmbeddingFunction(
function_name=name, reasons=tuple(reasons)
)
else:
func = cls(**resolved)
func._original_args = args
return func
def allow_remote_code(self, model: str, revision: Optional[str] = None) -> None:
"""
Trust the code in a model repository when an embedding function that
loads it is restored from table metadata.
Table metadata cannot enable `trust_remote_code` by itself, since anyone
who can write to the table could then run code in this process. Embedding
functions created in this process with `trust_remote_code=True` are
allowed automatically.
Parameters
----------
model : str
The model name, such as `"nomic-ai/nomic-embed-text-v1.5"`.
revision : Optional[str]
The revision the table's embedding function must pin. Pinning a
commit hash is recommended, since the code at a branch or tag can
change after you review it. If None, any revision is trusted,
including pull request refs that anyone can push to the repository.
"""
self._allowed_remote_code.add(
(model, _ANY_REVISION if revision is None else revision)
)
def allow_endpoint(self, url: str) -> None:
"""
Trust a custom endpoint, such as an OpenAI `base_url` or Ollama `host`,
when an embedding function that uses it is restored from table metadata.
Table metadata cannot pick an endpoint by itself, since it could send
your data and API keys to a server you don't control. Embedding functions
created in this process are allowed to use their endpoints automatically.
Parameters
----------
url : str
The endpoint URL. Trailing slashes are ignored.
"""
self._allowed_endpoints.add(_normalize_url(url))
def _remote_code_allowed(self, model: str, revision: Optional[str]) -> bool:
return (model, _ANY_REVISION) in self._allowed_remote_code or (
model,
revision,
) in self._allowed_remote_code
def _trust(self, func: EmbeddingFunction) -> None:
model = func.remote_code_model()
if model is not None and getattr(func, "trust_remote_code", False):
# Only the revision this code selected, where None means the default
# branch. Other revisions, such as PR refs, may hold other code.
self._allowed_remote_code.add(model)
for url in func.endpoints():
self.allow_endpoint(url)
def function_to_metadata(self, conf: "EmbeddingFunctionConfig"):
"""
Convert the given embedding function and source / vector column configs
into a config dictionary that can be serialized into arrow metadata
"""
func = conf.function
name = getattr(
func, "__embedding_function_registry_alias__", func.__class__.__name__
)
if isinstance(func, UntrustedEmbeddingFunction):
name = func.function_name
else:
name = getattr(
func, "__embedding_function_registry_alias__", func.__class__.__name__
)
json_data = func.safe_model_dump()
return {
"name": name,
@@ -180,6 +302,13 @@ class EmbeddingFunctionRegistry:
return self._variables[name]
_ANY_REVISION = object()
def _normalize_url(url: str) -> str:
return url.rstrip("/")
# Global instance
__REGISTRY__ = EmbeddingFunctionRegistry()
python/python/lancedb/embeddings/sentence_transformers.py
+21 -5
View File
@@ -2,7 +2,7 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
from typing import List, Union
from typing import List, Optional, Tuple, Union
import numpy as np
@@ -27,19 +27,32 @@ class SentenceTransformerEmbeddings(TextEmbeddingFunction):
The device to use for the model
normalize : bool, default True
Whether to normalize the embeddings
trust_remote_code : bool, default True
Whether to trust the remote code
trust_remote_code : bool, default False
Whether to run code from the model's repository. Only enable this for
repositories you trust. When restored from table metadata, this is only
enabled for models allowed with
[allow_remote_code][lancedb.embeddings.EmbeddingFunctionRegistry.allow_remote_code].
revision : Optional[str], default None
The model revision to load, such as a commit hash.
"""
name: str = "all-MiniLM-L6-v2"
device: str = "cpu"
normalize: bool = True
trust_remote_code: bool = True
trust_remote_code: bool = False
revision: Optional[str] = None
def __init__(self, **kwargs):
super().__init__(**kwargs)
self._ndims = None
def remote_code_model(self) -> Optional[Tuple[str, Optional[str]]]:
return (self.name, self.revision)
@classmethod
def _legacy_trust_remote_code_default(cls) -> bool:
return True
@property
def embedding_model(self):
"""
@@ -84,5 +97,8 @@ class SentenceTransformerEmbeddings(TextEmbeddingFunction):
"sentence_transformers", "sentence-transformers"
)
return sentence_transformers.SentenceTransformer(
self.name, device=self.device, trust_remote_code=self.trust_remote_code
self.name,
device=self.device,
trust_remote_code=self.trust_remote_code,
revision=self.revision,
)
python/python/lancedb/embeddings/transformers.py
+14 -4
View File
@@ -3,7 +3,7 @@
from functools import cached_property
from typing import List, Any
from typing import List, Any, Optional, Tuple
import numpy as np
@@ -34,7 +34,11 @@ class TransformersEmbeddingFunction(EmbeddingFunction):
Whether or not to allow for custom models defined on the HuggingFace
Hub in their own modeling files. This option should only be set to True
for repositories you trust and in which you have read the code, as it
will execute code present on the Hub on your local machine.
will execute code present on the Hub on your local machine. When
restored from table metadata, this is only enabled for models allowed with
[allow_remote_code][lancedb.embeddings.EmbeddingFunctionRegistry.allow_remote_code].
revision : Optional[str]
The model revision to load, such as a commit hash. Default is None.
to download package, run :
`pip install transformers`
@@ -45,6 +49,7 @@ class TransformersEmbeddingFunction(EmbeddingFunction):
name: str = "colbert-ir/colbertv2.0"
device: str = "cpu"
trust_remote_code: bool = False
revision: Optional[str] = None
_tokenizer: Any = PrivateAttr()
_model: Any = PrivateAttr()
@@ -52,12 +57,17 @@ class TransformersEmbeddingFunction(EmbeddingFunction):
super().__init__(*args, **kwargs)
self._ndims = None
transformers = attempt_import_or_raise("transformers")
self._tokenizer = transformers.AutoTokenizer.from_pretrained(self.name)
self._tokenizer = transformers.AutoTokenizer.from_pretrained(
self.name, revision=self.revision
)
self._model = transformers.AutoModel.from_pretrained(
self.name, trust_remote_code=self.trust_remote_code
self.name, trust_remote_code=self.trust_remote_code, revision=self.revision
)
self._model.to(self.device)
def remote_code_model(self) -> Optional[Tuple[str, Optional[str]]]:
return (self.name, self.revision)
model_config = ConfigDict(ignored_types=(cached_property,))
def ndims(self):
python/python/lancedb/embeddings/utils.py
+102 -17
View File
@@ -241,15 +241,30 @@ def retry_with_exponential_backoff(
exponential_base: float = 2,
jitter: bool = True,
max_retries: int = 7,
max_delay: float = 60,
):
"""Retry a function with exponential backoff.
Only errors that may succeed when repeated are retried. Errors that carry an
HTTP status code are retried when the status is 408, 409, 429 or >= 500; any
other 4xx status (for example 400, 401, 403, 404 or 422) is a permanent failure
that is raised immediately. Errors without a recognisable status code (for
example connection errors, timeouts, or exceptions raised by a custom embedding
function) are retried.
The delay before the n-th retry is ``initial_delay * exponential_base ** n``,
never more than ``max_delay``. With ``jitter`` enabled, a random amount of up to
half of that delay is removed ("equal jitter"), so the actual wait is always in
the range ``[delay / 2, delay]``.
Args:
func (function): The function to be retried.
initial_delay (float): Initial delay in seconds (default is 1).
exponential_base (float): The base for exponential backoff (default is 2).
jitter (bool): Whether to add jitter to the delay (default is True).
max_retries (int): Maximum number of retries (default is 7).
max_delay (float): Maximum delay between two attempts, in seconds
(default is 60).
Returns:
function: The decorated function.
@@ -265,13 +280,12 @@ def retry_with_exponential_backoff(
try:
return func(*args, **kwargs)
# Currently retrying on all exceptions as there is no way to know the
# format of the error msgs used by different APIs. We'll log the error
# and say that it is assumed that if this portion errors out, it's due
# to rate limit but the user should check the error message to be sure.
# Errors are retried unless they are known to be permanent (see
# `_is_non_retryable_error`). Custom embedding functions can raise
# arbitrary exceptions, so unknown errors are assumed to be transient.
except Exception as e: # noqa: PERF203
# Don't retry on authentication errors (e.g., OpenAI 401)
# These are permanent failures that won't be fixed by retrying
# Don't retry permanent failures (e.g., OpenAI 401 or 404). They
# won't be fixed by retrying and would only delay the real error.
if _is_non_retryable_error(e):
raise
@@ -280,42 +294,107 @@ def retry_with_exponential_backoff(
if num_retries > max_retries:
raise Exception(
f"Maximum number of retries ({max_retries}) exceeded.", e
)
) from e
delay *= exponential_base * (1 + jitter * random.random())
delay = min(delay * exponential_base, max_delay)
sleep_for = delay
if jitter:
sleep_for = delay / 2 + random.random() * delay / 2
logging.warning(
"Error occurred: %s \n Retrying in %s seconds (retry %s of %s) \n",
e,
delay,
sleep_for,
num_retries,
max_retries,
)
time.sleep(delay)
time.sleep(sleep_for)
return wrapper
# HTTP statuses in the 4xx range that are worth retrying: request timeout,
# conflict (the OpenAI SDK treats it as transient) and rate limiting.
_RETRYABLE_CLIENT_STATUSES = (408, 409, 429)
# AWS error codes that are reported with a 4xx status but are transient.
_RETRYABLE_AWS_ERROR_CODES = (
"Throttl",
"TooManyRequests",
"RequestLimitExceeded",
"ProvisionedThroughputExceeded",
"ModelNotReady",
"ModelTimeout",
)
def _get_http_status(error: Exception):
"""Best-effort extraction of an HTTP status code from an exception.
Provider SDKs expose it in different ways: ``status_code`` (openai, cohere),
``http_status`` (voyageai), ``response.status_code`` (requests, httpx),
``code`` (google-api-core, urllib) or ``response["ResponseMetadata"]``
(botocore). Returns None if no status can be found.
"""
candidates = [
getattr(error, "status_code", None),
getattr(error, "http_status", None),
]
response = getattr(error, "response", None)
if isinstance(response, dict):
candidates.append(response.get("ResponseMetadata", {}).get("HTTPStatusCode"))
else:
candidates.append(getattr(response, "status_code", None))
candidates.append(getattr(error, "code", None))
for status in candidates:
if isinstance(status, int) and not isinstance(status, bool):
if 100 <= status <= 599:
return status
return None
def _is_non_retryable_error(error: Exception) -> bool:
"""Check if an error should not be retried.
An error is non-retryable if it is an authentication error or if it carries an
HTTP client-error status (4xx) other than 408, 409 and 429. Errors without a
recognisable status are treated as retryable.
Args:
error: The exception to check
Returns:
True if the error should not be retried, False otherwise
"""
from .base import UntrustedEmbeddingConfigError
if isinstance(error, UntrustedEmbeddingConfigError):
return True
# Check for OpenAI authentication errors
error_type = type(error).__name__
if error_type == "AuthenticationError":
return True
# Check for other common non-retryable HTTP status codes
# 401 Unauthorized, 403 Forbidden
if hasattr(error, "status_code"):
if error.status_code in (401, 403):
return True
# A missing API key or a missing optional dependency fails the same way on
# every attempt
if isinstance(error, (MissingApiKeyError, ImportError)):
return True
return False
status = _get_http_status(error)
if status is None or not 400 <= status < 500:
return False
if status in _RETRYABLE_CLIENT_STATUSES:
return False
# AWS reports throttling with a 400 status; the error code tells them apart.
response = getattr(error, "response", None)
if isinstance(response, dict):
aws_code = str(response.get("Error", {}).get("Code", ""))
if aws_code.startswith(_RETRYABLE_AWS_ERROR_CODES):
return False
return True
def url_retrieve(url: str):
@@ -332,9 +411,15 @@ def url_retrieve(url: str):
raise ConnectionError("could not download {} due to {}".format(url, err))
class MissingApiKeyError(ValueError):
"""Raised when an embedding provider's API key is not set."""
def api_key_not_found_help(provider):
logging.error("Could not find API key for %s", provider)
raise ValueError(f"Please set the {provider.upper()}_API_KEY environment variable.")
raise MissingApiKeyError(
f"Please set the {provider.upper()}_API_KEY environment variable."
)
def is_flash_attn_2_available():
python/python/lancedb/embeddings/voyageai.py
+5 -3
View File
@@ -75,7 +75,7 @@ def transform_input(input_data: Union[str, bytes, Path]):
content = {"type": "text", "text": input_data}
elif isinstance(input_data, PIL_Image.Image):
buffered = BytesIO()
input_data.save(buffered, format="JPEG")
input_data.convert("RGB").save(buffered, format="JPEG")
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
content = {
"type": "image_base64",
@@ -84,7 +84,7 @@ def transform_input(input_data: Union[str, bytes, Path]):
elif isinstance(input_data, bytes):
img = PIL_Image.open(BytesIO(input_data))
buffered = BytesIO()
img.save(buffered, format="JPEG")
img.convert("RGB").save(buffered, format="JPEG")
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
content = {
"type": "image_base64",
@@ -103,7 +103,7 @@ def transform_input(input_data: Union[str, bytes, Path]):
else:
img = PIL_Image.open(input_data)
buffered = BytesIO()
img.save(buffered, format="JPEG")
img.convert("RGB").save(buffered, format="JPEG")
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
content = {
"type": "image_base64",
@@ -145,6 +145,8 @@ def sanitize_text_input(inputs: TEXT) -> List[str]:
"""
if isinstance(inputs, str):
inputs = [inputs]
elif isinstance(inputs, (list, np.ndarray)):
inputs = list(inputs)
elif isinstance(inputs, pa.Array):
inputs = inputs.to_pylist()
elif isinstance(inputs, pa.ChunkedArray):
python/python/lancedb/embeddings/watsonx.py
+3
View File
@@ -87,6 +87,9 @@ class WatsonxEmbeddings(TextEmbeddingFunction):
url: Optional[str] = None
params: Optional[Dict] = None
def endpoints(self) -> List[str]:
return [self.url] if self.url else []
@staticmethod
def sensitive_keys():
return ["api_key"]
python/python/lancedb/functions.py
+2 -5
View File
@@ -697,6 +697,7 @@ _GRAMMAR_PRIMITIVES = (
(pa.string(), "utf8"),
(pa.large_string(), "large_utf8"),
(pa.binary(), "binary"),
(pa.large_binary(), "large_binary"),
(pa.date32(), "date32"),
(pa.date64(), "date64"),
)
@@ -846,11 +847,7 @@ def _exact_blob_storage_type(field: pa.Field) -> dict[str, Any]:
{
"name": child.name,
"nullable": child.nullable,
"type": (
{"type": "large_binary"}
if pa.types.is_large_binary(child.type)
else _exact_arrow_type(child.type)
),
"type": _exact_arrow_type(child.type),
}
for child in storage
],
python/python/lancedb/integrations/pyarrow.py
+23 -1
View File
@@ -52,6 +52,9 @@ class PyarrowScannerAdapter(pa.dataset.Scanner):
raise NotImplementedError("memory_pool not supported")
def count_rows(self):
if isinstance(self.filter, pa.compute.Expression):
# stream the batches so memory does not grow with the matching rows
return sum(batch.num_rows for batch in self.to_reader())
return self.table.count_rows(self.filter)
def from_batches(self, **kwargs):
@@ -83,6 +86,8 @@ class PyarrowScannerAdapter(pa.dataset.Scanner):
return self.to_reader().read_all()
def to_reader(self, *, limit: Optional[int] = None):
if isinstance(self.filter, pa.compute.Expression):
return self._to_reader_with_expression(limit=limit)
query = self.table.search()
# Disable the builtin limit
if limit is None:
@@ -98,13 +103,30 @@ class PyarrowScannerAdapter(pa.dataset.Scanner):
query = query.where(self.filter, prefilter=True)
return query.to_batches(batch_size=self.batch_size)
def _to_reader_with_expression(self, *, limit: Optional[int] = None):
# LanceDB filters are SQL strings, but engines such as DuckDB push down
# pyarrow compute expressions. Scan the table and apply the filter and
# projection with pyarrow, since the filter may reference columns that
# are not part of the projection.
if limit is not None and limit <= 0:
raise ValueError("limit must be positive")
query = self.table.search().limit(self.table.count_rows())
scanner = pa.dataset.Scanner.from_batches(
query.to_batches(batch_size=self.batch_size),
columns=self.columns,
filter=self.filter,
)
if limit is None:
return scanner.to_reader()
return scanner.head(limit).to_reader()
class PyarrowDatasetAdapter(pa.dataset.Dataset):
def __init__(self, table: Table):
self.table = table
def count_rows(self, filter: Optional[Filter] = None):
return self.table.count_rows(filter)
return self.scanner(filter=filter).count_rows()
def get_fragments(self, filter: Optional[Filter] = None):
raise NotImplementedError
python/python/lancedb/merge.py
+8 -8
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_rte = None
self._validate_single_shard = None
def when_matched_update_all(
@@ -115,22 +115,22 @@ class LanceMergeInsertBuilder(object):
self._use_index = use_index
return self
def use_lsm(self, enable: bool) -> LanceMergeInsertBuilder:
def use_rte(self, enable: bool) -> LanceMergeInsertBuilder:
"""
Control MemWAL routing for this merge.
Control RTE 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 an RTE 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 RTE routing and errors if the table has no RTE
write spec. ``False`` forces the standard write path even when a spec
is set.
"""
self._use_lsm = enable
self._use_rte = enable
return self
def validate_single_shard(
@@ -139,13 +139,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 an RTE 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 an RTE write spec.
Parameters
----------
python/python/lancedb/namespace.py
+53
View File
@@ -11,6 +11,7 @@ through a namespace abstraction.
from __future__ import annotations
import sys
import warnings
from typing import Any, Dict, Iterable, List, Optional, Union
from uuid import UUID
@@ -82,6 +83,31 @@ from ._lancedb import Session
_MAX_QUERY_K = 2**31 - 1
def _create_table_storage_options(
storage_options: Optional[Dict[str, str]],
data_storage_version: Optional[str],
enable_v2_manifest_paths: Optional[bool],
) -> Optional[Dict[str, str]]:
if data_storage_version is None and enable_v2_manifest_paths is None:
return storage_options
options = dict(storage_options or {})
for name, value in (
("data_storage_version", data_storage_version),
("enable_v2_manifest_paths", enable_v2_manifest_paths),
):
if value is not None:
warnings.warn(
f"setting {name} directly on create_table is deprecated. "
"Use storage_options instead.",
DeprecationWarning,
stacklevel=3,
)
options[f"new_table_{name}"] = (
str(value).lower() if isinstance(value, bool) else value
)
return options
def _query_to_namespace_request(
table_id: List[str],
query: "Query",
@@ -105,6 +131,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:
@@ -582,6 +613,9 @@ class LanceNamespaceDBConnection(DBConnection):
if mode.lower() not in ["create", "overwrite"]:
raise ValueError("mode must be either 'create' or 'overwrite'")
validate_table_name(name)
storage_options = _create_table_storage_options(
storage_options, data_storage_version, enable_v2_manifest_paths
)
async_table = LOOP.run(
self._inner.create_table(
name,
@@ -660,6 +694,8 @@ class LanceNamespaceDBConnection(DBConnection):
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[float] = None,
) -> "MaterializedView":
"""Define a materialized view over a table in the root namespace.
See
@@ -675,6 +711,8 @@ class LanceNamespaceDBConnection(DBConnection):
where=where,
limit=limit,
with_no_data=with_no_data,
auto_refresh_enabled=auto_refresh_enabled,
auto_refresh_interval_seconds=auto_refresh_interval_seconds,
)
).name
)
@@ -690,6 +728,8 @@ class LanceNamespaceDBConnection(DBConnection):
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[float] = None,
) -> Job[None]:
job = LOOP.run(
self._inner.create_materialized_view_async(
@@ -699,6 +739,8 @@ class LanceNamespaceDBConnection(DBConnection):
where=where,
limit=limit,
with_no_data=with_no_data,
auto_refresh_enabled=auto_refresh_enabled,
auto_refresh_interval_seconds=auto_refresh_interval_seconds,
)
)
return Job(job)
@@ -1217,6 +1259,9 @@ class AsyncLanceNamespaceDBConnection:
if mode.lower() not in ["create", "overwrite"]:
raise ValueError("mode must be either 'create' or 'overwrite'")
validate_table_name(name)
storage_options = _create_table_storage_options(
storage_options, data_storage_version, enable_v2_manifest_paths
)
table = await self._inner.create_table(
name,
data,
@@ -1283,6 +1328,8 @@ class AsyncLanceNamespaceDBConnection:
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[float] = None,
) -> "AsyncMaterializedView":
"""Define a materialized view over a table in the root namespace."""
view = await self._inner.create_materialized_view(
@@ -1292,6 +1339,8 @@ class AsyncLanceNamespaceDBConnection:
where=where,
limit=limit,
with_no_data=with_no_data,
auto_refresh_enabled=auto_refresh_enabled,
auto_refresh_interval_seconds=auto_refresh_interval_seconds,
)
# Reopen through the namespace so the view's table carries the
# namespace client and pushdown configuration a bare inner table lacks.
@@ -1306,6 +1355,8 @@ class AsyncLanceNamespaceDBConnection:
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[float] = None,
) -> AsyncJob[None]:
"""Submit materialized-view creation and return its job."""
return await self._inner.create_materialized_view_async(
@@ -1315,6 +1366,8 @@ class AsyncLanceNamespaceDBConnection:
where=where,
limit=limit,
with_no_data=with_no_data,
auto_refresh_enabled=auto_refresh_enabled,
auto_refresh_interval_seconds=auto_refresh_interval_seconds,
)
async def open_materialized_view(self, name: str) -> "AsyncMaterializedView":
python/python/lancedb/permutation.py
+50 -4
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 an RTE write spec are rejected: unflushed rows have no row id.
"""
self._async = async_permutation_builder(table)
@@ -169,6 +169,50 @@ class PermutationBuilder:
)
return self
def bootstrap_sample(
self,
num_samples: int,
*,
ratios: Optional[list[float]] = None,
counts: Optional[list[int]] = None,
fixed: Optional[int] = None,
seed: Optional[int] = None,
split_names: Optional[list[str]] = None,
) -> "PermutationBuilder":
"""
Sample rows uniformly **with replacement**, independently per split.
Unlike the other split strategies, the number of output rows is
``num_samples``, not the number of (optionally filtered) input rows: the
same row can appear zero, one, or many times, and ``num_samples`` may be
larger or smaller than the input row count. Useful for bootstrap
resampling, oversampling rare classes, or any training workload that
wants a fixed-size stream of random draws rather than a single
deterministic pass over every row. Combine with :meth:`filter` to
restrict the eligible population.
One of ratios, counts, or fixed must be provided, with the same meaning
as in :meth:`split_random`, except they divide ``num_samples`` across
splits rather than the input row count: if fixed is 4, each split gets
``num_samples // 4`` samples.
Each split draws its own sample using a seed derived independently from
``(seed, split_index)``, so a split's content does not change if the samples are
later redistributed across a different number of splits.
The optional split_names can be provided to name the splits. If not provided,
the splits can only be referenced by their index.
"""
self._async.bootstrap_sample(
num_samples,
ratios=ratios,
counts=counts,
fixed=fixed,
seed=seed,
split_names=split_names,
)
return self
def split_calculated(
self, calculation: str, split_names: Optional[list[str]] = None
) -> "PermutationBuilder":
@@ -865,9 +909,11 @@ class Permutation:
f"Cannot rename column `{old_name}` to `{new_name}` because a column "
"with that name already exists"
)
new_selection = self.selection.copy()
new_selection[new_name] = new_selection[old_name]
del new_selection[old_name]
# Rebuild the mapping so the renamed column keeps its position
new_selection = {
(new_name if name == old_name else name): value
for name, value in self.selection.items()
}
return self._with_selection(new_selection)
def rename_columns(self, column_map: dict[str, str]) -> "Permutation":
python/python/lancedb/pydantic.py
+35 -4
View File
@@ -127,7 +127,11 @@ def Vector(
core_schema.list_schema(
min_length=dim,
max_length=dim,
items_schema=core_schema.float_schema(),
items_schema=(
core_schema.int_schema()
if pa.types.is_integer(value_type)
else core_schema.float_schema()
),
),
)
@@ -208,7 +212,11 @@ def MultiVector(
items_schema=core_schema.list_schema(
min_length=dim,
max_length=dim,
items_schema=core_schema.float_schema(),
items_schema=(
core_schema.int_schema()
if pa.types.is_integer(value_type)
else core_schema.float_schema()
),
),
),
)
@@ -260,7 +268,9 @@ def _py_type_to_arrow_type(py_type: Type[Any], field: FieldInfo) -> pa.DataType:
def _pydantic_model_to_fields(model: pydantic.BaseModel) -> List[pa.Field]:
return [
_pydantic_to_field(name, field) for name, field in model.model_fields.items()
_pydantic_to_field(name, field)
for name, field in model.model_fields.items()
if field.exclude is not True
]
@@ -493,6 +503,9 @@ class LanceModel(pydantic.BaseModel):
vec_and_function = []
for name, field_info in cls.safe_get_fields().items():
if field_info.exclude is True:
continue
func = get_extras(field_info, "vector_column_for")
if func is not None:
vec_and_function.append([name, func])
@@ -500,6 +513,8 @@ class LanceModel(pydantic.BaseModel):
configs = []
for vec, func in vec_and_function:
for source, field_info in cls.safe_get_fields().items():
if field_info.exclude is True:
continue
src_func = get_extras(field_info, "source_column_for")
if src_func is func:
# note we can't use == here since the function is a pydantic
@@ -524,4 +539,20 @@ def get_extras(field_info: FieldInfo, key: str) -> Any:
def model_to_dict(model: pydantic.BaseModel) -> dict[str, Any]:
"""Convert a Pydantic model to a dictionary."""
return model.model_dump()
return _enum_members_to_values(model.model_dump())
def _enum_members_to_values(value: Any) -> Any:
# model_dump() keeps Enum members, and Arrow cannot convert a member of a
# plain Enum (one without a str or int mixin). Write its value instead,
# which is what the Arrow type of the field is derived from.
if isinstance(value, Enum):
return value.value
if isinstance(value, dict):
return {k: _enum_members_to_values(v) for k, v in value.items()}
if isinstance(value, (list, tuple)) and value:
# Skip lists of plain scalars, such as vectors, without walking them.
if isinstance(value[0], (int, float, str, bytes)):
return value
return [_enum_members_to_values(v) for v in value]
return value
python/python/lancedb/query.py
+549 -167
View File
File diff suppressed because it is too large Load Diff
python/python/lancedb/remote/__init__.py
+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
python/python/lancedb/remote/db.py
+9
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),
@@ -722,6 +723,8 @@ class RemoteDBConnection(DBConnection):
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[float] = None,
) -> MaterializedView:
from .table import RemoteTable
@@ -733,6 +736,8 @@ class RemoteDBConnection(DBConnection):
where=where,
limit=limit,
with_no_data=with_no_data,
auto_refresh_enabled=auto_refresh_enabled,
auto_refresh_interval_seconds=auto_refresh_interval_seconds,
)
)
return MaterializedView(
@@ -754,6 +759,8 @@ class RemoteDBConnection(DBConnection):
where: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[float] = None,
) -> Job[None]:
job = LOOP.run(
self._conn.create_materialized_view_async(
@@ -763,6 +770,8 @@ class RemoteDBConnection(DBConnection):
where=where,
limit=limit,
with_no_data=with_no_data,
auto_refresh_enabled=auto_refresh_enabled,
auto_refresh_interval_seconds=auto_refresh_interval_seconds,
)
)
return Job(job)
python/python/lancedb/remote/header.py
+1 -1
View File
@@ -141,7 +141,7 @@ class OAuthProvider(HeaderProvider):
# Set expiration if provided
expires_in = token_data.get("expires_in")
if expires_in:
if expires_in is not None:
self._token_expires_at = time.time() + expires_in
else:
# Token doesn't expire or expiration unknown
python/python/lancedb/remote/table.py
+51 -31
View File
@@ -31,7 +31,7 @@ from lancedb._lancedb import (
DropColumnsResult,
FtsToken,
IndexConfig,
LsmWriteSpec,
RTEWriteSpec,
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]):
@@ -276,8 +279,12 @@ class RemoteTable(Table):
"""Tokenize a query using the tokenizer configured on an FTS index.
Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are
rebuilt in the client process from index metadata, so the same tokenizer
model files must exist locally.
rebuilt in the client process from index metadata. For remote tables,
this means the same tokenizer model files must also exist locally
(under ``LANCE_LANGUAGE_MODEL_HOME``) and must match the files the
server uses. A mismatch cannot be detected: the returned tokens may
differ from what the server's index produces. If the model files are
missing, an error naming the tokenizer is raised.
"""
return LOOP.run(
self._table.tokenize(query, column=column, index_name=index_name)
@@ -998,8 +1005,17 @@ class RemoteTable(Table):
| None = None,
*,
computed: Dict[str, str] | None = None,
auto_refresh_enabled: bool | None = None,
auto_refresh_interval_seconds: int | None = None,
) -> AddColumnsResult:
return LOOP.run(self._table.add_columns(transforms, computed=computed))
return LOOP.run(
self._table.add_columns(
transforms,
computed=computed,
auto_refresh_enabled=auto_refresh_enabled,
auto_refresh_interval_seconds=auto_refresh_interval_seconds,
)
)
def refresh_column(self, column: str):
return LOOP.run(self._table.refresh_column(column))
@@ -1031,20 +1047,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_rte(self, spec: "RTEWriteSpec") -> None:
"""Install an RTEWriteSpec."""
return LOOP.run(self._table.enable_rte(spec))
def unset_lsm_write_spec(self) -> None:
"""Remove the LsmWriteSpec."""
return LOOP.run(self._table.unset_lsm_write_spec())
def disable_rte(self) -> None:
"""Remove the RTEWriteSpec."""
return LOOP.run(self._table.disable_rte())
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_rte_spec(self) -> Optional["RTEWriteSpec"]:
"""Read the installed RTEWriteSpec, or ``None``."""
return LOOP.run(self._table.get_rte_spec())
# WAL-PK-FUSION: delete both hooks.
def _hybrid_pk_fusion_learned(self) -> bool:
@@ -1054,31 +1070,31 @@ class RemoteTable(Table):
def _note_hybrid_pk_fusion(self) -> None:
self._table._note_hybrid_pk_fusion()
def checkpoint_lsm(self) -> None:
def checkpoint_rte(self) -> None:
"""Synchronous version of
[`AsyncTable.checkpoint_lsm`][lancedb.AsyncTable.checkpoint_lsm]."""
return LOOP.run(self._table.checkpoint_lsm())
[`AsyncTable.checkpoint_rte`][lancedb.AsyncTable.checkpoint_rte]."""
return LOOP.run(self._table.checkpoint_rte())
def flush_lsm(self) -> None:
def flush_rte(self) -> None:
"""Synchronous version of
[`AsyncTable.flush_lsm`][lancedb.AsyncTable.flush_lsm]."""
return LOOP.run(self._table.flush_lsm())
[`AsyncTable.flush_rte`][lancedb.AsyncTable.flush_rte]."""
return LOOP.run(self._table.flush_rte())
def compact_lsm(self) -> None:
def compact_rte(self) -> None:
"""Synchronous version of
[`AsyncTable.compact_lsm`][lancedb.AsyncTable.compact_lsm]."""
return LOOP.run(self._table.compact_lsm())
[`AsyncTable.compact_rte`][lancedb.AsyncTable.compact_rte]."""
return LOOP.run(self._table.compact_rte())
def get_lsm_stats(self, *, include_generation_rows: bool = False) -> Optional[dict]:
def get_rte_stats(self, *, include_sstable_rows: bool = False) -> Optional[dict]:
"""Synchronous version of
[`AsyncTable.get_lsm_stats`][lancedb.AsyncTable.get_lsm_stats]."""
[`AsyncTable.get_rte_stats`][lancedb.AsyncTable.get_rte_stats]."""
return LOOP.run(
self._table.get_lsm_stats(include_generation_rows=include_generation_rows)
self._table.get_rte_stats(include_sstable_rows=include_sstable_rows)
)
def close_lsm_writers(self) -> None:
def close_rte_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_rte_writers())
def drop_index(self, index_name: str):
return LOOP.run(self._table.drop_index(index_name))
@@ -1154,7 +1170,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 +1180,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))
python/python/lancedb/rerankers/base.py
+4 -1
View File
@@ -241,6 +241,7 @@ class Reranker(ABC):
This op is slower than just keeping relevance score but can be useful
for debugging.
"""
combined_schema = pa.unify_schemas([vector_results.schema, fts_results.schema])
# add nulls to fts results for _distance
if "_distance" not in fts_results.column_names:
fts_results = fts_results.append_column(
@@ -265,7 +266,9 @@ class Reranker(ABC):
else:
vector_results_dict[key] = value
combined = pa.Table.from_pylist(list(vector_results_dict.values()))
combined = pa.Table.from_pylist(
list(vector_results_dict.values()), schema=combined_schema
)
return combined
def _keep_relevance_score(self, combined_results: pa.Table):
python/python/lancedb/rerankers/cohere.py
+1 -1
View File
@@ -62,7 +62,7 @@ class CohereReranker(Reranker):
"COHERE_API_KEY not set. Either set it in your environment or \
pass it as `api_key` argument to the CohereReranker."
)
return cohere.Client(os.environ.get("COHERE_API_KEY") or self.api_key)
return cohere.Client(self.api_key or os.environ.get("COHERE_API_KEY"))
def _rerank(self, result_set: pa.Table, query: str):
result_set = self._handle_empty_results(result_set)
python/python/lancedb/rerankers/openai.py
+12 -9
View File
@@ -70,15 +70,18 @@ class OpenaiReranker(Reranker):
],
)
results = json.loads(response.choices[0].message.content)["documents"]
docs, scores = list(
zip(*[(result["content"], result["relevance_score"]) for result in results])
) # tuples
# replace the self.column column with the docs
result_set = result_set.drop(self.column)
result_set = result_set.append_column(
self.column, pa.array(docs, type=pa.string())
)
# add the scores
# The model returns the documents in its own order, so match each score
# back to its row by content.
scores_by_content = {
result["content"]: result["relevance_score"] for result in results
}
missing = sum(doc not in scores_by_content for doc in docs)
if missing:
raise ValueError(
f"The model returned no relevance_score for {missing} of "
f"{len(docs)} documents"
)
scores = [scores_by_content[doc] for doc in docs]
result_set = result_set.append_column(
"_relevance_score", pa.array(scores, type=pa.float32())
)
python/python/lancedb/rerankers/voyageai.py
+2 -2
View File
@@ -32,7 +32,7 @@ class VoyageAIReranker(Reranker):
return_score : str, default "relevance"
options are "relevance" or "all". Only "relevance" is supported for now.
api_key : str, default None
The API key to use. If None, will use the OPENAI_API_KEY environment variable.
The API key to use. If None, will use the VOYAGE_API_KEY environment variable.
truncation : Optional[bool], default None
"""
@@ -64,7 +64,7 @@ class VoyageAIReranker(Reranker):
pass it as `api_key` argument to the VoyageAIReranker."
)
return voyageai.Client(
api_key=os.environ.get("VOYAGE_API_KEY") or self.api_key,
api_key=self.api_key or os.environ.get("VOYAGE_API_KEY"),
)
def _rerank(self, result_set: pa.Table, query: str):
python/python/lancedb/scannable.py
+18 -2
View File
@@ -276,6 +276,14 @@ def _from_iterable(data: Iterator) -> Scannable:
_registered_modules: set[str] = set()
def _hf_batches(data):
if data._indices is None:
yield from data.data.to_batches()
else:
for table in data.with_format("arrow").iter(batch_size=1000):
yield from table.to_batches()
def _register_optional_converters():
"""Register converters for optional dependencies that are already imported."""
@@ -323,7 +331,15 @@ def _register_optional_converters():
def _from_hf_dataset(data: HFDataset) -> Scannable:
table = data.data.table # Access underlying Arrow table
return Scannable(
schema=table.schema, num_rows=len(data), reader=table.to_reader
schema=table.schema,
num_rows=len(data),
reader=(
table.to_reader
if data._indices is None
else lambda: pa.RecordBatchReader.from_batches(
table.schema, _hf_batches(data)
)
),
)
@to_scannable.register(HFDatasetDict)
@@ -335,7 +351,7 @@ def _register_optional_converters():
def gen():
for split_name, dataset in data.items():
for batch in dataset.data.to_batches():
for batch in _hf_batches(dataset):
split_arr = pa.array(
[split_name] * len(batch), type=pa.string()
)
python/python/lancedb/schema.py
+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:
python/python/lancedb/streaming.py
+166 -8
View File
@@ -19,6 +19,7 @@ above.
"""
import ctypes
import hashlib
import heapq
import logging
import os
@@ -39,6 +40,7 @@ from torch.utils.data import DataLoader, IterableDataset, get_worker_info
from .permutation import (
Permutation,
PermutationBuilder,
Transforms,
permutation_builder,
_drop_base_version,
@@ -244,6 +246,20 @@ class StreamingDataset(IterableDataset):
default) rows are shuffled using ``shuffle_seed`` and ``epoch``.
When ``False`` rows are divided into splits sequentially in storage
order, which can be useful for deterministic debugging or evaluation.
Ignored when ``split_strategy`` is provided.
split_strategy:
Escape hatch for any ``PermutationBuilder`` split strategy this class
doesn't otherwise expose a dedicated parameter for -- ``split_hash``,
``split_calculated``, ``bootstrap_sample`` (uniform sampling with
replacement), or any strategy added later. A callable that receives a
fresh ``PermutationBuilder`` for ``table`` and returns a configured one;
it is called and ``.execute()``'d directly, taking over permutation
construction entirely. When provided, ``shuffle``, ``shuffle_seed``,
``epoch``, and ``shuffle_clump_size`` are ignored -- the callable is
solely responsible for how rows are split and ordered. For example,
``split_strategy=lambda b: b.bootstrap_sample(10_000, fixed=num_splits,
seed=42)`` draws 10,000 samples with replacement, divided evenly across
``num_splits`` splits.
shuffle_seed:
Base seed for the random permutation. Combined with ``epoch`` so
each epoch produces a different ordering. Pass ``None`` to generate
@@ -261,10 +277,15 @@ class StreamingDataset(IterableDataset):
storage) at the cost of higher memory usage per split buffer. Defaults
to ``DEFAULT_READ_BATCH_SIZE`` (64).
io_queue_depth:
Number of I/O batches to keep in flight per split. Higher values
overlap storage latency with transform and training compute at the cost
of more memory and threads. Must be greater than zero. Defaults to
``DEFAULT_PREFETCH_BATCHES`` (4).
Number of fetched batches to keep ahead of the transform stage per
split: this bounds the combined count of in-flight fetches and
already-fetched batches still waiting for a transform slot. Higher
values overlap storage latency with transform and training compute
at the cost of more memory and threads, but I/O stops running ahead
once this many batches are buffered, applying backpressure when the
transform stage (or ``transform_queue_depth``) is the bottleneck.
Must be greater than zero. Defaults to ``DEFAULT_PREFETCH_BATCHES``
(4).
columns:
Optional list of column names to read. When set, only those columns
are fetched from storage; all others are omitted. ``None`` (the
@@ -388,6 +409,9 @@ class StreamingDataset(IterableDataset):
*,
num_splits: Optional[int] = None,
shuffle: bool = True,
split_strategy: Optional[
Callable[[PermutationBuilder], PermutationBuilder]
] = None,
shuffle_seed: Optional[int] = 0,
epoch: int = 0,
rank: int = 0,
@@ -578,13 +602,70 @@ class StreamingDataset(IterableDataset):
builder = permutation_builder(table)
if filter is not None:
builder = builder.filter(filter)
if shuffle:
if split_strategy is not None:
# Escape hatch: lets any PermutationBuilder split strategy drive this
# dataset (split_hash, split_calculated, bootstrap_sample, or any future
# strategy) without StreamingDataset needing a bespoke parameter for each
# one. shuffle/shuffle_seed/epoch/shuffle_clump_size are ignored in this
# path -- the caller's split_strategy is solely responsible for ordering.
self._perm_table = split_strategy(builder).execute()
# _rank_splits below assumes split ids are exactly 0..num_splits-1 with no
# gaps (every other path guarantees this by always building with
# fixed=num_splits) -- a split_strategy callable has no such guarantee, and
# a mismatch would otherwise silently drop rows (ids >= num_splits) or leave
# some ranks reading an empty split (missing ids < num_splits).
split_value_counts = (
self._perm_table.to_arrow().column("split_id").value_counts()
)
counts_by_split = dict(
zip(
split_value_counts.field("values").to_pylist(),
split_value_counts.field("counts").to_pylist(),
)
)
expected_splits = set(range(num_splits))
if set(counts_by_split) != expected_splits:
raise ValueError(
f"split_strategy produced split ids {sorted(counts_by_split)}, "
f"but num_splits={num_splits} requires exactly "
f"{{0, ..., {num_splits - 1}}}. Pass num_splits matching the "
"strategy's own split count (e.g. `fixed=num_splits` on "
"bootstrap_sample/split_random/etc.)."
)
# The round-robin iteration below (and in _iter_owned) assumes every
# split has the same row count -- it stops as soon as any one owned
# split is exhausted (see its own comment: "all splits exhaust
# simultaneously"). Every other path guarantees this by always building
# with fixed=num_splits; a split_strategy using ratios/counts/split_hash
# weights has no such guarantee, and unequal sizes would otherwise
# silently drop rows from the larger splits -- or, with world_size > 1
# where each rank owns different splits, end ranks after different row
# counts and desync or hang a distributed run's end-of-epoch collectives.
if len(set(counts_by_split.values())) > 1:
raise ValueError(
f"split_strategy produced unequal split sizes "
f"{counts_by_split} (split id -> row count). Every split must "
"have the same row count. Use `fixed=num_splits` rather than "
"`ratios`/`counts`/weighted `split_hash`, which can produce "
"unequal sizes, or split unevenly-sized groups into separate "
"StreamingDatasets instead."
)
elif shuffle:
perm_seed = shuffle_seed + epoch * _EPOCH_PRIME
self._perm_table = builder.split_random(
fixed=num_splits, seed=perm_seed, clump_size=shuffle_clump_size
).execute()
else:
self._perm_table = builder.split_sequential(fixed=num_splits).execute()
# Only the split_strategy path needs a permutation fingerprint in the
# checkpoint (see `_checkpoint_snapshot`): split_random/split_sequential are
# already fully pinned by shuffle_seed/num_splits, and a fingerprint over raw
# row ids would spuriously reject an otherwise-valid resume after those row
# ids get renumbered by compaction (row ids are not stable by default).
self._uses_split_strategy = split_strategy is not None
# Lazily computed and cached by `_perm_fingerprint` -- `_perm_table` never
# changes after construction, so this is correct to compute once.
self._perm_fingerprint_cache: Optional[str] = None
if self._blocks_per_epoch == "auto":
self._blocks_per_epoch = self._estimate_blocks_per_epoch()
@@ -836,8 +917,17 @@ class StreamingDataset(IterableDataset):
io_pending[i].append((abs_start, io_pool.submit(_io_call, perm_i, indices)))
def _fill_io(i: int) -> None:
# Budget covers both in-flight fetches (io_pending) and batches
# that have already landed but are still waiting for a transform
# slot (raw_batches). Without raw_batches in the count, a slow
# transform stage never applies backpressure to Stage 1: each
# completed fetch is drained into raw_batches and immediately
# replaced by a new one, so I/O keeps racing arbitrarily far
# ahead of the transform stage instead of merely overlapping
# storage latency with it.
while (
len(io_pending[i]) < io_queue_depth and fetch_head[i] < split_sizes[i]
len(io_pending[i]) + len(raw_batches[i]) < io_queue_depth
and fetch_head[i] < split_sizes[i]
):
_submit_io(i)
@@ -1392,11 +1482,51 @@ class StreamingDataset(IterableDataset):
)
return state
@property
def _perm_fingerprint(self) -> str:
"""Cheap identity check for the built permutation, computed once and cached
(``_perm_table`` never changes after construction).
Only meaningful -- and only recorded in a checkpoint -- for the
``split_strategy`` path (see ``_checkpoint_snapshot``): ``num_splits``/
``shuffle_seed`` alone cannot detect a `load_state_dict` resuming saved
per-split progress against a *different* permutation than the one that
produced it, since a ``split_strategy`` callable ignores ``shuffle_seed``
entirely -- e.g. a ``bootstrap_sample`` with no seed (or any changed
callable) can rebuild a completely different permutation on the next run
while ``shuffle_seed`` keeps its same default value and the existing check
passes anyway. ``split_random``/``split_sequential`` are already fully
pinned by ``shuffle_seed``/``num_splits`` and must *not* additionally be
fingerprinted: this hashes physical row ids, which are not stable across
compaction by default, so checking it on those paths would reject an
otherwise-valid resume after the table gets compacted between save and
load even though the seeded permutation itself is unchanged.
"""
if self._perm_fingerprint_cache is None:
arrow_tbl = self._perm_table.to_arrow()
hasher = hashlib.sha256()
hasher.update(str(arrow_tbl.num_rows).encode())
for name in ("row_id", "split_id"):
column = arrow_tbl.column(name).combine_chunks()
hasher.update(column.to_numpy(zero_copy_only=False).tobytes())
self._perm_fingerprint_cache = hasher.hexdigest()
return self._perm_fingerprint_cache
def _checkpoint_snapshot(self) -> dict:
# Only recorded for split_strategy: split_random/split_sequential are
# already fully pinned by shuffle_seed/num_splits below, and row ids are not
# stable across compaction by default, so fingerprinting them unconditionally
# would reject an otherwise-valid resume after the table got compacted.
perm_fingerprint = (
{"perm_fingerprint": self._perm_fingerprint}
if (self._uses_split_strategy)
else {}
)
if self._pack_sequences is not None:
return {
"shuffle_seed": self._shuffle_seed,
"num_splits": self._num_splits,
**perm_fingerprint,
"epoch": self._epoch,
"pack_sequences": self._pack_sequences,
"eos_id": self._eos_id,
@@ -1417,6 +1547,7 @@ class StreamingDataset(IterableDataset):
return {
"shuffle_seed": self._shuffle_seed,
"num_splits": self._num_splits,
**perm_fingerprint,
"epoch": self._epoch,
"samples_consumed_per_split": samples,
"positions_consumed_per_split": positions,
@@ -1546,8 +1677,16 @@ class StreamingDataset(IterableDataset):
Raises ``ValueError`` if ``num_splits`` or ``shuffle_seed`` differ
from the checkpoint, since a different split structure or shuffle order
makes mid-epoch resumption meaningless. Packed checkpoints
pin ``pack_sequences``, ``eos_id``, ``pad_id``,
makes mid-epoch resumption meaningless. Also raises if the checkpoint
was built with ``split_strategy`` and records a permutation fingerprint
that does not match the current dataset's permutation -- this catches a
changed or non-deterministically-reseeded ``split_strategy`` that
``num_splits``/``shuffle_seed`` alone cannot detect, since ``shuffle_seed``
is unused on that path. Checkpoints from ``split_random``/
``split_sequential`` never carry this fingerprint (row ids are not stable
across compaction, so it would reject an otherwise-valid resume), and
continue to rely on ``shuffle_seed``/``num_splits`` alone as before.
Packed checkpoints pin ``pack_sequences``, ``eos_id``, ``pad_id``,
``blocks_per_epoch``, and ``epoch``.
"""
if state["num_splits"] != self._num_splits:
@@ -1560,6 +1699,19 @@ class StreamingDataset(IterableDataset):
f"shuffle_seed mismatch: checkpoint has {state['shuffle_seed']}, "
f"current dataset has {self._shuffle_seed}"
)
checkpoint_fingerprint = state.get("perm_fingerprint")
if (
checkpoint_fingerprint is not None
and checkpoint_fingerprint != self._perm_fingerprint
):
raise ValueError(
"permutation mismatch: this checkpoint was saved against a "
"different permutation (e.g. a split_strategy that resamples "
"differently between runs, such as bootstrap_sample with no "
"seed, or a changed callable) than the current dataset built. "
"Resuming would apply saved per-split positions to the wrong "
"rows."
)
self._consumer_checkpoint_requires_uniform = False
if "pack_buffers" in state or self._pack_sequences is not None:
@@ -1661,6 +1813,10 @@ class StreamingDataset(IterableDataset):
first = states[0]
packed = "pack_buffers" in first
config_keys = ["shuffle_seed", "num_splits", "epoch"]
if "perm_fingerprint" in first:
# Absent only in checkpoints saved before this key existed; skip rather
# than KeyError so old checkpoints keep merging.
config_keys.append("perm_fingerprint")
if packed:
config_keys.extend(
["pack_sequences", "eos_id", "pad_id", "blocks_per_epoch"]
@@ -1670,6 +1826,8 @@ class StreamingDataset(IterableDataset):
if ("pack_buffers" in state) != packed:
raise ValueError("cannot merge packed and unpacked state dicts")
for key in config_keys:
if key not in state:
continue
if state[key] != first[key]:
raise ValueError(
f"{key} mismatch across state dicts: "
python/python/lancedb/table.py
+475 -158
View File
File diff suppressed because it is too large Load Diff
python/python/lancedb/util.py
+51 -9
View File
@@ -10,6 +10,7 @@ import os
import pathlib
import warnings
from datetime import date, datetime
from decimal import Decimal
from functools import singledispatch
from typing import Tuple, Union, Optional, Any, List
from urllib.parse import urlparse
@@ -33,20 +34,20 @@ def safe_import_adlfs():
adlfs = safe_import_adlfs()
def get_uri_scheme(uri: str) -> str:
def get_uri_scheme(uri: Union[str, pathlib.Path]) -> str:
"""
Get the scheme of a URI. If the URI does not have a scheme, assume it is a file URI.
Parameters
----------
uri : str
uri : str or pathlib.Path
The URI to parse.
Returns
-------
str: The scheme of the URI.
"""
parsed = urlparse(uri)
parsed = urlparse(str(uri))
scheme = parsed.scheme
if not scheme:
scheme = "file"
@@ -60,25 +61,26 @@ def get_uri_scheme(uri: str) -> str:
return scheme
def get_uri_location(uri: str) -> str:
def get_uri_location(uri: Union[str, pathlib.Path]) -> str:
"""
Get the location of a URI. If the parameter is not a url, assumes it is just a path
Parameters
----------
uri : str
uri : str or pathlib.Path
The URI to parse.
Returns
-------
str: Location part of the URL, without scheme
"""
parsed = urlparse(uri)
uri_str = str(uri)
parsed = urlparse(uri_str)
if len(parsed.scheme) == 1:
# Windows drive names are parsed as the scheme
# e.g. "c:\path" -> ParseResult(scheme="c", netloc="", path="/path", ...)
# So we add special handling here for schemes that are a single character
return uri
return uri_str
if not parsed.netloc:
return parsed.path
@@ -235,6 +237,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. "
@@ -295,7 +310,16 @@ def _query_vector_dim(query: Optional[Any]) -> Optional[int]:
if query.ndim == 0:
return None
return query.shape[-1]
if isinstance(query, list) and query:
if isinstance(query, (pa.Array, pa.ChunkedArray)):
if len(query) == 0:
return None
if pa.types.is_fixed_size_list(query.type):
return query.type.list_size
if pa.types.is_list(query.type) or pa.types.is_large_list(query.type):
first = query[0].as_py()
return len(first) if first is not None else None
return len(query)
if isinstance(query, (list, tuple)) and query:
first = query[0]
if isinstance(first, (list, tuple, np.ndarray)):
return len(first)
@@ -373,6 +397,14 @@ def _(value: float):
return str(value)
@value_to_sql.register(Decimal)
def _(value: Decimal):
if not value.is_finite():
raise ValueError("Non-finite Decimal values cannot be converted to SQL")
# Quote exact decimal text so the target-column cast avoids float rounding.
return value_to_sql(format(value, "f"))
@value_to_sql.register(bool)
def _(value: bool):
return str(value).upper()
@@ -394,7 +426,8 @@ def _(value: date):
@value_to_sql.register(list)
def _(value: list):
@value_to_sql.register(tuple)
def _(value: Union[list, tuple]):
return "[" + ", ".join(map(value_to_sql, value)) + "]"
@@ -439,6 +472,15 @@ def _(value: np.floating):
return value_to_sql(float(value))
@value_to_sql.register(np.datetime64)
def _(value: np.datetime64):
if np.isnat(value):
return value_to_sql(None)
if np.datetime_data(value.dtype)[0] in ("ps", "fs", "as"):
raise ValueError("SQL timestamps cannot preserve sub-nanosecond precision")
return value_to_sql(str(np.datetime_as_string(value, unit="ns", timezone="UTC")))
def deprecated(func):
"""This is a decorator which can be used to mark functions
as deprecated. It will result in a warning being emitted
python/python/tests/test_authz.py
+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
)
python/python/tests/test_blob.py
+1108 -8
View File
File diff suppressed because it is too large Load Diff
python/python/tests/test_catalog.py
+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 == []
python/python/tests/test_closed_table_query.py
+32
View File
@@ -0,0 +1,32 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import pytest
from lancedb.db import AsyncConnection
@pytest.mark.asyncio
@pytest.mark.parametrize("operation", ["count_rows", "query", "head"])
async def test_closed_table_queries_raise_runtime_error(
mem_db_async: AsyncConnection, operation
):
table = await mem_db_async.create_table("closed_query", data=[{"id": 0}])
assert (await table.query().to_arrow()).to_pylist() == [{"id": 0}]
table.close()
table.close()
assert not table.is_open()
with pytest.raises(RuntimeError, match="Table closed_query is closed"):
if operation == "query":
table.query()
else:
await getattr(table, operation)()
@pytest.mark.asyncio
async def test_open_table_queries(mem_db_async: AsyncConnection):
table = await mem_db_async.create_table("open_query", data=[{"id": 0}])
assert await table.count_rows() == 1
assert (await table.query().to_arrow()).to_pylist() == [{"id": 0}]
assert (await table.head()).to_pylist() == [{"id": 0}]
table.close()
python/python/tests/test_db.py
+3
View File
@@ -86,6 +86,7 @@ def test_sync_debugger_inspection_does_not_use_background_loop(tmp_path, monkeyp
assert members["uri"] == str(tmp_path)
assert members["read_consistency_interval"] is None
assert members["session"] is None
assert repr(db) == f"LanceDBConnection(uri={str(tmp_path)!r})"
assert repr(table) == f"LanceTable(name='test', _conn={db!r})"
@@ -105,6 +106,8 @@ def test_read_consistency_interval_does_not_use_background_loop(tmp_path, monkey
assert db.read_consistency_interval == consistency_interval
assert db_from_inner.read_consistency_interval == consistency_interval
assert db.session is None
assert db_from_inner.session is None
def test_ingest_pd(tmp_path):
python/python/tests/test_duckdb.py
+16
View File
@@ -22,3 +22,19 @@ def test_basic_query(tmp_path):
results = duck_conn.sql("SELECT SUM(y) FROM adapter").fetchall()
assert results[0][0] == 26
def test_query_with_filter(tmp_path):
data = pa.table({"x": [1, 2, 3, 4], "y": [5, 6, 7, 8]})
conn = lancedb.connect(tmp_path)
tbl = conn.create_table("test", data)
adapter = PyarrowDatasetAdapter(tbl) # noqa: F841
duck_conn = duckdb.connect()
results = duck_conn.sql("SELECT SUM(y) FROM adapter WHERE x > 2").fetchall()
assert results[0][0] == 15
results = duck_conn.sql("SELECT y FROM adapter WHERE x > 2 ORDER BY y").fetchall()
assert results == [(7,), (8,)]
python/python/tests/test_elastic_dataloader.py
+329
View File
@@ -2036,6 +2036,92 @@ def test_transform_queue_depth_bounds_cooked_rows(lance_table):
)
def test_raw_queue_depth_bounded_when_transform_is_the_bottleneck(tmp_path):
"""io_queue_depth bounds raw_batches even when transform is the bottleneck.
Regression test for a bug where Stage 1 (I/O) applied no backpressure at
all: _fill_io only capped io_pending (in-flight fetches) at
io_queue_depth, but every completed fetch was drained into raw_batches
unconditionally (_drain_io) and immediately replaced by a new fetch
(_fill_io again), so I/O kept racing arbitrarily far ahead of a slow
transform stage -- raw_batches had no cap of its own.
transform_queue_depth did not help either: its admission check in
_try_submit_tx only counts cooked rows and in-flight transform futures
(`len(cooked[i]) + len(tx_pending[i]) * batch_size`); it never looks at
len(raw_batches[i]), so it cannot throttle Stage 1.
The fix makes _fill_io count both io_pending and raw_batches against the
io_queue_depth budget, so raw_queue_depth should never exceed
io_queue_depth * read_batch_size per split, regardless of how slow the
transform stage is.
I/O is forced to run synchronously (same technique as
test_transform_queue_depth_does_not_admit_at_capacity_minus_one) so the
assertion is deterministic regardless of machine load, rather than
relying on a real-time race between I/O and transform threads. Only the
transform runs on a real background thread (with a short sleep) so the
consumer genuinely blocks waiting for it, letting Stage 1 attempt to run
ahead between blocking waits.
"""
import concurrent.futures as cf
from concurrent.futures import ThreadPoolExecutor
from unittest.mock import patch
import pyarrow as pa
import time
n_rows = 200
read_batch_size = 10
io_queue_depth = 2
transform_queue_depth = 2
max_allowed = io_queue_depth * read_batch_size
db = lancedb.connect(tmp_path)
table = db.create_table("t", pa.table({"id": list(range(n_rows))}))
def slow_transform(batch: pa.RecordBatch) -> list:
time.sleep(0.005)
return batch.column("id").to_pylist()
original_submit = ThreadPoolExecutor.submit
def synchronous_io_submit(self, fn, *args, **kwargs):
if getattr(fn, "__name__", "") == "_io_call":
fut = cf.Future()
try:
fut.set_result(fn(*args, **kwargs))
except Exception as exc:
fut.set_exception(exc)
return fut
return original_submit(self, fn, *args, **kwargs)
with patch.object(ThreadPoolExecutor, "submit", synchronous_io_submit):
ds = StreamingDataset(
table,
num_splits=1,
shuffle_seed=SHUFFLE_SEED,
read_batch_size=read_batch_size,
io_queue_depth=io_queue_depth,
transform_queue_depth=transform_queue_depth,
transform_parallelism=1,
transform=slow_transform,
)
peak_raw = 0
ids = []
for item in ds:
ids.append(item) # slow_transform yields bare ids, not dicts
if ds.raw_queue_depth > peak_raw:
peak_raw = ds.raw_queue_depth
assert peak_raw <= max_allowed, (
f"raw_queue_depth peaked at {peak_raw}, expected <= {max_allowed} "
"(io_queue_depth * read_batch_size)"
)
# The backpressure fix must not come at the cost of dropping or
# duplicating rows.
assert sorted(ids) == list(range(n_rows))
def test_transform_queue_depth_does_not_admit_at_capacity_minus_one(tmp_path):
"""Admission requires a full read_batch_size of free space, not just one slot.
@@ -3173,6 +3259,22 @@ def test_doc_example_checkpoint(lance_table):
)
def test_streaming_dataset_over_table_with_schema_metadata(tmp_path):
"""StreamingDataset builds its permutation at construction; base table schema
metadata (e.g. Hugging Face's `source-column-name-map`) must not break it."""
schema = pa.schema(
[pa.field("id", pa.int64())],
metadata={b"source-column-name-map": b"{}", b"huggingface": b"{}"},
)
table = lancedb.connect(tmp_path).create_table(
"data", pa.table({"id": list(range(NUM_ROWS))}, schema=schema)
)
ds = StreamingDataset(table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED)
assert sorted(s["id"] for s in ds) == list(range(NUM_ROWS))
# ---------------------------------------------------------------------------
# Remote tables (LanceDB Cloud / Enterprise)
# ---------------------------------------------------------------------------
@@ -3195,3 +3297,230 @@ def test_streaming_dataset_over_remote_table():
assert len(server.scans) == 1, "the permutation is built with one row-id scan"
assert server.takes, "rows must be fetched with row-id takes"
assert_server_safe_row_id_requests(server)
# ---------------------------------------------------------------------------
# split_strategy escape hatch
# ---------------------------------------------------------------------------
def test_split_strategy_bootstrap_sample(lance_table):
"""split_strategy lets StreamingDataset use bootstrap_sample directly, producing
more rows than the base table has (sampling with replacement)."""
num_samples = 4 * NUM_ROWS
ds = StreamingDataset(
lance_table,
split_strategy=lambda b: b.bootstrap_sample(
num_samples, fixed=NUM_SPLITS, seed=SHUFFLE_SEED
),
num_splits=NUM_SPLITS,
)
ids = [row["id"] for row in ds]
assert len(ids) == num_samples
# 4x oversampling a small table must produce at least one repeat.
assert len(set(ids)) < len(ids)
assert set(ids).issubset(set(range(NUM_ROWS)))
def test_split_strategy_matches_direct_permutation_builder_use(lance_table):
"""StreamingDataset(split_strategy=...), read single-process (no DataLoader
workers, so this one process owns every split and round-robins across all of
them), must yield exactly what calling the same strategy directly on a
PermutationBuilder would read, interleaved one row per split per cycle.
Reads each split via __getitems__ (explicit indices), not plain iteration:
Permutation's sequential __iter__ returns empty for a permutation with repeated
row_ids (confirmed directly -- len() is correct, __getitems__ returns the right
rows, but `for row in perm` yields nothing). That's a pre-existing limitation of
Permutation's iterator, unrelated to Bootstrap sampling's correctness: it doesn't
affect StreamingDataset or this benchmark, since both only ever read via
explicit-index batched fetches (take_offsets/__getitems__), never the plain
iterator -- flagged separately, not fixed here (out of scope for this change).
"""
from lancedb.permutation import Permutation, permutation_builder
def strategy(builder):
return builder.bootstrap_sample(2 * NUM_ROWS, fixed=NUM_SPLITS, seed=1)
direct_perm_tbl = strategy(permutation_builder(lance_table)).execute()
expected_by_split = []
for split in range(NUM_SPLITS):
perm = Permutation.from_tables(lance_table, direct_perm_tbl, split=split)
rows = perm.__getitems__(list(range(len(perm))))
expected_by_split.append([row["id"] for row in rows])
samples_per_split = len(expected_by_split[0])
expected = [
expected_by_split[split][cycle]
for cycle in range(samples_per_split)
for split in range(NUM_SPLITS)
]
ds = StreamingDataset(lance_table, split_strategy=strategy, num_splits=NUM_SPLITS)
actual = [row["id"] for row in ds]
assert actual == expected
def test_split_strategy_is_reproducible_across_instances(lance_table):
def strategy(builder):
return builder.bootstrap_sample(3 * NUM_ROWS, fixed=NUM_SPLITS, seed=99)
ds1 = StreamingDataset(lance_table, split_strategy=strategy, num_splits=NUM_SPLITS)
ds2 = StreamingDataset(lance_table, split_strategy=strategy, num_splits=NUM_SPLITS)
assert [row["id"] for row in ds1] == [row["id"] for row in ds2]
def test_split_strategy_num_splits_mismatch_raises(lance_table):
"""StreamingDataset(split_strategy=...) must reject a strategy whose split
count doesn't match num_splits. `_rank_splits` assumes split ids are exactly
0..num_splits-1 with no gaps; a mismatch would otherwise silently drop rows
(split ids >= num_splits) or leave some ranks reading an empty split (missing
ids < num_splits)."""
with pytest.raises(ValueError, match="num_splits"):
StreamingDataset(
lance_table,
split_strategy=lambda b: b.bootstrap_sample(
2 * NUM_ROWS, fixed=NUM_SPLITS, seed=SHUFFLE_SEED
),
num_splits=NUM_SPLITS + 1, # the strategy above always builds NUM_SPLITS
)
def test_split_strategy_unequal_split_sizes_raises(lance_table):
"""StreamingDataset(split_strategy=...) must reject a strategy whose splits
all exist (ids 0..num_splits-1) but aren't the same size. The round-robin
iteration (see _iter_owned) assumes every owned split has the same row count
and stops as soon as any one is exhausted -- `ratios`/`counts` (unlike
`fixed`) can produce unequal split sizes, which would otherwise silently
drop rows from the larger splits."""
with pytest.raises(ValueError, match="unequal split sizes"):
StreamingDataset(
lance_table,
split_strategy=lambda b: b.bootstrap_sample(
100, ratios=[0.7, 0.3], seed=SHUFFLE_SEED
),
num_splits=2,
)
def test_split_strategy_permutation_mismatch_on_resume_raises(lance_table):
"""load_state_dict must reject resuming against a different permutation than
the one that produced the checkpoint. num_splits/shuffle_seed alone can't catch
this for split_strategy, since that path ignores shuffle_seed entirely: a
bootstrap_sample reseeded differently resamples completely differently while
shuffle_seed keeps whatever value the caller passed either way, so the
pre-existing checks would pass despite the mismatch."""
def strategy(seed):
return lambda b: b.bootstrap_sample(2 * NUM_ROWS, fixed=NUM_SPLITS, seed=seed)
ds1 = StreamingDataset(
lance_table,
split_strategy=strategy(1),
num_splits=NUM_SPLITS,
shuffle_seed=SHUFFLE_SEED,
)
it = iter(ds1)
for _ in range(NUM_SPLITS): # one full cycle: every split advances uniformly
next(it)
checkpoint = ds1.state_dict()
ds2 = StreamingDataset(
lance_table,
split_strategy=strategy(2), # different seed -> a different permutation
num_splits=NUM_SPLITS,
shuffle_seed=SHUFFLE_SEED, # unchanged -- irrelevant on this path
)
with pytest.raises(ValueError, match="permutation mismatch"):
ds2.load_state_dict(checkpoint)
def test_split_strategy_resume_with_matching_permutation_succeeds(lance_table):
"""The common case: resuming against the same split_strategy (same seed) must
work exactly like the non-split_strategy checkpoint/resume path, continuing
from exactly the same position with no gaps or repeats."""
def strategy(b):
return b.bootstrap_sample(2 * NUM_ROWS, fixed=NUM_SPLITS, seed=7)
ds1 = StreamingDataset(
lance_table,
split_strategy=strategy,
num_splits=NUM_SPLITS,
shuffle_seed=SHUFFLE_SEED,
)
it = iter(ds1)
for _ in range(NUM_SPLITS):
next(it)
checkpoint = ds1.state_dict()
remaining_original = [row["id"] for row in it]
ds2 = StreamingDataset(
lance_table,
split_strategy=strategy,
num_splits=NUM_SPLITS,
shuffle_seed=SHUFFLE_SEED,
)
ds2.load_state_dict(checkpoint) # must not raise
remaining_resumed = [row["id"] for row in ds2]
assert remaining_original == remaining_resumed
def test_default_checkpoint_omits_perm_fingerprint(lance_table):
"""split_random/split_sequential checkpoints must not carry a
perm_fingerprint -- see test_default_checkpoint_resumes_after_compaction for
why checking one on this path would be wrong, not just unnecessary."""
ds = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED)
assert "perm_fingerprint" not in ds.state_dict()
def test_split_strategy_checkpoint_includes_perm_fingerprint(lance_table):
ds = StreamingDataset(
lance_table,
split_strategy=lambda b: b.bootstrap_sample(
2 * NUM_ROWS, fixed=NUM_SPLITS, seed=SHUFFLE_SEED
),
num_splits=NUM_SPLITS,
)
assert "perm_fingerprint" in ds.state_dict()
def test_default_checkpoint_resumes_after_compaction(lance_table):
"""split_random/split_sequential checkpoints must resume normally even after
the table is compacted between save and load. Lance row ids are not stable
across compaction by default (confirmed empirically: a 2-fragment table's
row ids differ before and after `optimize()`), so fingerprinting raw row ids
on this path -- instead of relying on shuffle_seed/num_splits alone, as it
does -- would reject this resume even though the seeded permutation itself
is unchanged."""
lance_table.add(pa.table({"id": [NUM_ROWS]})) # a second fragment to compact
ds = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED)
it = iter(ds)
for _ in range(NUM_SPLITS):
next(it)
checkpoint = ds.state_dict()
assert "perm_fingerprint" not in checkpoint
before_compaction = set(
lance_table.to_lance()
.to_table(columns=[], with_row_id=True)["_rowid"]
.to_pylist()
)
lance_table.optimize()
after_compaction = set(
lance_table.to_lance()
.to_table(columns=[], with_row_id=True)["_rowid"]
.to_pylist()
)
assert before_compaction != after_compaction, (
"test setup didn't actually renumber row ids -- fix the fixture, not "
"the assertion below"
)
ds_resumed = StreamingDataset(
lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED
)
ds_resumed.load_state_dict(checkpoint) # must not raise
python/python/tests/test_embedding_field_metadata.py
+84
View File
@@ -0,0 +1,84 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import copy
import lancedb
import pytest
from lancedb.embeddings.base import TextEmbeddingFunction
from lancedb.embeddings.registry import get_registry
from lancedb.pydantic import LanceModel, Vector
class LocalEmbedding(TextEmbeddingFunction):
def ndims(self):
return 2
def generate_embeddings(self, texts):
return [[1.0, 0.0] if text == "alpha" else [0.0, 1.0] for text in texts]
@pytest.mark.parametrize(
"method,tag",
[("SourceField", "source_column_for"), ("VectorField", "vector_column_for")],
)
@pytest.mark.parametrize("extra", [None, {}, {"x-purpose": "retrieval"}])
def test_embedding_field_preserves_extra_metadata(method, tag, extra):
function = LocalEmbedding.create(max_retries=0)
original = copy.deepcopy(extra)
field = getattr(function, method)(json_schema_extra=extra)
assert extra == original
assert field.json_schema_extra[tag] is function
for key, value in (extra or {}).items():
assert field.json_schema_extra[key] == value
@pytest.mark.parametrize(
"method,tag",
[("SourceField", "source_column_for"), ("VectorField", "vector_column_for")],
)
def test_extra_metadata_cannot_replace_embedding_identity(method, tag):
function = LocalEmbedding.create(max_retries=0)
supplied = {tag: "caller value", "x-purpose": "retrieval"}
field = getattr(function, method)(json_schema_extra=supplied)
assert supplied[tag] == "caller value"
assert field.json_schema_extra[tag] is function
assert field.json_schema_extra["x-purpose"] == "retrieval"
@pytest.mark.parametrize("method", ["SourceField", "VectorField"])
def test_embedding_field_preserves_other_field_arguments(method):
function = LocalEmbedding.create(max_retries=0)
field = getattr(function, method)(title="Custom title", description="Custom field")
assert field.title == "Custom title"
assert field.description == "Custom field"
def test_extra_metadata_preserves_schema_and_local_embedding(tmp_path, monkeypatch):
registry = get_registry()
monkeypatch.setattr(registry, "_functions", registry._functions.copy())
registry.register("test-embedding-field-extra")(LocalEmbedding)
function = LocalEmbedding.create(max_retries=0)
class Row(LanceModel):
text: str = function.SourceField(json_schema_extra={"x-purpose": "source"})
vector: Vector(2) = function.VectorField(
json_schema_extra={"x-purpose": "embedding"}
)
json_schema = Row.model_json_schema()
assert json_schema["properties"]["text"]["x-purpose"] == "source"
assert json_schema["properties"]["vector"]["x-purpose"] == "embedding"
configs = Row.parse_embedding_functions()
assert len(configs) == 1
assert configs[0].function is function
assert configs[0].source_column == "text"
assert configs[0].vector_column == "vector"
db = lancedb.connect(tmp_path)
table = db.create_table(
"rows", schema=Row, data=[{"text": "alpha"}, {"text": "beta"}]
)
assert table.search("alpha").limit(1).to_list()[0]["text"] == "alpha"
reopened = lancedb.connect(tmp_path).open_table("rows")
assert isinstance(reopened.embedding_functions["vector"].function, LocalEmbedding)
assert reopened.search("alpha").limit(1).to_list()[0]["text"] == "alpha"
python/python/tests/test_embedding_registry.py
+88
View File
@@ -0,0 +1,88 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import lancedb
import pytest
from lancedb.embeddings.base import TextEmbeddingFunction
from lancedb.embeddings.registry import EmbeddingFunctionRegistry, get_registry
from lancedb.pydantic import LanceModel, Vector
class FirstEmbedding(TextEmbeddingFunction):
def ndims(self):
return 2
def generate_embeddings(self, texts):
return [[1.0, 0.0] if text == "alpha" else [0.0, 1.0] for text in texts]
class SecondEmbedding(FirstEmbedding):
def generate_embeddings(self, texts):
return [[0.0, 1.0] if text == "alpha" else [1.0, 0.0] for text in texts]
@pytest.mark.parametrize(
"first_alias,second_alias,second_class",
[
("shared", "shared", SecondEmbedding),
(None, "FirstEmbedding", SecondEmbedding),
("SecondEmbedding", None, SecondEmbedding),
(None, None, FirstEmbedding),
("", "", FirstEmbedding),
],
)
def test_duplicate_registration_preserves_original(
first_alias, second_alias, second_class
):
registry = EmbeddingFunctionRegistry()
registry.register(first_alias)(FirstEmbedding)
key = first_alias or "FirstEmbedding"
original_alias = FirstEmbedding.__embedding_function_registry_alias__
second_attributes = dict(second_class.__dict__)
with pytest.raises(KeyError, match=f"{key} was already registered"):
registry.register(second_alias)(second_class)
assert registry.get(key) is FirstEmbedding
assert FirstEmbedding.__embedding_function_registry_alias__ == original_alias
assert dict(second_class.__dict__) == second_attributes
def test_class_name_does_not_collide_with_distinct_alias():
registry = EmbeddingFunctionRegistry()
registry.register()(FirstEmbedding)
same_name = type("FirstEmbedding", (SecondEmbedding,), {})
registry.register("separate")(same_name)
assert registry.get("FirstEmbedding") is FirstEmbedding
assert registry.get("separate") is same_name
def test_same_class_explicit_alias_registration_is_idempotent():
registry = EmbeddingFunctionRegistry()
assert registry.register("shared")(FirstEmbedding) is FirstEmbedding
attributes = dict(FirstEmbedding.__dict__)
assert registry.register("shared")(FirstEmbedding) is FirstEmbedding
assert registry.get("shared") is FirstEmbedding
assert dict(FirstEmbedding.__dict__) == attributes
def test_alias_collision_does_not_change_persisted_table(tmp_path, monkeypatch):
registry = get_registry()
monkeypatch.setattr(registry, "_functions", registry._functions.copy())
alias = "test-alias-collision-persisted-table"
registry.register(alias)(FirstEmbedding)
function = FirstEmbedding.create(max_retries=0)
class Row(LanceModel):
text: str = function.SourceField()
vector: Vector(2) = function.VectorField()
db = lancedb.connect(tmp_path)
table = db.create_table(
"rows", schema=Row, data=[{"text": "alpha"}, {"text": "beta"}]
)
assert table.search("alpha").limit(1).to_list()[0]["text"] == "alpha"
assert registry.register(alias)(FirstEmbedding) is FirstEmbedding
with pytest.raises(KeyError, match=f"{alias} was already registered"):
registry.register(alias)(SecondEmbedding)
reopened = lancedb.connect(tmp_path).open_table("rows")
assert isinstance(reopened.embedding_functions["vector"].function, FirstEmbedding)
assert reopened.search("alpha").limit(1).to_list()[0]["text"] == "alpha"
python/python/tests/test_embeddings.py
+415
View File
@@ -23,6 +23,12 @@ from lancedb.embeddings.utils import retry
from lancedb.pydantic import LanceModel, Vector
@pytest.fixture(autouse=True)
def isolated_embedding_registry(monkeypatch):
registry = get_registry()
monkeypatch.setattr(registry, "_functions", registry._functions.copy())
def mock_embed_func(input_data):
return [np.random.randn(128).tolist() for _ in range(len(input_data))]
@@ -64,6 +70,37 @@ def test_embedding_function(tmp_path):
assert np.allclose(actual, expected)
def test_embedding_function_registered_without_alias(tmp_path):
registry = EmbeddingFunctionRegistry.get_instance()
@registry.register()
class NoAliasEmbeddingFunction(TextEmbeddingFunction):
def ndims(self):
return 2
def generate_embeddings(self, texts):
return [np.array([1.0, 2.0]) for _ in texts]
assert registry.get("NoAliasEmbeddingFunction") is NoAliasEmbeddingFunction
func = NoAliasEmbeddingFunction.create()
class Schema(LanceModel):
text: str = func.SourceField()
vector: Vector(2) = func.VectorField()
db = lancedb.connect(tmp_path)
db.create_table("test", schema=Schema)
# The class name is written to the table metadata, so the function can be
# resolved again when the table is reopened.
table = lancedb.connect(tmp_path).open_table("test")
table.add([{"text": "hello"}])
conf = table.embedding_functions["vector"]
assert isinstance(conf.function, NoAliasEmbeddingFunction)
assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0]]
def test_instructor_ndims_uses_instruction():
instructor = get_registry().get("instructor").create()
model = MagicMock()
@@ -596,6 +633,359 @@ def test_openai_no_retry_on_401(mock_sleep):
assert mock_sleep.call_count == 0
class _HttpStatusError(Exception):
"""Stand-in for provider SDK errors (openai, cohere, ...) with a status code."""
def __init__(self, status_code, message="error"):
super().__init__(message)
self.status_code = status_code
class _ResponseStatusError(Exception):
"""Stand-in for requests/httpx errors, which expose `response.status_code`."""
def __init__(self, status_code):
super().__init__(f"HTTP {status_code}")
self.response = MagicMock(status_code=status_code)
def _run_with_retry(side_effect, **kwargs):
"""Run a function through the retry helper with a mocked `time.sleep`.
Returns (mock_func, sleeps, error).
"""
from lancedb.embeddings.utils import retry_with_exponential_backoff
mock_func = MagicMock(side_effect=side_effect)
sleeps = []
error = None
with patch("time.sleep", side_effect=sleeps.append):
try:
retry_with_exponential_backoff(mock_func, **kwargs)()
except Exception as e:
error = e
return mock_func, sleeps, error
@pytest.mark.parametrize("status", [400, 401, 403, 404, 405, 410, 413, 422])
@pytest.mark.parametrize("make_error", [_HttpStatusError, _ResponseStatusError])
def test_retry_does_not_retry_permanent_http_errors(status, make_error):
err = make_error(status)
mock_func, sleeps, error = _run_with_retry(err, max_retries=7)
assert error is err
assert mock_func.call_count == 1
assert sleeps == []
@pytest.mark.parametrize("status", [408, 409, 429, 500, 502, 503, 504])
@pytest.mark.parametrize("make_error", [_HttpStatusError, _ResponseStatusError])
def test_retry_retries_transient_http_errors(status, make_error):
mock_func, sleeps, error = _run_with_retry(
[make_error(status), make_error(status), "result"], max_retries=7
)
assert error is None
assert mock_func.call_count == 3
assert len(sleeps) == 2
@pytest.mark.parametrize(
"exc",
[
ConnectionError("connection reset"),
TimeoutError("timed out"),
OSError("network is unreachable"),
ValueError("custom embedding function failure"),
RuntimeError("boom"),
Exception("plain exception"),
],
)
def test_retry_retries_errors_without_http_status(exc):
mock_func, sleeps, error = _run_with_retry([exc, exc, "result"], max_retries=7)
assert error is None
assert mock_func.call_count == 3
assert len(sleeps) == 2
def test_retry_error_status_attributes_of_other_types_are_ignored():
# `code` is a string for openai errors ("model_not_found"), and arbitrary
# exceptions may carry unrelated numeric attributes. Neither is an HTTP status.
class WeirdError(Exception):
code = "model_not_found"
status_code = None
class ExitCodeError(Exception):
code = 2
for exc in (WeirdError("a"), ExitCodeError("b")):
mock_func, sleeps, error = _run_with_retry(exc, max_retries=2)
assert mock_func.call_count == 3
assert len(sleeps) == 2
def test_retry_botocore_style_errors():
# botocore ClientError exposes `response["ResponseMetadata"]["HTTPStatusCode"]`.
# Bedrock reports throttling with a 400, which must stay retryable.
class ClientError(Exception):
def __init__(self, status, code):
super().__init__(code)
self.response = {
"ResponseMetadata": {"HTTPStatusCode": status},
"Error": {"Code": code},
}
throttled = ClientError(400, "ThrottlingException")
mock_func, sleeps, error = _run_with_retry([throttled, "result"], max_retries=3)
assert error is None
assert mock_func.call_count == 2
invalid = ClientError(400, "ValidationException")
mock_func, sleeps, error = _run_with_retry(invalid, max_retries=3)
assert error is invalid
assert mock_func.call_count == 1
assert sleeps == []
def test_retry_authentication_error_by_name_not_retried():
class AuthenticationError(Exception):
pass
mock_func, sleeps, error = _run_with_retry(AuthenticationError("bad key"))
assert isinstance(error, AuthenticationError)
assert mock_func.call_count == 1
assert sleeps == []
def test_retry_gives_up_after_max_retries_with_original_error_attached():
original = _HttpStatusError(503, "overloaded")
mock_func, sleeps, error = _run_with_retry(original, max_retries=3)
assert mock_func.call_count == 4
assert len(sleeps) == 3
assert "Maximum number of retries (3) exceeded" in str(error)
assert error.args[1] is original
assert error.__cause__ is original
def test_retry_delay_is_capped_by_max_delay():
# Default settings (7 retries, 60s cap): the old implementation could sleep for
# up to ~47 minutes in a single retry and hours in total.
exc = _HttpStatusError(503)
for seed in range(50):
import random
random.seed(seed)
_, sleeps, _ = _run_with_retry(exc)
assert len(sleeps) == 7
assert max(sleeps) <= 60
assert sum(sleeps) <= 7 * 60
# Custom cap
_, sleeps, _ = _run_with_retry(exc, max_retries=10, max_delay=5)
assert len(sleeps) == 10
assert max(sleeps) <= 5
def test_retry_delay_without_jitter_is_exact_exponential_then_capped():
_, sleeps, _ = _run_with_retry(
_HttpStatusError(500),
initial_delay=1,
exponential_base=2,
jitter=False,
max_retries=8,
max_delay=30,
)
assert sleeps == [2, 4, 8, 16, 30, 30, 30, 30]
def test_retry_jitter_is_bounded():
# Each wait is within [delay / 2, delay], where delay grows by exponential_base
# per retry (no random factor is compounded across retries).
exc = _HttpStatusError(500)
for seed in range(50):
import random
random.seed(seed)
_, sleeps, _ = _run_with_retry(exc, max_retries=8, max_delay=100)
for i, slept in enumerate(sleeps):
delay = min(2 ** (i + 1), 100)
assert delay / 2 <= slept <= delay
# Extremes of the random source map to the ends of the range
with patch("random.random", return_value=0.0):
_, sleeps, _ = _run_with_retry(exc, max_retries=3)
assert sleeps == [1, 2, 4]
with patch("random.random", return_value=1.0):
_, sleeps, _ = _run_with_retry(exc, max_retries=3)
assert sleeps == [2, 4, 8]
def test_retry_large_max_retries_does_not_overflow():
_, sleeps, _ = _run_with_retry(_HttpStatusError(503), max_retries=1500)
assert len(sleeps) == 1500
assert max(sleeps) <= 60
def _register_failing_embedding(name, error):
"""Register an embedding function that always raises `error`.
Returns the class and a list that records every call. The function is
re-created from table metadata, so the call count has to live outside of it.
"""
calls = []
@register(name)
class FailingEmbeddingFunction(TextEmbeddingFunction):
def ndims(self):
return 4
def generate_embeddings(self, texts):
calls.append(len(texts))
raise error
return FailingEmbeddingFunction, calls
@pytest.mark.parametrize("status", [400, 404, 422])
def test_permanent_embedding_error_is_not_retried(tmp_path, status):
"""A permanent provider error (e.g. a misspelled model name) must surface
immediately from both `table.add` and `table.search` instead of being retried
with minutes of backoff."""
error = _HttpStatusError(status, "The model `typo-embed` does not exist")
cls, calls = _register_failing_embedding(f"failing-embedding-{status}", error)
func = cls.create() # default max_retries=7
class Schema(LanceModel):
text: str = func.SourceField()
vector: Vector(func.ndims()) = func.VectorField()
db = lancedb.connect(tmp_path)
table = db.create_table("t", schema=Schema)
with patch("time.sleep") as mock_sleep:
with pytest.raises(Exception, match="The model `typo-embed` does not exist"):
table.add([{"text": "hello"}])
assert len(calls) == 1
with pytest.raises(_HttpStatusError, match="does not exist"):
table.search("hello").limit(1).to_list()
assert len(calls) == 2
mock_sleep.assert_not_called()
def test_transient_embedding_error_is_retried_in_table_add(tmp_path):
cls, calls = _register_failing_embedding(
"failing-embedding-503", _HttpStatusError(503)
)
func = cls.create(max_retries=2)
class Schema(LanceModel):
text: str = func.SourceField()
vector: Vector(func.ndims()) = func.VectorField()
db = lancedb.connect(tmp_path)
table = db.create_table("t", schema=Schema)
with patch("time.sleep") as mock_sleep:
with pytest.raises(Exception, match="Maximum number of retries"):
table.add([{"text": "hello"}])
assert len(calls) == 3
assert mock_sleep.call_count == 2
assert all(call.args[0] <= 60 for call in mock_sleep.call_args_list)
@patch("time.sleep")
@pytest.mark.parametrize(
"error",
[
pytest.param("missing_key", id="missing-api-key"),
pytest.param(ImportError("No module named 'voyageai'"), id="missing-package"),
],
)
def test_no_retry_on_deterministic_errors(mock_sleep, error):
from lancedb.embeddings.utils import (
api_key_not_found_help,
retry_with_exponential_backoff,
)
def func():
if error == "missing_key":
api_key_not_found_help("voyageai")
raise error
mock_func = MagicMock(side_effect=func)
wrapped_func = retry_with_exponential_backoff(mock_func, max_retries=3)
with pytest.raises((ValueError, ImportError)):
wrapped_func()
assert mock_func.call_count == 1
assert mock_sleep.call_count == 0
@patch("time.sleep")
def test_retry_on_wrapped_provider_error(mock_sleep):
"""Providers such as Bedrock wrap throttling in a plain ValueError."""
from lancedb.embeddings.utils import retry_with_exponential_backoff
mock_func = MagicMock(
side_effect=[
ValueError("Error raised by boto3 client: ThrottlingException"),
[0.1],
]
)
wrapped_func = retry_with_exponential_backoff(mock_func, max_retries=3)
assert wrapped_func() == [0.1]
assert mock_func.call_count == 2
assert mock_sleep.call_count == 1
def test_openai_bad_request_partial_batch_recovery():
# Regression test for issue #1677
# Fake exception types, avoids requiring the optional openai/httpx deps
class FakeBadRequestError(Exception):
pass
class FakeAuthenticationError(Exception):
pass
fake_openai_module = MagicMock()
fake_openai_module.BadRequestError = FakeBadRequestError
fake_openai_module.AuthenticationError = FakeAuthenticationError
registry = get_registry()
model = registry.get("openai").create(name="text-embedding-ada-002")
def fake_create(input, model=None, **kwargs):
# Full batch always rejected
if len(input) > 1:
raise FakeBadRequestError("bad request")
if input[0] == "bad text":
raise FakeBadRequestError("bad request")
embedding = MagicMock()
embedding.embedding = [0.1, 0.2, 0.3]
result = MagicMock()
result.data = [embedding]
return result
mock_client = MagicMock()
mock_client.embeddings.create.side_effect = fake_create
with (
patch(
"lancedb.embeddings.openai.attempt_import_or_raise",
return_value=fake_openai_module,
),
patch.object(type(model), "_openai_client", new_callable=lambda: mock_client),
):
embeddings = model.generate_embeddings(["good text", "bad text"])
assert embeddings[0] == [0.1, 0.2, 0.3]
assert embeddings[1] is None
def test_ollama_embeddings_pickle():
"""OllamaEmbeddings must pickle even after the cached client is created."""
registry = get_registry()
@@ -663,3 +1053,28 @@ def test_jina_generate_image_input_dict_local_path(tmp_path):
image_dict = JinaEmbeddings._generate_image_input_dict(image)
assert "image" in image_dict
assert isinstance(image_dict["image"], str) and len(image_dict["image"]) > 0
def test_jina_embeddings_use_their_own_api_key(monkeypatch):
"""Each JinaEmbeddings instance must send its own API key."""
pytest.importorskip("requests")
monkeypatch.delenv("JINA_API_KEY", raising=False)
registry = get_registry()
registry.set_var("jina_key_a", "key-a")
registry.set_var("jina_key_b", "key-b")
func_a = registry.get("jina").create(api_key="$var:jina_key_a")
func_b = registry.get("jina").create(api_key="$var:jina_key_b")
sent = []
def fake_post(session, url, json=None, **kwargs):
sent.append(session.headers["Authorization"])
response = MagicMock()
response.json.return_value = {"data": [{"index": 0, "embedding": [0.0]}]}
return response
with patch("requests.Session.post", fake_post):
func_a.compute_source_embeddings(["hello"])
func_b.compute_source_embeddings(["hello"])
assert sent == ["Bearer key-a", "Bearer key-b"]
python/python/tests/test_embeddings_trust.py
+291
View File
@@ -0,0 +1,291 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""
Restoring an embedding function from table metadata must not let the table grant
itself trust: remote model code and custom endpoints need approval from the
process that opens the table.
"""
import json
from unittest.mock import patch
import lancedb
import pyarrow as pa
import pytest
from lancedb.embeddings import (
EmbeddingFunctionConfig,
OllamaEmbeddings,
OpenAIEmbeddings,
SentenceTransformerEmbeddings,
UntrustedEmbeddingConfigError,
get_registry,
)
from lancedb.embeddings.base import UntrustedEmbeddingFunction
@pytest.fixture(autouse=True)
def fresh_trust(monkeypatch):
registry = get_registry()
monkeypatch.setattr(registry, "_allowed_remote_code", set())
monkeypatch.setattr(registry, "_allowed_endpoints", set())
def metadata_for(name: str, model: dict) -> dict:
entry = {
"name": name,
"model": model,
"source_column": "text",
"vector_column": "vector",
}
return {b"embedding_functions": json.dumps([entry]).encode()}
def restore(name: str, model: dict):
return get_registry().parse_functions(metadata_for(name, model))["vector"].function
def test_sentence_transformers_does_not_trust_remote_code_by_default():
assert SentenceTransformerEmbeddings.create().trust_remote_code is False
@pytest.mark.parametrize("name", ["sentence-transformers", "huggingface", "colbert"])
def test_restored_remote_code_requires_allowlist(name):
model = {"name": "attacker/model", "trust_remote_code": True}
func = restore(name, model)
with pytest.raises(UntrustedEmbeddingConfigError, match="attacker/model"):
func.compute_query_embeddings("hello")
with pytest.raises(UntrustedEmbeddingConfigError):
func.ndims()
def test_blocked_function_does_not_construct_model():
# huggingface loads its model in __init__, so the check has to happen first.
with patch(
"lancedb.embeddings.transformers.attempt_import_or_raise"
) as attempt_import:
restore("huggingface", {"name": "attacker/model", "trust_remote_code": True})
attempt_import.assert_not_called()
def test_allowlisted_model_restores_with_remote_code():
get_registry().allow_remote_code("org/model")
func = restore(
"sentence-transformers", {"name": "org/model", "trust_remote_code": True}
)
assert isinstance(func, SentenceTransformerEmbeddings)
assert func.trust_remote_code is True
def test_allowlist_does_not_cover_other_models():
get_registry().allow_remote_code("org/model")
func = restore(
"sentence-transformers", {"name": "org/other", "trust_remote_code": True}
)
with pytest.raises(UntrustedEmbeddingConfigError, match="org/other"):
func.compute_query_embeddings("hello")
def test_pinned_revision_must_match():
get_registry().allow_remote_code("org/model", revision="abc123")
pinned = restore(
"sentence-transformers",
{"name": "org/model", "revision": "abc123", "trust_remote_code": True},
)
assert pinned.trust_remote_code is True
for model in [
{"name": "org/model", "revision": "def456", "trust_remote_code": True},
{"name": "org/model", "trust_remote_code": True},
]:
func = restore("sentence-transformers", model)
with pytest.raises(UntrustedEmbeddingConfigError, match="org/model"):
func.compute_query_embeddings("hello")
def test_restored_trust_flag_follows_allowlist_when_absent():
# Tables written before the default changed don't record trust_remote_code.
assert restore(
"sentence-transformers", {"name": "org/model"}
).trust_remote_code is (False)
get_registry().allow_remote_code("org/model")
assert restore("sentence-transformers", {"name": "org/model"}).trust_remote_code
def test_created_function_records_trust_flag():
# Otherwise a new table using the default would look like a pre-fix table,
# where a missing flag meant True.
conf = EmbeddingFunctionConfig(
source_column="text",
vector_column="vector",
function=SentenceTransformerEmbeddings.create(name="org/model"),
)
assert get_registry().function_to_metadata(conf)["model"] == {
"name": "org/model",
"trust_remote_code": False,
}
get_registry().allow_remote_code("org/model")
metadata = get_registry().get_table_metadata([conf])
restored = get_registry().parse_functions(metadata)["vector"].function
assert restored.trust_remote_code is False
def test_huggingface_missing_flag_keeps_false_default():
# huggingface always defaulted to False, so approval must not flip it.
get_registry().allow_remote_code("org/model")
with patch("lancedb.embeddings.transformers.attempt_import_or_raise"):
func = restore("huggingface", {"name": "org/model"})
assert func.trust_remote_code is False
def test_restored_explicit_false_stays_false():
get_registry().allow_remote_code("org/model")
func = restore(
"sentence-transformers", {"name": "org/model", "trust_remote_code": False}
)
assert func.trust_remote_code is False
def test_create_trusts_its_own_config(tmp_path):
func = SentenceTransformerEmbeddings.create(
name="org/model", trust_remote_code=True
)
db = lancedb.connect(tmp_path)
schema = pa.schema(
[pa.field("text", pa.string()), pa.field("vector", pa.list_(pa.float32(), 4))]
)
db.create_table(
"t",
schema=schema,
embedding_functions=[
EmbeddingFunctionConfig(
source_column="text", vector_column="vector", function=func
)
],
)
restored = db.open_table("t").embedding_functions["vector"].function
assert isinstance(restored, SentenceTransformerEmbeddings)
assert restored.trust_remote_code is True
def test_create_trusts_only_the_revision_it_selected():
# Anyone can push a PR ref to a Hugging Face repo, so trusting a model's
# default branch must not extend to its other revisions.
SentenceTransformerEmbeddings.create(name="org/model", trust_remote_code=True)
func = restore(
"sentence-transformers",
{"name": "org/model", "revision": "refs/pr/123", "trust_remote_code": True},
)
assert isinstance(func, UntrustedEmbeddingFunction)
same = restore(
"sentence-transformers", {"name": "org/model", "trust_remote_code": True}
)
assert same.trust_remote_code is True
def test_explicit_unpinned_approval_covers_any_revision():
get_registry().allow_remote_code("org/model")
func = restore(
"sentence-transformers",
{"name": "org/model", "revision": "abc123", "trust_remote_code": True},
)
assert func.trust_remote_code is True
def test_missing_flag_does_not_inherit_unapproved_trust_default():
class LegacySentenceTransformer(SentenceTransformerEmbeddings):
trust_remote_code: bool = True
registry = get_registry()
registry.register("legacy-trust-default")(LegacySentenceTransformer)
try:
func = restore("legacy-trust-default", {"name": "attacker/model"})
assert not func.trust_remote_code
finally:
del registry._functions["legacy-trust-default"]
def test_trusting_one_model_does_not_trust_a_swapped_name():
SentenceTransformerEmbeddings.create(name="org/model", trust_remote_code=True)
func = restore(
"sentence-transformers", {"name": "attacker/model", "trust_remote_code": True}
)
with pytest.raises(UntrustedEmbeddingConfigError):
func.compute_query_embeddings("hello")
ATTACKER_URL = "https://attacker.example/v1"
@pytest.mark.parametrize(
"name,model",
[
("openai", {"base_url": ATTACKER_URL}),
("ollama", {"host": ATTACKER_URL}),
("ollama", {"ollama_client_kwargs": {"proxy": ATTACKER_URL}}),
("watsonx", {"url": ATTACKER_URL}),
],
)
def test_restored_custom_endpoint_requires_allowlist(name, model):
func = restore(name, model)
with pytest.raises(UntrustedEmbeddingConfigError, match="attacker.example"):
func.compute_query_embeddings("hello")
get_registry().allow_endpoint(ATTACKER_URL)
assert isinstance(restore(name, model), get_registry().get(name))
def test_default_endpoints_need_no_allowlist():
assert isinstance(restore("openai", {}), OpenAIEmbeddings)
assert isinstance(
restore("ollama", {"host": "http://localhost:11434"}), OllamaEmbeddings
)
def test_endpoint_allowlist_ignores_trailing_slash():
get_registry().allow_endpoint("http://my-host:11434/")
assert isinstance(
restore("ollama", {"host": "http://my-host:11434"}), OllamaEmbeddings
)
def test_create_trusts_its_own_endpoint():
OpenAIEmbeddings.create(base_url="http://my-proxy/v1")
assert isinstance(
restore("openai", {"base_url": "http://my-proxy/v1"}), OpenAIEmbeddings
)
def test_blocked_function_round_trips_metadata():
model = {"name": "attacker/model", "trust_remote_code": True}
func = restore("sentence-transformers", model)
conf = EmbeddingFunctionConfig(
source_column="text", vector_column="vector", function=func
)
entry = get_registry().function_to_metadata(conf)
assert entry["name"] == "sentence-transformers"
assert entry["model"] == model
def test_table_with_blocked_function_still_opens(tmp_path):
db = lancedb.connect(tmp_path)
schema = pa.schema(
[pa.field("text", pa.string()), pa.field("vector", pa.list_(pa.float32(), 2))],
metadata=metadata_for(
"sentence-transformers",
{"name": "attacker/model", "trust_remote_code": True},
),
)
db.create_table("t", schema=schema)
tbl = db.open_table("t")
tbl.add([{"text": "hi", "vector": [1.0, 2.0]}])
assert tbl.search([1.0, 2.0]).limit(1).to_list()[0]["text"] == "hi"
with pytest.raises(UntrustedEmbeddingConfigError, match="allow_remote_code"):
tbl.search("hi").limit(1).to_list()
python/python/tests/test_expr.py
+61
View File
@@ -80,6 +80,67 @@ class TestExprConstruction:
e = lit(d)
assert isinstance(e, Expr)
@pytest.mark.parametrize(
"value, decimal_type",
[
("1E-7", pa.decimal128(12, 7)),
("-1.23E-7", pa.decimal128(12, 9)),
("1E+3", pa.decimal128(12, 0)),
("-1.23E+5", pa.decimal128(12, 0)),
("0E-7", pa.decimal128(12, 7)),
("-0E+3", pa.decimal128(12, 0)),
("1E-38", pa.decimal128(38, 38)),
("1E+37", pa.decimal128(38, 0)),
("1.234567890123456789012345678901234567E+37", pa.decimal128(38, 0)),
("0E+128", pa.decimal128(1, 0)),
],
)
def test_lit_decimal_exponent_query(self, tmp_path, value, decimal_type):
expected = Decimal(value)
table = lancedb.connect(tmp_path).create_table(
"decimals",
pa.table(
{
"id": [1, 2],
# Arrow's Python ingestion rejects 0E+128 even though its
# numeric value is representable. Store canonical zero to
# exercise the expression literal independently of ingestion.
"price": pa.array(
[Decimal(0) if expected.is_zero() else expected, None],
type=decimal_type,
),
}
),
)
predicate = col("price") == lit(expected)
result = table.search().where(predicate).to_arrow()
assert result["id"].to_pylist() == [1]
assert result["price"].to_pylist() == [expected]
# Exercise the SQL serializer and parser as well as the typed query path.
sql_result = table.search().where(predicate.to_sql()).to_arrow()
assert sql_result["id"].to_pylist() == [1]
assert sql_result["price"].to_pylist() == [expected]
@pytest.mark.parametrize(
"value",
[
"NaN",
"sNaN",
"Infinity",
"-Infinity",
"1E-39",
"1E+38",
"1E+129",
"1E+1000000000",
"0E+129",
"12345678901234567890123456789012345678E+1",
"9" * 39,
],
)
def test_lit_decimal_unrepresentable_raises(self, value):
with pytest.raises(ValueError):
lit(Decimal(value))
class TestExprOperators:
def test_eq_operator(self):
python/python/tests/test_first_class_function_slice1.py
+20 -3
View File
@@ -398,8 +398,21 @@ class _FunctionDeclarationInner:
def __init__(self):
self.calls = []
async def add_function_columns(self, application_json, output_name):
self.calls.append((json.loads(application_json), output_name))
async def add_function_columns(
self,
application_json,
output_name,
auto_refresh_enabled,
auto_refresh_interval_seconds,
):
self.calls.append(
(
json.loads(application_json),
output_name,
auto_refresh_enabled,
auto_refresh_interval_seconds,
)
)
return "declared"
@@ -420,14 +433,18 @@ async def test_add_columns_routes_struct_as_one_and_multi_output_binding_atomica
)
assert result == "declared"
assert inner.calls[-1][1] == "features"
assert inner.calls[-1][2:] == (None, None)
bare = application._copy(update={"columns": {}}).rename(
columns={"normalized_text": "search_text"}
)
result = await table.add_columns(bare)
result = await table.add_columns(
bare, auto_refresh_enabled=False, auto_refresh_interval_seconds=300
)
assert result == "declared"
assert inner.calls[-1][1] is None
assert inner.calls[-1][0]["columns"] == {"normalized_text": "search_text"}
assert inner.calls[-1][2:] == (False, 300)
@pytest.mark.asyncio
python/python/tests/test_first_class_function_slice2.py
+22 -1
View File
@@ -622,7 +622,6 @@ def test_canonical_arrow_type_prefers_the_compact_grammar():
for outside in [
pa.timestamp("us"),
pa.decimal128(10, 2),
pa.large_binary(),
pa.binary(4),
pa.duration("s"),
pa.list_(pa.float32(), 0),
@@ -1280,6 +1279,28 @@ def test_explicit_large_utf8_schemas_use_the_canonical_function_name():
}
def test_large_binary_schemas_use_the_canonical_function_name():
payload = pa.field("data", pa.large_binary(), nullable=True)
@udf(
input_schema=pa.schema([payload]),
output_schema=pa.field("result", pa.large_binary(), nullable=False),
)
def preserve(data):
return data
signature = preserve.registration_request.signature
assert signature.inputs[0].arrow_type == "large_binary"
assert signature.output.arrow_type == "large_binary"
cells = pa.large_list(pa.field("item", pa.struct([payload]), nullable=True))
assert json.loads(_canonical_arrow_type(cells))["fields"][0]["type"] == {
"type": "struct",
"fields": [
{"name": "data", "nullable": True, "type": {"type": "large_binary"}}
],
}
def test_nested_struct_output_uses_canonical_exact_json():
token = pa.struct(
[
python/python/tests/test_header_provider.py
+24
View File
@@ -114,6 +114,30 @@ class TestOAuthProvider:
"x-lancedb-credential-type": "oidc",
}
def test_zero_lifetime_token_is_refreshed(self):
"""A token that expires immediately must not be cached indefinitely."""
call_count = 0
def fetcher():
nonlocal call_count
call_count += 1
return {
"access_token": f"token{call_count}",
"expires_in": 0,
}
provider = OAuthProvider(fetcher, refresh_buffer_seconds=0)
assert provider.get_headers() == {
"Authorization": "Bearer token1",
"x-lancedb-credential-type": "oidc",
}
assert provider.get_headers() == {
"Authorization": "Bearer token2",
"x-lancedb-credential-type": "oidc",
}
assert call_count == 2
def test_missing_access_token(self):
"""Test error handling when access_token is missing."""
python/python/tests/test_huggingface.py
+103
View File
@@ -129,3 +129,106 @@ def test_generator(tmp_path: Path):
assert len(tbl) == 2
assert tbl.schema == ds.features.arrow_schema
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("operation", ["create", "add"])
@pytest.mark.parametrize(
"selection",
["normal", "contiguous", "selected", "duplicates", "filtered", "shuffled", "empty"],
)
async def test_write_selected_hf_dataset(tmp_path, asynchronous, operation, selection):
source = datasets.Dataset.from_dict({"id": list(range(5)), "text": list("abcde")})
data = {
"normal": source,
"contiguous": source.select(range(1, 4)),
"selected": source.select([4, 1]),
"duplicates": source.select([2, 2, 0]),
"filtered": source.filter(lambda row: row["id"] % 2 == 1),
"shuffled": source.shuffle(seed=42),
"empty": source.select([]),
}[selection]
expected = data.to_list()
if asynchronous:
db = await lancedb.connect_async(tmp_path)
if operation == "create":
table = await db.create_table("selected", data)
else:
table = await db.create_table("selected", schema=data.features.arrow_schema)
await table.add(data)
actual = await table.to_arrow()
else:
db = lancedb.connect(tmp_path)
if operation == "create":
table = db.create_table("selected", data)
else:
table = db.create_table("selected", schema=data.features.arrow_schema)
table.add(data)
actual = table.to_arrow()
assert actual.to_pylist() == expected
assert actual.schema.equals(data.features.arrow_schema, check_metadata=True)
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("operation", ["create", "add"])
async def test_write_selected_hf_splits(tmp_path, asynchronous, operation):
source = datasets.Dataset.from_dict({"id": list(range(5)), "text": list("abcde")})
data = datasets.DatasetDict(
{
"train": source.select([4, 1]),
"test": source.select([3, 0, 2]),
"empty": source.select([]),
}
)
expected = [
dict(row, split=split)
for split, dataset in data.items()
for row in dataset.to_list()
]
schema = source.features.arrow_schema.append(pa.field("split", pa.string()))
if asynchronous:
db = await lancedb.connect_async(tmp_path)
if operation == "create":
table = await db.create_table("splits", data)
else:
table = await db.create_table("splits", schema=schema)
await table.add(data)
actual = await table.to_arrow()
else:
db = lancedb.connect(tmp_path)
if operation == "create":
table = db.create_table("splits", data)
else:
table = db.create_table("splits", schema=schema)
table.add(data)
actual = table.to_arrow()
assert actual.to_pylist() == expected
assert actual.schema.equals(schema, check_metadata=True)
@pytest.mark.parametrize("format_type", [None, "numpy", "pandas", "arrow"])
@pytest.mark.parametrize("empty", [False, True])
def test_selected_hf_reader_is_rescannable(format_type, empty):
from lancedb.scannable import _register_optional_converters, to_scannable
source = datasets.Dataset.from_dict(
{"id": list(range(2105)), "text": [str(i) for i in range(2105)]}
)
indices = [] if empty else list(reversed(range(2105))) + [1, 1]
data = source.select(indices).with_format(format_type, columns=["id"])
original_format = data.format
expected = source.select(indices).to_list()
_register_optional_converters()
scannable = to_scannable(data)
assert scannable.rescannable
assert scannable.num_rows == len(expected)
for _ in range(2):
batches = list(scannable.reader())
if not empty:
assert all(batch.num_rows <= 1000 for batch in batches)
actual = pa.Table.from_batches(batches, schema=scannable.schema)
assert actual.to_pylist() == expected
assert actual.schema.equals(data.features.arrow_schema, check_metadata=True)
assert data.format == original_format
python/python/tests/test_hybrid_query.py
+12
View File
@@ -536,3 +536,15 @@ def test_normalize_scores():
assert pc.equal(result, expected), (
f"Expected {expected} but got {result} for invert={invert}"
)
def test_hybrid_query_rejects_query_with_vector(sync_table: Table):
# Pins the full message: the space between "hybrid search." and "But" is
# load-bearing (see the _validate_query literal concatenation).
with pytest.raises(ValueError) as excinfo:
sync_table.search("dog", query_type="hybrid").vector([0.5, -0.5]).to_arrow()
assert str(excinfo.value) == (
"You can either provide a string query in search() method "
"or set `vector()` and `text()` explicitly for hybrid search. "
"But not both."
)
python/python/tests/test_index.py
+22
View File
@@ -430,6 +430,28 @@ async def test_create_ivfrq_index(some_table: AsyncTable):
assert indices[0].name == "vector_idx"
@pytest.mark.asyncio
async def test_read_ivfrq_index_created_with_lance(db_async, some_table: AsyncTable):
pytest.importorskip("lance")
dataset = await some_table.to_lance()
dataset.create_index(
"vector",
index_type="IVF_RQ",
metric="cosine",
num_partitions=1,
)
# Open a fresh LanceDB handle so the index is discovered from disk.
table = await db_async.open_table("some_table")
indices = await table.list_indices()
assert len(indices) == 1
assert indices[0].index_type == "IvfRq"
assert indices[0].columns == ["vector"]
results = await table.query().nearest_to(list(range(DIM))).limit(8).to_arrow()
assert len(results) == 8
@pytest.mark.asyncio
async def test_create_hnswpq_index(some_table: AsyncTable):
await some_table.create_index("vector", config=HnswPq(num_partitions=10))
python/python/tests/test_materialized_views.py
+103 -2
View File
@@ -129,8 +129,21 @@ def test_sync_remote_create_uses_public_async_connection():
where=None,
limit=None,
with_no_data=False,
auto_refresh_enabled=None,
auto_refresh_interval_seconds=None,
):
calls.append((name, source, select, where, limit, with_no_data))
calls.append(
(
name,
source,
select,
where,
limit,
with_no_data,
auto_refresh_enabled,
auto_refresh_interval_seconds,
)
)
return StubAsyncMaterializedView()
async def drop_materialized_view(self, name, *, namespace_path=None):
@@ -148,9 +161,11 @@ def test_sync_remote_create_uses_public_async_connection():
where="age >= 18",
limit=10,
with_no_data=True,
auto_refresh_enabled=True,
auto_refresh_interval_seconds=0.5,
)
assert view.name == "adults"
assert calls == [("adults", "people", ["name"], "age >= 18", 10, True)]
assert calls == [("adults", "people", ["name"], "age >= 18", 10, True, True, 0.5)]
db.drop_materialized_view("adults", namespace_path=["analytics"])
assert calls[-1] == ("drop", "adults", ["analytics"])
@@ -550,3 +565,89 @@ def test_local_materialized_view_listing_pages_and_resume(tmp_path):
assert sorted([first, *resumed]) == ["a_view", "z_view"]
assert resumed.page_token() is None
assert list(resumed) == []
@pytest.mark.parametrize(
"options",
[
{},
{"auto_refresh_enabled": True},
{"auto_refresh_enabled": False, "auto_refresh_interval_seconds": 0.5},
],
)
def test_remote_create_forwards_refresh_options_without_interval_default(options):
with mock_remote_materialized_view_create() as (host, requests):
db = lancedb.connect("db://dev", api_key="fake", host_override=host)
job = db.create_materialized_view_async("adults", "people", **options)
assert job.id == "mv-create-123"
assert requests == [
(
"/v1/materialized_view/adults/create",
{
"query": 'SELECT * FROM "people"',
**options,
},
)
]
@pytest.mark.asyncio
async def test_async_remote_create_keeps_interval_unspecified():
with mock_remote_materialized_view_create() as (host, requests):
db = await lancedb.connect_async("db://dev", api_key="fake", host_override=host)
job = await db.create_materialized_view_async(
"adults",
"people",
auto_refresh_enabled=True,
)
assert job.id == "mv-create-123"
assert requests[0][1] == {
"query": 'SELECT * FROM "people"',
"auto_refresh_enabled": True,
}
@pytest.mark.parametrize(
"options",
[
{"auto_refresh_enabled": True},
{"auto_refresh_interval_seconds": 300},
],
)
def test_embedded_create_refuses_automatic_refresh_before_creating_view(
tmp_path, options
):
db = make_db(tmp_path)
with pytest.raises(NotImplementedError, match="Cloud or Enterprise"):
db.create_materialized_view("adults", "people", **options)
assert "adults" not in db.table_names()
@pytest.mark.parametrize("seconds", [0, -1, float("nan"), float("inf")])
def test_remote_create_rejects_invalid_interval_without_request(seconds):
with mock_remote_materialized_view_create() as (host, requests):
db = lancedb.connect("db://dev", api_key="fake", host_override=host)
with pytest.raises(ValueError, match="positive, finite"):
db.create_materialized_view_async(
"adults", "people", auto_refresh_interval_seconds=seconds
)
assert requests == []
def test_namespace_forwards_automatic_refresh_request(tmp_path):
db = lancedb.connect_namespace("dir", {"root": str(tmp_path)})
with pytest.raises(NotImplementedError, match="Cloud or Enterprise"):
db.create_materialized_view("mv", "source", auto_refresh_enabled=True)
@pytest.mark.asyncio
async def test_async_namespace_forwards_automatic_refresh_request(tmp_path):
db = lancedb.connect_namespace_async("dir", {"root": str(tmp_path)})
with pytest.raises(NotImplementedError, match="Cloud or Enterprise"):
await db.create_materialized_view("mv", "source", auto_refresh_enabled=True)
def test_embedded_explicitly_disabled_refresh_is_supported(tmp_path):
db = make_db(tmp_path)
view = db.create_materialized_view("mv", "people", auto_refresh_enabled=False)
assert view.table.count_rows() == 3
python/python/tests/test_merge_insert_lsm.py
+50 -58
View File
@@ -1,14 +1,14 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Tests for the MemWAL LSM ``merge_insert`` dispatch."""
"""Tests for the RTE LSM ``merge_insert`` dispatch."""
from datetime import timedelta
import lancedb
import pyarrow as pa
import pytest
from lancedb._lancedb import LsmWriteSpec
from lancedb._lancedb import RTEWriteSpec
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 RTE 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_rte(RTEWriteSpec.bucket("id", 1))
return table
@@ -67,7 +67,7 @@ def test_lsm_merge_insert_bucket(tmp_path):
.when_not_matched_insert_all()
.execute(_reader([3, 4, 5]))
)
# LSM path: rows go to the MemWAL, so only num_rows is populated.
# LSM path: rows go to the RTE, so only num_rows is populated.
assert result.num_rows == 3
assert result.version == 0
assert result.num_inserted_rows == 0
@@ -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_rte(RTEWriteSpec.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_rte(RTEWriteSpec.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_rte_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_rte(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_rte(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_rte_true_without_spec_errors(tmp_path):
# A table with a primary key but no RTE 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_rte(True) demands RTE routing; without a spec it errors.
with pytest.raises(Exception, match="use_rte"):
(
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.use_lsm(True)
.use_rte(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 RTE 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_rte_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_rte(RTEWriteSpec.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_rte_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_rte(RTEWriteSpec.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_rte(False) bypasses the RTE and reads the base table only.
base_only = table.search().use_rte(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_rte(RTEWriteSpec.unsharded())
_lsm_upsert(table, [2, 3, 4]) # ids 2,3,4 -> values 0,1,2
@@ -258,31 +258,31 @@ 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 RTE write spec
# No spec: default read and use_lsm(False) both read the base table, no error.
# No spec: default read and use_rte(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_rte(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_rte_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_rte(RTEWriteSpec.unsharded())
_lsm_upsert(table, [4])
# with_row_id is unsupported by the LSM scanner; on a MemWAL table the default
# with_row_id is unsupported by the LSM scanner; on an RTE table the default
# (auto-routed) read hard-errors instead of silently reading a stale base.
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_rte(False) is the escape hatch: it reads the base table only.
base = table.search().with_row_id(True).use_rte(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_rte(RTEWriteSpec.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 RTE 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,9 +353,7 @@ 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_rte(RTEWriteSpec.unsharded().with_maintained_indexes([index_name]))
return table
@@ -379,9 +377,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_rte(False) bypasses the RTE, 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_rte(False).limit(100).to_arrow()["id"].to_pylist()
)
assert {1, 2, 3} <= base_ids
assert 1000 not in base_ids
@@ -453,9 +451,7 @@ 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_rte(RTEWriteSpec.unsharded().with_maintained_indexes([index_name]))
# in-flight doc 4 lands in the memtable's maintained FTS index.
(
@@ -498,8 +494,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_rte(False) is the escape hatch: the base-only standard path honors them.
base = table.search(query).distance_range(0.0, 100.0).use_rte(False).to_arrow()
assert 1000 not in set(base["id"].to_pylist())
@@ -544,7 +540,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_rte(RTEWriteSpec.unsharded().with_maintained_indexes([]))
with pytest.raises(Exception, match="maintained"):
table.search("fox", query_type="fts", fts_columns="text").to_arrow()
@@ -553,7 +549,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_rte(RTEWriteSpec.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 +558,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_rte(False) reads the base table at the pinned version.
base = table.search().use_rte(False).to_arrow()
assert sorted(base["id"].to_pylist()) == [1, 2, 3]
@@ -571,14 +567,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_rte(RTEWriteSpec.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_rte(False) is the escape hatch: it reads the base table.
base = table.take_row_ids([0, 1]).use_rte(False).to_arrow()
assert base.num_rows == 2
@@ -588,9 +584,7 @@ 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_rte(RTEWriteSpec.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).
with pytest.raises(Exception, match="postfilter"):
@@ -608,9 +602,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_rte(RTEWriteSpec.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 +623,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_rte(RTEWriteSpec.unsharded().with_maintained_indexes([]))
with pytest.raises(Exception, match="maintained"):
table.search([1.0] * VECTOR_DIM).to_arrow()
@@ -644,7 +636,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_rte(RTEWriteSpec.unsharded().with_maintained_indexes([name]))
# Multiple physical segments of one logical index must not be miscounted as
# multiple indexes and rejected.
ids = set(
python/python/tests/test_namespace.py
+132
View File
@@ -6,6 +6,7 @@
import tempfile
import shutil
import importlib
import warnings
import http.server
import json
import threading
@@ -25,6 +26,125 @@ PUSHDOWN_DATA = pa.table(
)
class TestNamespaceCreateTableOptions:
@pytest.mark.parametrize("version", ["legacy", "stable"])
@pytest.mark.parametrize("empty", [False, True])
def test_deprecated_storage_version(self, tmp_path, version, empty):
db = lancedb.connect_namespace("dir", {"root": str(tmp_path / "namespace")})
data = None if empty else [{"id": 1}]
schema = pa.schema([("id", pa.int64())])
other_version = "stable" if version == "legacy" else "legacy"
options = {"new_table_data_storage_version": other_version}
control = lancedb.connect(tmp_path / "control").create_table(
"test",
data=data,
schema=schema,
storage_options={"new_table_data_storage_version": version},
)
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always", DeprecationWarning)
table = db.create_table(
"test",
data=data,
schema=schema,
storage_options=options,
data_storage_version=version,
)
assert (
table.to_lance().data_storage_version
== control.to_lance().data_storage_version
)
assert table.to_arrow().to_pylist() == ([] if empty else data)
assert options == {"new_table_data_storage_version": other_version}
assert any("data_storage_version" in str(w.message) for w in caught)
@pytest.mark.asyncio
@pytest.mark.parametrize("version", ["legacy", "stable"])
@pytest.mark.parametrize("empty", [False, True])
async def test_deprecated_storage_version_async(self, tmp_path, version, empty):
root = str(tmp_path / "namespace")
db = lancedb.connect_namespace_async("dir", {"root": root})
data = None if empty else [{"id": 1}]
schema = pa.schema([("id", pa.int64())])
other_version = "stable" if version == "legacy" else "legacy"
options = {"new_table_data_storage_version": other_version}
control = lancedb.connect(tmp_path / "control").create_table(
"test",
data=data,
schema=schema,
storage_options={"new_table_data_storage_version": version},
)
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always", DeprecationWarning)
table = await db.create_table(
"test",
data=data,
schema=schema,
storage_options=options,
data_storage_version=version,
)
reopened = lancedb.connect_namespace("dir", {"root": root}).open_table("test")
assert (
reopened.to_lance().data_storage_version
== control.to_lance().data_storage_version
)
assert (await table.to_arrow()).to_pylist() == ([] if empty else data)
assert options == {"new_table_data_storage_version": other_version}
assert any("data_storage_version" in str(w.message) for w in caught)
@pytest.mark.parametrize("enabled", [False, True])
@pytest.mark.parametrize("empty", [False, True])
def test_deprecated_manifest_paths(self, tmp_path, enabled, empty):
db = lancedb.connect_namespace("dir", {"root": str(tmp_path)})
data = None if empty else [{"id": 1}]
options = {"new_table_enable_v2_manifest_paths": str(not enabled).lower()}
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always", DeprecationWarning)
table = db.create_table(
"test",
data=data,
schema=pa.schema([("id", pa.int64())]),
storage_options=options,
enable_v2_manifest_paths=enabled,
)
assert table.uses_v2_manifest_paths() == enabled
assert table.to_arrow().to_pylist() == ([] if empty else data)
assert db.open_table("test").uses_v2_manifest_paths() == enabled
assert options == {
"new_table_enable_v2_manifest_paths": str(not enabled).lower()
}
assert any("enable_v2_manifest_paths" in str(w.message) for w in caught)
@pytest.mark.asyncio
@pytest.mark.parametrize("enabled", [False, True])
@pytest.mark.parametrize("empty", [False, True])
async def test_deprecated_manifest_paths_async(self, tmp_path, enabled, empty):
db = lancedb.connect_namespace_async("dir", {"root": str(tmp_path)})
data = None if empty else [{"id": 1}]
options = {"new_table_enable_v2_manifest_paths": str(not enabled).lower()}
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always", DeprecationWarning)
table = await db.create_table(
"test",
data=data,
schema=pa.schema([("id", pa.int64())]),
storage_options=options,
enable_v2_manifest_paths=enabled,
)
assert await table.uses_v2_manifest_paths() == enabled
assert (await table.to_arrow()).to_pylist() == ([] if empty else data)
reopened = await db.open_table("test")
assert await reopened.uses_v2_manifest_paths() == enabled
assert options == {
"new_table_enable_v2_manifest_paths": str(not enabled).lower()
}
assert any("enable_v2_manifest_paths" in str(w.message) for w in caught)
def _ipc_file(table: pa.Table = PUSHDOWN_DATA) -> bytes:
sink = pa.BufferOutputStream()
with pa.ipc.new_file(sink, table.schema) as writer:
@@ -1209,3 +1329,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)
)
python/python/tests/test_ollama_embeddings.py
+126
View File
@@ -0,0 +1,126 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import json
import httpx
import numpy as np
import pyarrow as pa
import pytest
from pydantic import ValidationError
from lancedb.embeddings import get_registry
@pytest.fixture
def ollama_embeddings():
pytest.importorskip("ollama")
requests = []
def respond(request):
body = json.loads(request.content)
requests.append(body)
if body["model"] == "missing-model":
return httpx.Response(404, json={"error": "model not found"})
return httpx.Response(
200,
json={
"embeddings": [
[float(index), float(len(text))]
for index, text in enumerate(body["input"])
]
},
)
function = (
get_registry()
.get("ollama")
.create(
name="nomic-embed-text",
options={"num_ctx": 4096},
keep_alive="10m",
ollama_client_kwargs={
"transport": httpx.MockTransport(respond),
"trust_env": False,
},
)
)
yield function, requests
function._ollama_client._client.close()
@pytest.mark.parametrize(
"texts",
[
pytest.param(["hello", "世界"], id="list"),
pytest.param(np.array(["hello", "世界"]), id="numpy-unicode"),
pytest.param(np.array(["hello", "世界"], dtype=object), id="numpy-object"),
pytest.param(
np.array([b"hello", "世界".encode("utf-8")]), id="numpy-utf8-bytes"
),
pytest.param([b"hello", "世界".encode("utf-8")], id="list-utf8-bytes"),
pytest.param(np.array(["hello", "unused", "世界"])[::2], id="numpy-slice"),
pytest.param(pa.array(["hello", "世界"]), id="arrow"),
pytest.param(pa.chunked_array([["hello"], ["世界"]]), id="chunked-arrow"),
],
)
def test_ollama_source_embeddings_inputs(ollama_embeddings, texts):
function, requests = ollama_embeddings
assert function.compute_source_embeddings(texts) == [[0.0, 5.0], [1.0, 2.0]]
assert requests == [
{
"model": "nomic-embed-text",
"input": ["hello", "世界"],
"options": {"num_ctx": 4096},
"keep_alive": "10m",
}
]
@pytest.mark.parametrize(
"texts",
[
pytest.param(np.array([["hello"], ["world"]]), id="numpy-two-dimensional"),
pytest.param([["hello"], ["world"]], id="list-two-dimensional"),
pytest.param(np.array([b"\xff"]), id="numpy-invalid-utf8"),
pytest.param([b"\xff"], id="list-invalid-utf8"),
],
)
def test_ollama_source_embeddings_rejects_invalid_input_locally(
ollama_embeddings, texts
):
function, requests = ollama_embeddings
with pytest.raises(ValidationError, match="valid string"):
function.compute_source_embeddings(texts)
assert requests == []
def test_ollama_generate_embeddings_empty_array(ollama_embeddings):
function, requests = ollama_embeddings
assert function.generate_embeddings(np.array([], dtype=str)) == []
assert len(requests) == 1
assert requests[0]["input"] == []
def test_ollama_query_embeddings(ollama_embeddings):
function, requests = ollama_embeddings
assert function.compute_query_embeddings("hello") == [[0.0, 5.0]]
assert len(requests) == 1
assert requests[0]["input"] == ["hello"]
def test_ollama_source_embeddings_preserves_api_error(ollama_embeddings):
ollama = pytest.importorskip("ollama")
function, requests = ollama_embeddings
function.name = "missing-model"
with pytest.raises(ollama.ResponseError, match="model not found") as error:
function.compute_source_embeddings(np.array(["hello"]))
assert error.value.status_code == 404
assert len(requests) == 1
python/python/tests/test_openai_reranker.py
+66
View File
@@ -0,0 +1,66 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import json
from types import SimpleNamespace
import pyarrow as pa
import pytest
from lancedb.rerankers import OpenaiReranker
def reranker_replying(documents):
"""Return an OpenaiReranker whose chat client replies with these documents."""
reply = json.dumps({"documents": documents})
def create(**kwargs):
message = SimpleNamespace(content=reply)
return SimpleNamespace(choices=[SimpleNamespace(message=message)])
reranker = OpenaiReranker(api_key="test-key")
reranker.__dict__["_client"] = SimpleNamespace(
chat=SimpleNamespace(completions=SimpleNamespace(create=create))
)
return reranker
def vector_results():
return pa.table(
{
"id": [1, 2, 3],
"text": ["dogs bark", "cats purr", "birds sing"],
"_distance": [0.1, 0.2, 0.3],
}
)
def test_scores_stay_with_their_rows_when_the_model_reorders():
reranker = reranker_replying(
[
{"content": "cats purr", "relevance_score": 0.9},
{"content": "birds sing", "relevance_score": 0.5},
{"content": "dogs bark", "relevance_score": 0.1},
]
)
result = reranker.rerank_vector("which animal purrs?", vector_results())
assert result.select(["id", "text"]).to_pylist() == [
{"id": 2, "text": "cats purr"},
{"id": 3, "text": "birds sing"},
{"id": 1, "text": "dogs bark"},
]
assert result["_relevance_score"].to_pylist() == pytest.approx([0.9, 0.5, 0.1])
def test_a_row_the_model_leaves_out_raises():
reranker = reranker_replying(
[
{"content": "cats purr", "relevance_score": 0.9},
{"content": "dogs bark", "relevance_score": 0.1},
]
)
with pytest.raises(ValueError, match="1 of 3"):
reranker.rerank_vector("which animal purrs?", vector_results())
python/python/tests/test_permutation.py
+106
View File
@@ -115,6 +115,35 @@ def test_split_random_fixed(mem_db):
assert split_ids.count(split_id) == 25
def test_permutation_over_table_with_schema_metadata(tmp_path):
"""Schema-level metadata on the base table must not break the build.
Hugging Face datasets carry keys like `huggingface` or `source-column-name-map`.
"""
schema = pa.schema(
[pa.field("x", pa.int64())],
# `split_names` collides with a permutation config key; the config must win.
metadata={b"huggingface": b"{}", b"split_names": b'["base"]'},
)
tbl = connect(tmp_path).create_table(
"t", pa.table({"x": list(range(100))}, schema=schema)
)
permutation_tbl = (
permutation_builder(tbl)
.split_random(fixed=2, seed=1, split_names=["train", "test"])
.execute()
)
permutations = Permutations(tbl, permutation_tbl)
assert permutations.split_names == ["train", "test"]
train = Permutation.from_tables(tbl, permutation_tbl, "train")
test = Permutation.from_tables(tbl, permutation_tbl, "test")
assert train.num_rows + test.num_rows == 100
train_x = [row["x"] for batch in train.iter(16) for row in batch]
test_x = [row["x"] for batch in test.iter(16) for row in batch]
assert sorted(train_x + test_x) == list(range(100))
def test_split_random_with_seed(mem_db):
"""Test that seeded random splits are reproducible."""
tbl = mem_db.create_table("test_table", pa.table({"x": range(50), "y": range(50)}))
@@ -132,6 +161,76 @@ def test_split_random_with_seed(mem_db):
assert data1["split_id"] == data2["split_id"]
def test_bootstrap_sample_fixed(mem_db):
"""Output row count is num_samples, not the input row count, and splits divide
it evenly."""
tbl = mem_db.create_table("test_table", pa.table({"x": range(20)}))
permutation_tbl = (
permutation_builder(tbl).bootstrap_sample(1000, fixed=4, seed=42).execute()
)
assert permutation_tbl.count_rows() == 1000
data = permutation_tbl.search(None).to_arrow().to_pydict()
split_ids = data["split_id"]
assert set(split_ids) == {0, 1, 2, 3}
for split_id in range(4):
assert split_ids.count(split_id) == 250
# Every sampled row_id must still be one of the 20 real rows.
assert set(data["row_id"]).issubset(set(range(20)))
def test_bootstrap_sample_draws_with_replacement(mem_db):
"""Sampling more rows than exist must produce repeats."""
tbl = mem_db.create_table("test_table", pa.table({"x": range(5)}))
permutation_tbl = (
permutation_builder(tbl).bootstrap_sample(500, fixed=1, seed=1).execute()
)
data = permutation_tbl.search(None).to_arrow().to_pydict()
assert len(data["row_id"]) == 500
# Only 5 distinct rows exist, so with 500 draws every one of them must repeat.
for row_id in range(5):
assert data["row_id"].count(row_id) > 1
def test_bootstrap_sample_with_seed_is_reproducible(mem_db):
tbl = mem_db.create_table("test_table", pa.table({"x": range(50)}))
perm1 = permutation_builder(tbl).bootstrap_sample(200, fixed=2, seed=42).execute()
perm2 = permutation_builder(tbl).bootstrap_sample(200, fixed=2, seed=42).execute()
data1 = perm1.search(None).to_arrow().to_pydict()
data2 = perm2.search(None).to_arrow().to_pydict()
assert data1["row_id"] == data2["row_id"]
assert data1["split_id"] == data2["split_id"]
def test_bootstrap_sample_respects_filter(mem_db):
tbl = mem_db.create_table("test_table", pa.table({"x": range(20)}))
permutation_tbl = (
permutation_builder(tbl)
.filter("x >= 15")
.bootstrap_sample(100, fixed=1, seed=7)
.execute()
)
data = permutation_tbl.search(None).to_arrow().to_pydict()
assert len(data["row_id"]) == 100
assert set(data["row_id"]).issubset(set(range(15, 20)))
def test_bootstrap_sample_requires_exactly_one_size_arg(mem_db):
tbl = mem_db.create_table("test_table", pa.table({"x": range(10)}))
with pytest.raises(ValueError):
permutation_builder(tbl).bootstrap_sample(10).execute()
with pytest.raises(ValueError):
permutation_builder(tbl).bootstrap_sample(
10, ratios=[0.5, 0.5], fixed=2
).execute()
def test_split_hash(mem_db):
"""Test hash-based splitting."""
tbl = mem_db.create_table(
@@ -712,6 +811,13 @@ def test_rename_column(some_permutation: Permutation):
some_permutation.rename_column("non_existent", "new_value")
def test_rename_column_keeps_position(some_permutation: Permutation):
renamed = some_permutation.rename_column("id", "new_id")
assert renamed.schema == pa.schema([("new_id", pa.int64()), ("value", pa.int64())])
batch = renamed.with_format("arrow").__getitems__([0])
assert batch.schema.names == ["new_id", "value"]
def test_rename_columns(some_permutation: Permutation):
assert some_permutation.rename_columns({"value": "new_value"}).schema == pa.schema(
[("id", pa.int64()), ("new_value", pa.int64())]
python/python/tests/test_pyarrow.py
+17
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import pyarrow as pa
import pyarrow.compute as pc
import lancedb
from lancedb.integrations.pyarrow import PyarrowDatasetAdapter
@@ -48,3 +49,19 @@ def test_dataset_adapter(tmp_path):
assert adapter.head(10).num_rows == 0
assert adapter.scanner().projected_schema == pa.schema({"x": pa.int64()})
def test_dataset_adapter_expression_filter(tmp_path):
data = pa.table({"x": [1, 2, 3, 4], "y": [5, 6, 7, 8]})
conn = lancedb.connect(tmp_path)
tbl = conn.create_table("test", data)
adapter = PyarrowDatasetAdapter(tbl)
filt = pc.field("x") > 2
assert adapter.count_rows(filt) == 2
assert adapter.to_table(filter=filt) == data.slice(2, 2)
assert adapter.head(1, filter=filt) == data.slice(2, 1)
# The filter can reference a column that is not projected
assert adapter.to_table(columns=["y"], filter=filt) == pa.table({"y": [7, 8]})
assert adapter.scanner(filter=filt).count_rows() == 2
python/python/tests/test_pydantic.py
+110
View File
@@ -415,6 +415,67 @@ def test_nullable_vector():
assert schema == pa.schema([pa.field("vec", pa.list_(pa.float32(), 16), True)])
@pytest.mark.parametrize("vector_type", [Vector, MultiVector])
@pytest.mark.parametrize(
"value_type, values",
[
(pa.int64(), [-(2**53 + 1), 2**53 + 3]),
(pa.uint64(), [2**53 + 1, 2**64 - 1]),
(pa.int8(), [-128, 127]),
(pa.uint8(), [0, 255]),
(pa.float32(), [1.25, 2.5]),
(pa.float64(), [1.25, 2.5]),
],
)
def test_typed_vector_round_trip(tmp_path, vector_type, value_type, values):
import lancedb
annotation = vector_type(2, value_type=value_type)
expected = [values, values] if vector_type is MultiVector else values
class Row(LanceModel):
id: int
payload: annotation
row = Row(id=1, payload=expected)
table = lancedb.connect(tmp_path).create_table("typed_vectors", [row], schema=Row)
table.add([Row(id=2, payload=expected)])
assert table.schema == Row.to_arrow_schema()
assert table.search().to_arrow().to_pylist() == [
{"id": 1, "payload": expected},
{"id": 2, "payload": expected},
]
assert json.loads(row.model_dump_json())["payload"] == expected
@pytest.mark.parametrize("vector_type", [Vector, MultiVector])
@pytest.mark.parametrize("nullable", [False, True])
def test_integer_vector_validation(vector_type, nullable):
annotation = vector_type(2, value_type=pa.int64(), nullable=nullable)
class Row(LanceModel):
payload: annotation
def shape(values):
return [values] if vector_type is MultiVector else values
assert Row.to_arrow_schema().field("payload").nullable is nullable
assert Row(payload=shape([1.0, -2.0])).model_dump()["payload"] == shape([1, -2])
with pytest.raises(pydantic.ValidationError, match="valid integer"):
Row(payload=shape([1.5, 2]))
with pytest.raises(pydantic.ValidationError):
Row(payload=shape([1]))
with pytest.raises(pydantic.ValidationError):
Row(payload=shape([1, 2, 3]))
with pytest.raises(pydantic.ValidationError):
Row(payload=None)
class OptionalRow(LanceModel):
payload: Optional[annotation] = None
assert OptionalRow(payload=None).payload is None
def test_bare_vector_raises_clear_error():
namespace = {
"__name__": "test_model_without_pyarrow",
@@ -733,3 +794,52 @@ def test_enum_types():
assert schema.field("priority").type == pa.int64()
assert schema.field("opt_status").type == pa.dictionary(pa.int32(), pa.utf8())
assert schema.field("opt_status").nullable
def test_plain_enum_values_are_written(mem_db):
"""Members of an Enum without a str/int mixin are written as their values."""
class Color(Enum):
RED = "red"
BLUE = "blue"
class Priority(Enum):
LOW = 1
HIGH = 2
class Tag(LanceModel):
color: Color
class TestModel(LanceModel):
color: Color
priority: Priority
opt_color: Optional[Color] = None
colors: List[Color]
tag: Tag
vector: Vector(2)
rows = [
TestModel(
color=Color.RED,
priority=Priority.HIGH,
opt_color=Color.BLUE,
colors=[Color.BLUE, Color.RED],
tag=Tag(color=Color.BLUE),
vector=[1.0, 2.0],
),
TestModel(
color=Color.BLUE,
priority=Priority.LOW,
colors=[],
tag=Tag(color=Color.RED),
vector=[3.0, 4.0],
),
]
tbl = mem_db.create_table("test", schema=TestModel)
tbl.add(rows)
assert tbl.to_arrow().select(["color", "priority", "opt_color"]).to_pylist() == [
{"color": "red", "priority": 2, "opt_color": "blue"},
{"color": "blue", "priority": 1, "opt_color": None},
]
assert tbl.search().to_pydantic(TestModel) == rows
python/python/tests/test_pydantic_excluded_fields.py
+128
View File
@@ -0,0 +1,128 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import lancedb
import pytest
from pydantic import BaseModel, Field
from lancedb.pydantic import LanceModel
class Details(BaseModel):
label: str
transient: str = Field(default="nested default", exclude=True)
class Row(LanceModel):
id: int
details: Details
transient: str = Field(default="local default", exclude=True)
generated: list[str] = Field(default_factory=lambda: ["local"], exclude=True)
retained: str = Field(default="persisted", exclude=False)
def rows():
return [
Row(
id=i,
details=Details(label=str(i), transient="private"),
transient="private",
generated=["private"],
)
for i in (1, 2)
]
def assert_schema(schema):
assert schema.names == ["id", "details", "retained"]
assert [field.name for field in schema.field("details").type] == ["label"]
assert Row.field_names() == ["id", "details", "transient", "generated", "retained"]
def assert_restored(restored):
assert sorted(row.id for row in restored) == [1, 2]
for row in restored:
assert row.transient == "local default"
assert row.generated == ["local"]
assert row.details.transient == "nested default"
assert row.details.label == str(row.id)
assert row.retained == "persisted"
@pytest.mark.parametrize("explicit_schema", [False, True])
def test_excluded_fields_round_trip(tmp_path, explicit_schema):
data = rows()
kwargs = {"schema": Row} if explicit_schema else {}
db = lancedb.connect(tmp_path)
table = db.create_table("rows", data=data[:1], **kwargs)
table.add(data[1:])
assert_schema(table.schema)
reopened = lancedb.connect(tmp_path).open_table("rows")
assert_restored(reopened.search().to_pydantic(Row))
assert data[0].transient == "private"
assert data[0].generated == ["private"]
@pytest.mark.asyncio
@pytest.mark.parametrize("explicit_schema", [False, True])
async def test_excluded_fields_round_trip_async(tmp_path, explicit_schema):
data = rows()
kwargs = {"schema": Row} if explicit_schema else {}
db = await lancedb.connect_async(tmp_path)
table = await db.create_table("rows", data=data[:1], **kwargs)
await table.add(data[1:])
assert_schema(await table.schema())
reopened_db = await lancedb.connect_async(tmp_path)
reopened = await reopened_db.open_table("rows")
assert_restored(await reopened.query().to_pydantic(Row))
assert data[0].details.transient == "private"
@pytest.mark.parametrize("excluded", ["transient", "source", "vector"])
def test_excluded_fields_keep_embedding_metadata_consistent(
tmp_path, monkeypatch, excluded
):
from lancedb.embeddings.base import TextEmbeddingFunction
from lancedb.embeddings.registry import get_registry
from lancedb.pydantic import Vector
calls = []
class LocalEmbedding(TextEmbeddingFunction):
def ndims(self):
return 2
def generate_embeddings(self, texts):
calls.extend(texts)
return [[1.0, 0.0] for _ in texts]
registry = get_registry()
monkeypatch.setattr(registry, "_functions", registry._functions.copy())
registry.register("test-excluded-fields")(LocalEmbedding)
function = LocalEmbedding.create(max_retries=0)
class EmbeddedRow(LanceModel):
text: str = function.SourceField(default="alpha", exclude=excluded == "source")
vector: Vector(2) = function.VectorField(
default_factory=lambda: [1.0, 0.0], exclude=excluded == "vector"
)
transient: str = Field(default="local", exclude=True)
configs = EmbeddedRow.parse_embedding_functions()
assert len(configs) == (1 if excluded == "transient" else 0)
db = lancedb.connect(tmp_path)
data = [{"text": "alpha"}] if excluded == "transient" else [EmbeddedRow()]
table = db.create_table("embeddings", schema=EmbeddedRow, data=data)
assert "transient" not in table.schema.names
assert ("text" in table.schema.names) == (excluded != "source")
assert ("vector" in table.schema.names) == (excluded != "vector")
reopened = lancedb.connect(tmp_path).open_table("embeddings")
assert len(reopened.embedding_functions) == len(configs)
restored = reopened.search().to_pydantic(EmbeddedRow)
assert restored[0].text == "alpha"
assert restored[0].vector == [1.0, 0.0]
assert restored[0].transient == "local"
if excluded == "transient":
assert calls == ["alpha"]
assert reopened.search("alpha").limit(1).to_list()[0]["text"] == "alpha"
else:
assert calls == []
python/python/tests/test_query.py
+370 -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):
@@ -614,6 +833,28 @@ def test_with_row_id(table: lancedb.table.Table):
assert rs["_rowid"].to_pylist() == [0, 1]
def test_filter_preserves_deeply_nested_struct_fields(tmp_db):
# Regression test for https://github.com/lancedb/lancedb/issues/2217.
data = [
{
"item": "test",
"sub_obj": {
"second_sub_obj": [
{
"third_sub_obj": {"int_val": 1},
"str_val": "test",
}
]
},
}
]
table = tmp_db.create_table("deeply_nested_struct", data=data)
result = table.search().where("item = 'test'").limit(1).to_arrow()
assert result.to_pylist() == data
def test_blob_v2_query_omits_auto_row_id(tmp_db):
table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_omits_auto_rowid")
@@ -2265,3 +2506,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)
python/python/tests/test_read_consistency_interval.py
+83
View File
@@ -0,0 +1,83 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
from datetime import timedelta
import lancedb
import pytest
from lance_namespace import connect as connect_namespace_client
from lancedb import _lancedb
from lancedb.namespace import LanceNamespaceDBConnection
@pytest.mark.parametrize("seconds", [-1.0, -0.000001])
@pytest.mark.parametrize("kind", ["sync", "namespace", "namespace_async", "client"])
def test_negative_read_consistency_interval(tmp_path, seconds, kind):
interval = timedelta(seconds=seconds)
with pytest.raises(ValueError, match="read_consistency_interval"):
if kind == "sync":
lancedb.connect(tmp_path, read_consistency_interval=interval)
elif kind == "client":
client = connect_namespace_client("dir", {"root": str(tmp_path)})
LanceNamespaceDBConnection(client, read_consistency_interval=interval)
else:
connect = (
lancedb.connect_namespace
if kind == "namespace"
else lancedb.connect_namespace_async
)
connect("dir", {"root": str(tmp_path)}, read_consistency_interval=interval)
@pytest.mark.asyncio
@pytest.mark.parametrize("seconds", [-1.0, -0.000001])
async def test_negative_async_read_consistency_interval(tmp_path, seconds):
with pytest.raises(ValueError, match="read_consistency_interval"):
await lancedb.connect_async(
tmp_path, read_consistency_interval=timedelta(seconds=seconds)
)
@pytest.mark.asyncio
@pytest.mark.parametrize("seconds", [float("nan"), float("inf"), -float("inf"), 1e30])
@pytest.mark.parametrize("kind", ["connect", "namespace", "client"])
async def test_binding_rejects_invalid_read_consistency_interval(
tmp_path, seconds, kind
):
with pytest.raises(ValueError, match="read_consistency_interval"):
if kind == "connect":
await _lancedb.connect(str(tmp_path), read_consistency_interval=seconds)
elif kind == "namespace":
_lancedb.connect_namespace(
"dir", {"root": str(tmp_path)}, read_consistency_interval=seconds
)
else:
client = connect_namespace_client("dir", {"root": str(tmp_path)})
_lancedb.connect_namespace_client(client, read_consistency_interval=seconds)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"interval", [None, timedelta(0), timedelta(microseconds=1), timedelta(seconds=5)]
)
async def test_valid_read_consistency_intervals(tmp_path, interval):
db = lancedb.connect(tmp_path / "sync", read_consistency_interval=interval)
assert db.read_consistency_interval == interval
db.close()
db = await lancedb.connect_async(
tmp_path / "async", read_consistency_interval=interval
)
assert await db.get_read_consistency_interval() == interval
db.close()
db = lancedb.connect_namespace(
"dir", {"root": str(tmp_path / "namespace")}, read_consistency_interval=interval
)
assert db.read_consistency_interval == interval
db.close()
client = connect_namespace_client("dir", {"root": str(tmp_path / "client")})
db = LanceNamespaceDBConnection(client, read_consistency_interval=interval)
assert db.read_consistency_interval == interval
db.close()
python/python/tests/test_remote_db.py
+347 -31
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,8 +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"),
]
@@ -1794,6 +1793,79 @@ def test_table_wait_for_index_timeout():
table.wait_for_index(["id_idx"], timedelta(seconds=1))
def _jieba_fts_handler(request):
"""Serve a table with a ``jieba/default`` FTS index and no other state."""
if request.path == "/v1/table/test/describe/":
body = {
"version": 1,
"schema": {
"fields": [
{"name": "text", "type": {"type": "string"}, "nullable": True}
]
},
}
elif request.path == "/v1/table/test/index/list/":
index_details = {
"base_tokenizer": "jieba/default",
"language": "English",
"with_position": False,
"max_token_length": 40,
"lower_case": True,
"stem": False,
"remove_stop_words": False,
"ascii_folding": False,
}
body = {
"indexes": [
{
"index_name": "text_idx",
"columns": ["text"],
"index_type": "FTS",
"index_details": json.dumps(index_details),
}
]
}
else:
request.send_response(404)
request.end_headers()
return
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(json.dumps(body).encode())
_MODEL_FILES_ERROR = (
r"(?s)Invalid directory path.*"
r"full text search index 'text_idx'.*"
r"tokenize\(\) runs on the client.*'jieba/default'.*"
r"LANCE_LANGUAGE_MODEL_HOME.*match the files the server uses"
)
def test_remote_tokenize_model_backed_without_model_files(tmp_path, monkeypatch):
monkeypatch.setenv("LANCE_LANGUAGE_MODEL_HOME", str(tmp_path))
with mock_lancedb_connection(_jieba_fts_handler) as db:
table = db.open_table("test")
with pytest.raises(ValueError, match=_MODEL_FILES_ERROR):
table.tokenize("我们是中国人", index_name="text_idx")
with pytest.raises(ValueError, match=_MODEL_FILES_ERROR):
table.tokenize("我们是中国人", column="text")
@pytest.mark.asyncio
async def test_remote_tokenize_model_backed_without_model_files_async(
tmp_path, monkeypatch
):
monkeypatch.setenv("LANCE_LANGUAGE_MODEL_HOME", str(tmp_path))
async with mock_lancedb_connection_async(_jieba_fts_handler) as db:
table = await db.open_table("test")
with pytest.raises(ValueError, match=_MODEL_FILES_ERROR):
await table.tokenize("我们是中国人", index_name="text_idx")
with pytest.raises(ValueError, match=_MODEL_FILES_ERROR):
await table.tokenize("我们是中国人", column="text")
def test_stats():
stats = {
"total_bytes": 38,
@@ -1845,10 +1917,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_rte, compact_rte, get_rte_stats, and is
responsible for writing the response.
"""
routes = ("flush_lsm", "compact_lsm", "get_lsm_stats")
routes = ("flush_rte", "compact_rte", "get_rte_stats")
def handler(request):
match = re.fullmatch(r"/v1/table/test/(\w+)/", request.path)
@@ -1880,9 +1952,9 @@ def send_json(request, payload, status=200):
request.wfile.write(json.dumps(payload).encode())
def test_get_lsm_stats_sync():
def test_get_rte_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 +1962,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 +1977,29 @@ def test_get_lsm_stats_sync():
seen_bodies = []
def lsm_handler(request, route):
assert route == "get_lsm_stats"
assert route == "get_rte_stats"
seen_bodies.append(read_json_body(request))
send_json(request, {"lsm_stats": {"buckets": [bucket]}})
send_json(request, {"rte_stats": {"table_shards": [table_shard]}})
with lsm_test_table(lsm_handler) as table:
assert table.get_lsm_stats() == {"buckets": [bucket]}
assert table.get_rte_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_rte_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_rte_stats_sync_returns_none_when_rte_disabled():
"""A null envelope means the RTE write path is not enabled, not an error."""
def lsm_handler(request, route):
send_json(request, {"lsm_stats": None})
send_json(request, {"rte_stats": None})
with lsm_test_table(lsm_handler) as table:
assert table.get_lsm_stats() is None
assert table.get_rte_stats() is None
def test_flush_and_compact_lsm_sync():
def test_flush_and_compact_rte_sync():
"""Both are one-shot POSTs answered 202 with no body."""
called = []
@@ -1937,13 +2009,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_rte() is None
assert table.compact_rte() is None
assert called == ["flush_rte", "compact_rte"]
def test_checkpoint_lsm_sync():
"""Seal, read the watermark, and return once L0 holds nothing.
def test_checkpoint_rte_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 +2024,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_rte_stats":
# An empty SSTable tier yields no target watermark, so the loop is done
# after the freeze without ever polling compaction.
send_json(request, {"rte_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_rte() is None
assert called == ["flush_rte", "get_rte_stats"]
@contextlib.contextmanager
@@ -2728,6 +2800,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)
@@ -3250,6 +3327,89 @@ BLOB_DESCRIBE_RESPONSE = {
}
@pytest.mark.asyncio
@pytest.mark.parametrize("use_async", [False, True])
@pytest.mark.parametrize("mode", ["merge", "overwrite"])
@pytest.mark.parametrize("target_with_ranges", [False, True])
async def test_remote_writes_normalize_blob_layouts(
use_async, mode, target_with_ranges
):
description = json.loads(json.dumps(BLOB_DESCRIBE_RESPONSE))
if target_with_ranges:
description["schema"]["fields"][1]["type"]["fields"].extend(
{"name": name, "type": {"type": "uint64"}, "nullable": True}
for name in ("position", "size")
)
received = []
def handler(request):
if request.path == "/v1/table/test/describe/":
response = description
else:
expected = "merge_insert" if mode == "merge" else "insert"
assert urlparse(request.path).path == f"/v1/table/test/{expected}/"
if request.headers.get("Transfer-Encoding") == "chunked":
body = bytearray()
while size := int(request.rfile.readline().split(b";")[0], 16):
body.extend(request.rfile.read(size))
assert request.rfile.read(2) == b"\r\n"
assert request.rfile.read(2) == b"\r\n"
else:
body = request.rfile.read(int(request.headers["Content-Length"]))
received.append(pa.ipc.open_stream(body).read_all())
response = {
"version": 2,
"num_updated_rows": 0,
"num_inserted_rows": 2,
"num_deleted_rows": 0,
}
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.send_header("phalanx-version", "0.5.0")
request.end_headers()
request.wfile.write(json.dumps(response).encode())
field = lancedb.blob("image")
storage_type = field.type.storage_type
if target_with_ranges:
storage_type = pa.struct(list(storage_type)[:2])
field = pa.field(
"image",
storage_type,
metadata={"ARROW:extension:name": "lance.blob.v2"},
)
images = pa.array([{"data": b"payload"}, None], type=storage_type)
if isinstance(field.type, pa.ExtensionType):
images = pa.ExtensionArray.from_storage(field.type, images)
data = pa.Table.from_arrays(
[pa.array([1, 2], type=pa.int64()), images],
schema=pa.schema([pa.field("id", pa.int64()), field]),
)
with mock_lancedb_connection(handler) as db:
remote_table = db.open_table("test")
table = remote_table._table if use_async else remote_table
if mode == "merge":
result = (
table.merge_insert("id").when_not_matched_insert_all().execute(data)
)
else:
result = table.add(data, mode="overwrite")
if use_async:
await result
assert len(received) == 1
images = received[0]["image"].combine_chunks()
if isinstance(images.type, pa.ExtensionType):
images = images.storage
with_ranges = target_with_ranges if mode == "merge" else not target_with_ranges
assert images.type.num_fields == (4 if with_ranges else 2)
expected = {"data": b"payload", "uri": None}
if with_ranges:
expected.update(position=None, size=None)
assert images.to_pylist() == [expected, None]
def blob_query_response_table():
image_field = pa.field(
"image",
@@ -3357,13 +3517,20 @@ def blob_remote_table(*, server_version=Version("0.5.0")):
else:
assert body["with_row_id"] is True
if columns == ["image"]:
# fetch_blob_files sizes its handles from a descriptor take.
# Blob handles are sized from a descriptor take.
assert body["filter"].startswith("_rowid IN")
response_table = blob_descriptor_take_table()
row_ids = [
int(value) for value in re.findall(r"\d+", body["filter"])
]
response_table = response_table.filter(
pa.compute.is_in(response_table["_rowid"], pa.array(row_ids))
)
else:
response_table = blob_query_response_table()
request.send_response(200)
request.send_header("Content-Type", "application/vnd.apache.arrow.file")
request.send_header("x-lancedb-version", "1")
request.end_headers()
with pa.ipc.new_file(request.wfile, response_table.schema) as writer:
writer.write_table(response_table)
@@ -3395,6 +3562,42 @@ def test_remote_blob_columns_and_fetch():
assert blobs.to_pylist() == [b"alpha", None, b"gamma"]
@pytest.mark.asyncio
@pytest.mark.parametrize("use_async", [False, True])
async def test_remote_fetch_blob_ranges(use_async):
with blob_remote_table() as remote_table:
table = remote_table._table if use_async else remote_table
requests = [
(30, 2, 3),
(10, 1, 3),
(20, 0, 1),
(10, 0, 2),
(40, 0, 0),
(10, 5, 0),
(30, 2, 3),
]
result = table.fetch_blob_ranges("image", requests)
if use_async:
result = await result
assert isinstance(result, pa.LargeBinaryArray)
assert result.to_pylist() == [b"mma", b"lph", None, b"al", b"", b"", b"mma"]
for invalid, message in [
([(10, 0, 1), (99, 0, 1)], "only 1 exist"),
([(10, 4, 2)], "exceeds blob size"),
([(20, 2**64 - 1, 1)], "overflows u64"),
]:
with pytest.raises(ValueError, match=message):
result = table.fetch_blob_ranges("image", invalid)
if use_async:
await result
result = table.fetch_blob_ranges("image", [])
if use_async:
result = await result
assert result.to_pylist() == []
@pytest.mark.asyncio
async def test_async_remote_to_lance_rejects_storage_access(monkeypatch):
with blob_remote_table() as remote_table:
@@ -3463,6 +3666,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 +3736,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 +3781,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 +3860,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():
python/python/tests/test_reranker_credentials.py
+74
View File
@@ -0,0 +1,74 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import sys
import types
import pytest
from lancedb.rerankers import CohereReranker, VoyageAIReranker
RERANKERS = [
pytest.param("cohere", "COHERE_API_KEY", CohereReranker, {}, id="cohere"),
pytest.param(
"voyageai",
"VOYAGE_API_KEY",
VoyageAIReranker,
{"model_name": "rerank-2.5"},
id="voyageai",
),
]
def fake_sdk(monkeypatch, module_name):
"""Replace the SDK module with one whose Client records the API key."""
keys = []
def client(api_key=None):
keys.append(api_key)
return object()
module = types.ModuleType(module_name)
module.Client = client
monkeypatch.setitem(sys.modules, module_name, module)
return keys
@pytest.mark.parametrize("module_name, env_var, reranker_cls, kwargs", RERANKERS)
@pytest.mark.parametrize(
"explicit_key, environment_key, expected_key",
[
("explicit-test-key", "environment-test-key", "explicit-test-key"),
("explicit-test-key", None, "explicit-test-key"),
(None, "environment-test-key", "environment-test-key"),
],
)
def test_reranker_api_key_precedence(
monkeypatch,
module_name,
env_var,
reranker_cls,
kwargs,
explicit_key,
environment_key,
expected_key,
):
keys = fake_sdk(monkeypatch, module_name)
if environment_key is None:
monkeypatch.delenv(env_var, raising=False)
else:
monkeypatch.setenv(env_var, environment_key)
reranker_cls(api_key=explicit_key, **kwargs)._client
assert keys == [expected_key]
@pytest.mark.parametrize("module_name, env_var, reranker_cls, kwargs", RERANKERS)
def test_reranker_requires_api_key(
monkeypatch, module_name, env_var, reranker_cls, kwargs
):
keys = fake_sdk(monkeypatch, module_name)
monkeypatch.delenv(env_var, raising=False)
with pytest.raises(ValueError, match=f"{env_var} not set"):
reranker_cls(**kwargs)._client
assert keys == []
python/python/tests/test_rerankers.py
+148
View File
@@ -785,6 +785,154 @@ def test_empty_hybrid_result_reranker():
assert "_rowid" in result.column_names
@pytest.mark.parametrize("norm", ["score", "rank"])
def test_hybrid_return_all_restores_original_scores(norm):
from lancedb.query import LanceHybridQueryBuilder
vector_results = pa.table(
{
"text": ["a", "b", "c"],
"_rowid": pa.array([0, 1, 2], type=pa.uint64()),
"_distance": pa.array([0.5, 2.5, 10.0], type=pa.float32()),
}
)
fts_results = pa.table(
{
"text": ["b", "a"],
"_rowid": pa.array([1, 0], type=pa.uint64()),
"_score": pa.array([3.25, 1.75], type=pa.float32()),
}
)
result = LanceHybridQueryBuilder._combine_hybrid_results(
fts_results=fts_results,
vector_results=vector_results,
norm=norm,
fts_query="query",
reranker=RRFReranker(return_score="all"),
limit=10,
with_row_ids=True,
)
# The raw distances and scores are returned, not the normalized values
# or ranks used for reranking.
rows = {row["_rowid"]: row for row in result.to_pylist()}
assert rows[0]["_distance"] == 0.5
assert rows[1]["_distance"] == 2.5
assert rows[2]["_distance"] == 10.0
assert rows[0]["_score"] == 1.75
assert rows[1]["_score"] == 3.25
assert rows[2]["_score"] is None
@pytest.mark.parametrize("return_score", ["relevance", "all"])
@pytest.mark.parametrize("first_row_id", [1, 2**63])
@pytest.mark.parametrize("score_type", [pa.float32(), pa.float64()])
def test_model_reranker_preserves_arrow_schema(
mocker, return_score, first_row_id, score_type
):
common_schema = pa.schema(
[
pa.field("_rowid", pa.uint64()),
pa.field("text", pa.string()),
pa.field("vector", pa.list_(pa.float32(), 2)),
],
metadata={b"source": b"reranker-test"},
)
vector_results = pa.Table.from_pylist(
[
{
"_rowid": first_row_id,
"text": "vector",
"vector": [1, 0],
"_distance": 0.1,
},
{
"_rowid": first_row_id + 1,
"text": "both",
"vector": [0, 1],
"_distance": 0.2,
},
],
schema=common_schema.append(pa.field("_distance", score_type)),
)
fts_results = pa.Table.from_pylist(
[
{
"_rowid": first_row_id + 1,
"text": "both",
"vector": [0, 1],
"_score": 0.9,
},
{
"_rowid": first_row_id + 2,
"text": "fts",
"vector": [1, 1],
"_score": 0.8,
},
],
schema=common_schema.append(pa.field("_score", score_type)),
)
reranker = CohereReranker(return_score=return_score)
reranker._client = mocker.Mock()
reranker._client.rerank.return_value.results = [
mocker.Mock(index=index, relevance_score=score)
for index, score in [(2, 0.9), (1, 0.8), (0, 0.7)]
]
result = reranker.rerank_hybrid("query", vector_results, fts_results)
assert result.select(common_schema.names).schema.equals(
common_schema, check_metadata=True
)
assert result["_rowid"].to_pylist() == [
first_row_id + 2,
first_row_id + 1,
first_row_id,
]
assert result["text"].to_pylist() == ["fts", "both", "vector"]
if return_score == "all":
assert (
result["_distance"]
.combine_chunks()
.equals(pa.array([None, 0.2, 0.1], type=score_type))
)
assert (
result["_score"]
.combine_chunks()
.equals(pa.array([0.8, 0.9, None], type=score_type))
)
@pytest.mark.parametrize("return_score", ["relevance", "all"])
def test_model_reranker_empty_results_preserve_schema(mocker, return_score):
common_schema = pa.schema(
[
("_rowid", pa.uint64()),
("text", pa.string()),
("vector", pa.list_(pa.float32(), 2)),
],
metadata={b"source": b"reranker-test"},
)
vector_results = pa.Table.from_pylist(
[], schema=common_schema.append(pa.field("_distance", pa.float32()))
)
fts_results = pa.Table.from_pylist(
[], schema=common_schema.append(pa.field("_score", pa.float32()))
)
reranker = CohereReranker(return_score=return_score)
reranker._client = mocker.Mock()
result = reranker.rerank_hybrid("query", vector_results, fts_results)
assert len(result) == 0
assert result.select(common_schema.names).schema.equals(
common_schema, check_metadata=True
)
assert result.schema.field("_relevance_score").type == pa.float32()
reranker._client.rerank.assert_not_called()
def test_cross_encoder_reranker_return_all(tmp_path):
pytest.importorskip("sentence_transformers")
reranker = CrossEncoderReranker(return_score="all")
python/python/tests/test_lsm_write_spec.py → python/python/tests/test_rte_spec.py
+69 -54
View File
@@ -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 RTEWriteSpec via
`Table.enable_rte` / `Table.disable_rte`.
"""
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 RTEWriteSpec
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_rte_spec_module_metadata():
assert lancedb.RTEWriteSpec is RTEWriteSpec
assert RTEWriteSpec.__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_rte_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_rte(RTEWriteSpec.bucket("id", 0))
with pytest.raises(Exception, match="num_buckets"):
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 1025))
table.enable_rte(RTEWriteSpec.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_rte(RTEWriteSpec.bucket("id", 4))
with pytest.raises(Exception, match="already set"):
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 8))
table.enable_rte(RTEWriteSpec.bucket("id", 8))
def test_unset_lsm_write_spec(tmp_path):
def test_disable_rte(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 RTE write spec"):
table.disable_rte()
# 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_rte(RTEWriteSpec.bucket("id", 4))
table.disable_rte()
# 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 RTE write spec"):
table.disable_rte()
table.enable_rte(RTEWriteSpec.bucket("id", 8))
def test_enable_rte_requires_primary_key(tmp_path):
_db, table = _make_table(tmp_path)
for spec in (
RTEWriteSpec.unsharded(),
RTEWriteSpec.bucket("id", 4),
RTEWriteSpec.identity("id"),
):
with pytest.raises(Exception, match="primary key"):
table.enable_rte(spec)
assert table.get_rte_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_rte(RTEWriteSpec.unsharded())
table.disable_rte()
def test_lsm_write_spec_repr():
s = LsmWriteSpec.bucket("id", 4)
def test_rte_spec_repr():
s = RTEWriteSpec.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 = RTEWriteSpec.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_rte_spec_with_maintained_indexes():
s = RTEWriteSpec.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_rte(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_rte(RTEWriteSpec.bucket("id", 4))
await table.disable_rte()
# 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 RTE write spec"):
await table.disable_rte()
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_rte(RTEWriteSpec.identity("v"))
table.disable_rte()
def test_lsm_write_spec_identity_and_writer_config_defaults():
s = LsmWriteSpec.identity("v")
def test_rte_spec_identity_and_writer_config_defaults():
s = RTEWriteSpec.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_rte_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_rte_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_rte(
RTEWriteSpec.bucket("id", 4)
.with_maintained_indexes([idx_name])
.with_writer_config_defaults({"durable_write": "false"})
)
spec = table.get_lsm_write_spec()
spec = table.get_rte_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_rte()
assert table.get_rte_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_rte(RTEWriteSpec.identity("id"))
spec = table.get_rte_spec()
assert spec.spec_type == "identity"
assert spec.column == "id"
assert spec.maintained_indexes is None
table.unset_lsm_write_spec()
table.disable_rte()
# 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_rte(RTEWriteSpec.unsharded().with_maintained_indexes([]))
spec = table.get_rte_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_rte_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_rte_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_rte(
RTEWriteSpec.bucket("id", 8).with_maintained_indexes([idx_name])
)
spec = await table.get_lsm_write_spec()
spec = await table.get_rte_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_rte()
assert await table.get_rte_spec() is None
python/python/tests/test_session.py
+10
View File
@@ -36,3 +36,13 @@ def test_session_cache_configuration(tmp_path):
assert final_cache_size > initial_cache_size # Cache should have grown
assert final_cache_items >= initial_cache_items # Items should not decrease
assert initial_cache_size < index_cache_size + metadata_cache_size
def test_connection_session_returns_supplied_session(tmp_path):
"""Sync connections must expose the session they were opened with."""
db = lancedb.connect(tmp_path)
assert db.session is None
session = lancedb.Session()
db = lancedb.connect(tmp_path, session=session)
assert db.session is session
python/python/tests/test_table.py
+443 -2
View File
@@ -12,10 +12,10 @@ import warnings
import weakref
from concurrent.futures import ThreadPoolExecutor
from datetime import date, datetime, timedelta
from decimal import Decimal
from decimal import Decimal, localcontext
from time import sleep
from typing import List
from unittest.mock import patch
from unittest.mock import MagicMock, patch
import lancedb
from lancedb import table as table_module
@@ -280,6 +280,30 @@ def test_create_table_infers_large_int_vectors(mem_db: DBConnection):
assert vector_column.to_pylist() == [[0.0, 300.0]]
def test_create_table_all_null_vector_column(mem_db: DBConnection):
# A vector column can legitimately arrive with every value null, e.g. when
# the embeddings have not been computed yet. The dimension is unknowable in
# that case, so schema inference must not crash and should leave the column
# as a variable-length list rather than a fixed-size one.
data = pa.table(
{
"id": pa.array([1, 2]),
"vector": pa.array([None, None], type=pa.list_(pa.float32())),
}
)
table = mem_db.create_table("all_null_vector", data=data)
vector_field = table.schema.field("vector")
assert pa.types.is_list(vector_field.type)
assert vector_field.type.value_type == pa.float32()
result = table.to_arrow()
assert result.num_rows == 2
assert result.column("vector").to_pylist() == [None, None]
assert table.count_rows() == 2
@pytest.mark.asyncio
async def test_create_table_async_infers_large_int_vectors(
mem_db_async: AsyncConnection,
@@ -518,6 +542,30 @@ def test_add(mem_db: DBConnection):
_add(table, schema)
def test_add_from_iterator_that_queries_table(mem_db: DBConnection):
source = mem_db.create_table("source", data=pa.table({"id": range(16)}))
target = mem_db.create_table("target", schema=source.schema)
def batches():
for _ in range(5):
yield source.search().limit(10).to_arrow()
target.add(batches())
assert target.count_rows() == 50
def test_create_table_from_iterator_that_queries_table(mem_db: DBConnection):
source = mem_db.create_table("source", data=pa.table({"id": range(16)}))
def batches():
yield source.search().limit(10).to_arrow()
target = mem_db.create_table("target", data=batches())
assert target.count_rows() == 10
def test_add_releases_arrow_buffers_without_gc(mem_db: DBConnection):
"""Regression test for https://github.com/lancedb/lancedb/issues/2512."""
schema = pa.schema([pa.field("x", pa.int64())])
@@ -1585,6 +1633,25 @@ def test_create_index_async_returns_done_job(mem_db: DBConnection):
job.cancel()
@pytest.mark.parametrize("method", ["create_index", "create_index_async"])
@pytest.mark.parametrize("wait_timeout", ["invalid", 1, []])
def test_create_index_invalid_wait_timeout(mem_db: DBConnection, method, wait_timeout):
table = mem_db.create_table("invalid_timeout", [{"id": 1}])
with pytest.raises(TypeError):
getattr(table, method)("id", config=BTree(), wait_timeout=wait_timeout)
assert table.list_indices() == []
@pytest.mark.parametrize("method", ["create_index", "create_index_async"])
@pytest.mark.parametrize("wait_timeout", [None, timedelta(seconds=1)])
def test_create_index_valid_wait_timeout(mem_db: DBConnection, method, wait_timeout):
table = mem_db.create_table("valid_timeout", [{"id": 1}])
result = getattr(table, method)("id", config=BTree(), wait_timeout=wait_timeout)
if result is not None:
result.wait()
assert len(table.list_indices()) == 1
@patch("lancedb.table.AsyncTable.create_index")
def test_create_index_method(mock_create_index, mem_db: DBConnection):
table = mem_db.create_table(
@@ -1725,6 +1792,51 @@ def test_create_index_method(mock_create_index, mem_db: DBConnection):
)
@pytest.mark.parametrize("api", ["legacy", "unified"])
@pytest.mark.parametrize("cuda_failure", ["unavailable", "init_error"])
def test_create_index_cuda_failure_suggests_cpu(
api, cuda_failure, mem_db: DBConnection
):
table = mem_db.create_table(
f"test_cuda_unavailable_{api}",
data=[{"vector": [3.1, 4.1]}, {"vector": [5.9, 26.5]}],
)
torch = MagicMock()
torch.cuda.is_available.return_value = cuda_failure == "init_error"
if cuda_failure == "init_error":
torch.cuda.init.side_effect = RuntimeError("CUDA driver initialization failed")
with (
patch.dict(sys.modules, {"torch": torch}),
patch.object(table, "to_lance") as to_lance,
):
if api == "legacy":
with (
pytest.warns(DeprecationWarning, match="create_index"),
pytest.raises(RuntimeError, match="accelerator=None.*CPU"),
):
table.create_index(
metric="dot",
num_partitions=1,
num_sub_vectors=1,
vector_column_name="vector",
accelerator="cuda",
)
else:
with pytest.raises(RuntimeError, match="accelerator=None.*CPU"):
table.create_index(
"vector",
config=IvfPq(
distance_type="dot",
num_partitions=1,
num_sub_vectors=1,
accelerator="cuda",
),
)
to_lance.return_value.create_index.assert_not_called()
@patch("lancedb.table.AsyncTable.create_index")
def test_create_index_name_and_train_parameters(
mock_create_index, mem_db: DBConnection
@@ -2677,6 +2789,51 @@ def test_update_with_arrow_scalar(mem_db: DBConnection):
assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0, 3.0, 4.0]]
@pytest.mark.parametrize("use_arrow_scalar", [False, True])
@pytest.mark.parametrize(
"decimal_type, value",
[
(pa.decimal128(10, 2), "9.99"),
(pa.decimal128(10, 2), "-9.99"),
(pa.decimal128(10, 2), "0.00"),
(pa.decimal128(10, 2), "1.23E+5"),
(pa.decimal128(38, 18), "1E-18"),
(pa.decimal128(38, 18), "12345678901234567.123456789012345678"),
(pa.decimal128(38, 18), "12345678901234567890.123456789012345678"),
(pa.decimal128(38, 18), "-12345678901234567890.123456789012345678"),
(
pa.decimal256(76, 38),
"12345678901234567890123456789012345678."
"12345678901234567890123456789012345678",
),
(
pa.decimal256(76, 38),
"-12345678901234567890123456789012345678."
"12345678901234567890123456789012345678",
),
],
)
def test_update_decimal(mem_db: DBConnection, decimal_type, value, use_arrow_scalar):
expected = Decimal(value)
table = mem_db.create_table(
"my_table",
pa.table(
{
"id": [1, 2],
"price": pa.array([Decimal("1"), expected], type=decimal_type),
}
),
)
update = table.to_arrow()["price"][1] if use_arrow_scalar else expected
with localcontext() as context:
context.prec = 6
result = table.update(where="id = 1", values={"price": update})
assert result.rows_updated == 1
actual = table.to_arrow()
assert actual.schema.field("price").type == decimal_type
assert actual["price"].to_pylist() == [expected, expected]
def test_update_types(mem_db: DBConnection):
table = mem_db.create_table(
"my_table",
@@ -3693,6 +3850,42 @@ def test_multiple_vector_columns(mem_db: DBConnection):
assert result1["text"][0] != result2["text"][0]
@pytest.mark.parametrize(
"query",
[
(0.0, 0.0),
pa.array([0.0, 0.0]),
pa.chunked_array([[0.0], [0.0]]),
],
)
@pytest.mark.parametrize("query_type", ["auto", "vector"])
def test_sync_search_accepts_supported_vector_containers(
mem_db: DBConnection, query, query_type
):
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field("small", pa.list_(pa.float32(), 2)),
pa.field("large", pa.list_(pa.float32(), 3)),
]
)
table = mem_db.create_table(
"supported_vector_containers",
data=pa.table(
{
"id": [1, 2],
"small": [[0.0, 0.0], [1.0, 1.0]],
"large": [[1.0, 1.0, 1.0], [0.0, 0.0, 0.0]],
},
schema=schema,
),
)
result = table.search(query, query_type=query_type).limit(1).to_arrow()
assert result["id"].to_pylist() == [1]
def test_create_scalar_index(mem_db: DBConnection):
vec_array = pa.array(
[[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]], pa.list_(pa.float32(), 2)
@@ -4036,6 +4229,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(
@@ -4378,6 +4591,58 @@ async def test_alter_columns_rejects_invalid_async(
assert await table.schema() == initial_schema
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("list_type", [pa.list_, pa.large_list])
@pytest.mark.parametrize(
"value_type, values, inferred_type",
[
(pa.float64(), [1.25, 2.5], pa.float32()),
(pa.int32(), [0, 255], pa.uint8()),
(pa.int32(), [-1, 256], pa.float32()),
],
)
async def test_inferred_vector_field_metadata(
tmp_db: DBConnection,
tmp_db_async: AsyncConnection,
asynchronous,
list_type,
value_type,
values,
inferred_type,
):
metadata = {
"lancedb:description": "Embedding values",
"lancedb:tag:model": "test-model",
}
schema = pa.schema(
[
pa.field(
"vector", list_type(value_type), nullable=False, metadata=metadata
),
pa.field("text", pa.string(), metadata={"description": "Source text"}),
],
metadata={"description": "Source dataset"},
)
rows = [{"vector": values, "text": "first"}, {"vector": values, "text": "second"}]
data = pa.Table.from_pylist(rows, schema=schema)
if asynchronous:
await tmp_db_async.create_table("vector_metadata", data=data)
table = await tmp_db_async.open_table("vector_metadata")
actual_schema = await table.schema()
actual_data = await table.to_arrow()
else:
tmp_db.create_table("vector_metadata", data=data)
table = tmp_db.open_table("vector_metadata")
actual_schema = table.schema
actual_data = table.to_arrow()
expected_schema = schema.set(
0, schema.field("vector").with_type(pa.list_(inferred_type, 2))
)
assert actual_schema.equals(expected_schema, check_metadata=True)
assert actual_data.to_pylist() == rows
def test_update_field_metadata(mem_db: DBConnection):
data = pa.table({"id": [0, 1], "category": ["a", "b"]})
table = mem_db.create_table("my_table", data=data)
@@ -4592,6 +4857,57 @@ def test_optimize_warns_on_zero_cleanup_sync(mem_db: DBConnection):
assert all("concurrent" in str(w.message) for w in seen)
def test_optimize_retains_tagged_versions(tmp_db: DBConnection):
# A tagged version older than the cleanup cutoff must be retained, not make
# optimize() fail part-way (after compaction, before the index step).
table = tmp_db.create_table(
"test", data=pa.table({"id": list(range(100)), "text": ["a"] * 100})
)
table.create_index("id", config=BTree())
table.tags.create("release-1", table.version)
tagged_version = table.version
for i in range(3):
table.add(pa.table({"id": [1000 + i], "text": ["new"]}))
assert table.stats()["fragment_stats"]["num_fragments"] == 4
assert table.index_stats("id_idx").num_unindexed_rows == 3
sleep(1.1)
table.optimize(cleanup_older_than=timedelta(seconds=1))
# The index step ran after the prune step
assert table.index_stats("id_idx").num_unindexed_rows == 0
# It stays optimizable
table.optimize(cleanup_older_than=timedelta(seconds=1))
assert table.stats()["fragment_stats"]["num_fragments"] == 1
# Untagged old versions were pruned, the tagged one was retained
versions = [v["version"] for v in table.list_versions()]
assert tagged_version in versions
for added in range(tagged_version + 1, tagged_version + 4):
assert added not in versions
assert len(table) == 103
table.checkout("release-1")
assert len(table) == 100
assert table.to_arrow().num_rows == 100
@pytest.mark.asyncio
async def test_optimize_retains_tagged_versions_async(tmp_db_async: AsyncConnection):
table = await tmp_db_async.create_table("test", data=[{"x": [1]}])
await table.tags.create("keep", await table.version())
await table.add(data=[{"x": [2]}])
await table.add(data=[{"x": [3]}])
sleep(1.1)
stats = await table.optimize(cleanup_older_than=timedelta(seconds=1))
assert stats.prune.old_versions_removed >= 1
await table.checkout("keep")
assert await table.count_rows() == 1
assert await table.query().to_arrow() == pa.table({"x": [[1]]})
def test_replace_field_metadata(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("my_table", data=[{"x": 0}])
@@ -4647,6 +4963,81 @@ def test_stats(mem_db: DBConnection):
assert stats_with_index["fragment_stats"]["num_fragments"] == 1
@pytest.mark.parametrize("iterable_type", [iter, tuple])
@pytest.mark.parametrize("empty", [False, True])
@pytest.mark.parametrize("with_schema", [False, True])
def test_create_table_recordbatch_iterable(
mem_db: DBConnection, iterable_type, empty, with_schema
):
schema = pa.schema([pa.field("id", pa.int64())])
batches = [] if empty else [pa.RecordBatch.from_pydict({"id": [1]}, schema)]
data = iterable_type(batches)
if empty and not with_schema:
with pytest.raises(ValueError, match="empty iterable without a schema"):
mem_db.create_table("test", data=data)
return
table = mem_db.create_table(
"test", data=data, schema=schema if with_schema else None
)
assert table.schema == schema
assert table.to_arrow().to_pylist() == ([] if empty else [{"id": 1}])
@pytest.mark.asyncio
@pytest.mark.parametrize("iterable_type", [iter, tuple])
@pytest.mark.parametrize("empty", [False, True])
@pytest.mark.parametrize("with_schema", [False, True])
async def test_create_table_recordbatch_iterable_async(
mem_db_async: AsyncConnection, iterable_type, empty, with_schema
):
schema = pa.schema([pa.field("id", pa.int64())])
batches = [] if empty else [pa.RecordBatch.from_pydict({"id": [1]}, schema)]
data = iterable_type(batches)
if empty and not with_schema:
with pytest.raises(ValueError, match="empty iterable without a schema"):
await mem_db_async.create_table("test", data=data)
return
table = await mem_db_async.create_table(
"test", data=data, schema=schema if with_schema else None
)
assert await table.schema() == schema
assert (await table.to_arrow()).to_pylist() == ([] if empty else [{"id": 1}])
@pytest.mark.parametrize("iterable_type", [iter, tuple])
@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
def test_add_empty_recordbatch_iterable(
mem_db: DBConnection, iterable_type, on_bad_vectors
):
table = mem_db.create_table("test", [{"id": 1}])
if on_bad_vectors == "error":
# The default scannable path still requires an input schema. This
# fix does not change its validation policy for schema-less input.
with pytest.raises(ValueError, match="empty iterator"):
table.add(iterable_type([]), on_bad_vectors=on_bad_vectors)
else:
# Sanitization has the existing table schema available.
table.add(iterable_type([]), on_bad_vectors=on_bad_vectors)
assert table.to_arrow().to_pylist() == [{"id": 1}]
@pytest.mark.asyncio
@pytest.mark.parametrize("iterable_type", [iter, tuple])
@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
async def test_add_empty_recordbatch_iterable_async(
mem_db_async: AsyncConnection, iterable_type, on_bad_vectors
):
table = await mem_db_async.create_table("test", [{"id": 1}])
if on_bad_vectors == "error":
with pytest.raises(ValueError, match="empty iterator"):
await table.add(iterable_type([]), on_bad_vectors=on_bad_vectors)
else:
await table.add(iterable_type([]), on_bad_vectors=on_bad_vectors)
assert (await table.to_arrow()).to_pylist() == [{"id": 1}]
def test_create_table_empty_list_with_schema(mem_db: DBConnection):
"""Test creating table with empty list data and schema
@@ -4869,3 +5260,53 @@ async def test_refresh_column_async_job_async_table(tmp_path):
assert result.published_version == 4
assert await job.status() == "finished"
assert (await table.to_arrow())["tripled"].to_pylist() == [9]
@pytest.mark.parametrize("interval", [None, 300])
def test_computed_auto_refresh_options(tmp_path, interval):
db = lancedb.connect(tmp_path)
table = db.create_table("refresh_options", [{"x": 1}])
before = table.version
table.add_columns(
computed={"doubled": "x * 2", "tripled": "x * 3"},
auto_refresh_enabled=False,
auto_refresh_interval_seconds=interval,
)
assert table.version == before + 1
for column in ("doubled", "tripled"):
metadata = table.schema.field(column).metadata
assert metadata[b"computed_column.auto_refresh_enabled"] == b"false"
if interval is None:
assert b"computed_column.auto_refresh_interval_seconds" not in metadata
else:
assert metadata[b"computed_column.auto_refresh_interval_seconds"] == b"300"
table.refresh_column("doubled")
assert table.to_arrow()["doubled"].to_pylist() == [2]
@pytest.mark.parametrize("value", [0, -1, True, 1.5, "300", 2**64])
def test_computed_auto_refresh_rejects_invalid_interval(tmp_path, value):
table = lancedb.connect(tmp_path).create_table("invalid_options", [{"x": 1}])
with pytest.raises(ValueError, match="auto_refresh_interval_seconds"):
table.add_columns(
computed={"doubled": "x * 2"}, auto_refresh_interval_seconds=value
)
assert table.schema.names == ["x"]
def test_auto_refresh_options_require_computed_columns(tmp_path):
table = lancedb.connect(tmp_path).create_table("invalid_kind", [{"x": 1}])
with pytest.raises(ValueError, match="require computed"):
table.add_columns({"doubled": "x * 2"}, auto_refresh_enabled=True)
@pytest.mark.asyncio
async def test_async_computed_auto_refresh_uses_server_interval(tmp_path):
db = await lancedb.connect_async(tmp_path)
table = await db.create_table("server_interval", [{"x": 1}])
before = await table.version()
await table.add_columns(computed={"doubled": "x * 2"})
assert await table.version() == before + 1
metadata = (await table.schema()).field("doubled").metadata
assert metadata[b"computed_column.auto_refresh_enabled"] == b"true"
assert b"computed_column.auto_refresh_interval_seconds" not in metadata
python/python/tests/test_util.py
+145 -1
View File
@@ -4,6 +4,7 @@
import os
import pathlib
from decimal import Decimal, localcontext
from typing import Optional
import lance
@@ -28,7 +29,14 @@ 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_location,
get_uri_scheme,
inf_vector_column_query,
join_uri,
value_to_sql,
)
from utils import exception_output
@@ -80,6 +88,22 @@ def test_normalize_uri():
assert parsed_scheme == expected_scheme
def test_get_uri_scheme_and_location_pathlib():
# Test pathlib.Path instances
abs_path = pathlib.Path("/tmp/data")
rel_path = pathlib.Path("relative/data")
assert get_uri_scheme(abs_path) == "file"
assert get_uri_scheme(rel_path) == "file"
assert get_uri_location(abs_path) == str(abs_path)
assert get_uri_location(rel_path) == str(rel_path)
# String URIs still behave identically
assert get_uri_scheme("s3://bucket/data") == "s3"
assert get_uri_location("s3://bucket/data") == "bucket/data"
def test_join_uri_remote():
schemes = ["s3", "az", "gs"]
for scheme in schemes:
@@ -158,6 +182,41 @@ def test_value_to_sql_string(tmp_path):
assert table.to_pandas().query("search == @value")["replace"].item() == value
@pytest.mark.parametrize(
"value, expected",
[
("9.99", "'9.99'"),
("-9.99", "'-9.99'"),
("0", "'0'"),
("-0.00", "'-0.00'"),
("1.23E+5", "'123000'"),
("1E-18", "'0.000000000000000001'"),
(
"12345678901234567890.123456789012345678",
"'12345678901234567890.123456789012345678'",
),
(
"12345678901234567890123456789012345678."
"12345678901234567890123456789012345678",
"'12345678901234567890123456789012345678."
"12345678901234567890123456789012345678'",
),
],
)
def test_value_to_sql_decimal(value, expected):
with localcontext() as context:
context.prec = 6
assert value_to_sql(Decimal(value)) == expected
@pytest.mark.parametrize("value", ["NaN", "sNaN", "NaN123456", "Infinity", "-Infinity"])
def test_value_to_sql_decimal_non_finite(value):
with pytest.raises(
ValueError, match="^Non-finite Decimal values cannot be converted to SQL$"
):
value_to_sql(Decimal(value))
def test_value_to_sql_dict():
# Simple flat struct
assert value_to_sql({"a": 1, "b": "hello"}) == "named_struct('a', 1, 'b', 'hello')"
@@ -471,6 +530,20 @@ def test_sanitize_vectors_with_only_null_elements(on_bad_vectors):
assert sanitize().to_pylist() == expected[on_bad_vectors]
@pytest.mark.parametrize("dtype", [pa.float16(), pa.float32(), pa.float64()])
@pytest.mark.parametrize("large", [False, True])
@pytest.mark.parametrize("chunked", [False, True])
def test_has_nan_values_skips_hidden_children_of_null_lists(dtype, large, chunked):
array_type = pa.LargeListArray if large else pa.ListArray
values = pa.array([float("nan"), None, float("nan"), 3.0, 4.0, 5.0], type=dtype)
arr = array_type.from_arrays(
[0, 2, 4, 6], values, mask=pa.array([True, False, False])
)
if chunked:
arr = pa.chunked_array([arr.slice(0, 2), arr.slice(2, 1)])
assert has_nan_values(arr).to_pylist() == [False, True, False]
def test_handle_bad_vectors_noop():
# ChunkedArray should be preserved as-is
vector = pa.chunked_array(
@@ -992,6 +1065,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(
@@ -1223,3 +1329,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)
python/python/tests/test_util_update_values.py
+215
View File
@@ -0,0 +1,215 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
from datetime import date, datetime
import lancedb
import numpy as np
import pyarrow as pa
import pytest
ARRAY_VALUES = [
((1.0, 2.0), [1.0, 2.0]),
([1.0, 2.0], [1.0, 2.0]),
(np.array([1.0, 2.0]), [1.0, 2.0]),
(None, None),
]
@pytest.mark.parametrize("value,expected", ARRAY_VALUES)
def test_update_array_values(tmp_path, value, expected):
dtype = pa.list_(pa.float64())
db = lancedb.connect(tmp_path)
table = db.create_table(
"values", pa.table({"array": pa.array([[0.0, 0.0]], type=dtype)})
)
table.update(values={"array": value})
assert table.to_arrow()["array"].to_pylist() == [expected]
@pytest.mark.asyncio
@pytest.mark.parametrize("value,expected", ARRAY_VALUES)
async def test_update_array_values_async(tmp_path, value, expected):
dtype = pa.list_(pa.float64())
db = await lancedb.connect_async(tmp_path)
table = await db.create_table(
"values", pa.table({"array": pa.array([[0.0, 0.0]], type=dtype)})
)
await table.update({"array": value})
assert (await table.to_arrow())["array"].to_pylist() == [expected]
TIME_VALUES = [
np.datetime64("2025", "Y"),
np.datetime64("2025-01", "M"),
np.datetime64("2025-01-02", "D"),
np.datetime64("2025-01-02T03", "h"),
np.datetime64("2025-01-02T03:04", "m"),
np.datetime64("2025-01-02T03:04:05", "s"),
np.datetime64("2025-01-02T03:04:05.123", "ms"),
np.datetime64("2025-01-02T03:04:05.123456", "us"),
np.datetime64("2025-01-02T03:04:05.123456789", "ns"),
np.datetime64("NaT", "ns"),
datetime(2025, 1, 2, 3, 4, 5, 123456),
None,
]
def expected_time(value):
if value is None or isinstance(value, np.datetime64) and np.isnat(value):
return None
return int(np.datetime64(value, "ns").astype(np.int64))
@pytest.mark.parametrize("value", TIME_VALUES)
def test_update_timestamp_values(tmp_path, value):
db = lancedb.connect(tmp_path)
data = pa.table({"time": pa.array([datetime(2025, 1, 1)], type=pa.timestamp("ns"))})
table = db.create_table("times", data)
table.update(values={"time": value})
assert table.to_arrow()["time"].cast(pa.int64()).to_pylist() == [
expected_time(value)
]
@pytest.mark.asyncio
@pytest.mark.parametrize("value", TIME_VALUES)
async def test_update_timestamp_values_async(tmp_path, value):
db = await lancedb.connect_async(tmp_path)
data = pa.table({"time": pa.array([datetime(2025, 1, 1)], type=pa.timestamp("ns"))})
table = await db.create_table("times", data)
await table.update({"time": value})
assert (await table.to_arrow())["time"].cast(pa.int64()).to_pylist() == [
expected_time(value)
]
@pytest.mark.parametrize("values", [[[1, 2], [3, 4]], []])
def test_nested_array_tuple_matches_list(tmp_path, values):
db = lancedb.connect(tmp_path)
data = pa.table(
{"array": pa.array([[[0, 0]]], type=pa.list_(pa.list_(pa.int64())))}
)
control = db.create_table("list_values", data)
candidate = db.create_table("tuple_values", data)
tuples = tuple(tuple(value) for value in values)
try:
control.update(values={"array": values})
except ValueError as error:
reason = (
"Expected a literal value in array" if values else "concat requires input"
)
assert reason in str(error)
with pytest.raises(ValueError, match=reason):
candidate.update(values={"array": tuples})
assert control.to_arrow().equals(data)
assert candidate.to_arrow().equals(data)
else:
candidate.update(values={"array": tuples})
assert candidate.to_arrow().equals(control.to_arrow())
@pytest.mark.asyncio
@pytest.mark.parametrize("values", [[[1, 2], [3, 4]], []])
async def test_nested_array_tuple_matches_list_async(tmp_path, values):
db = await lancedb.connect_async(tmp_path)
data = pa.table(
{"array": pa.array([[[0, 0]]], type=pa.list_(pa.list_(pa.int64())))}
)
control = await db.create_table("list_values", data)
candidate = await db.create_table("tuple_values", data)
tuples = tuple(tuple(value) for value in values)
try:
await control.update({"array": values})
except ValueError as error:
reason = (
"Expected a literal value in array" if values else "concat requires input"
)
assert reason in str(error)
with pytest.raises(ValueError, match=reason):
await candidate.update({"array": tuples})
assert (await control.to_arrow()).equals(data)
assert (await candidate.to_arrow()).equals(data)
else:
await candidate.update({"array": tuples})
assert (await candidate.to_arrow()).equals(await control.to_arrow())
@pytest.mark.parametrize("unit", ["ps", "fs", "as"])
def test_subnanosecond_update_does_not_round(tmp_path, unit):
db = lancedb.connect(tmp_path)
data = pa.table({"time": pa.array([datetime(1970, 1, 1)], type=pa.timestamp("ns"))})
table = db.create_table("precise", data)
with pytest.raises(ValueError, match="cannot preserve sub-nanosecond precision"):
table.update(values={"time": np.datetime64(1, unit)})
assert table.to_arrow()["time"].cast(pa.int64()).to_pylist() == [0]
@pytest.mark.asyncio
@pytest.mark.parametrize("unit", ["ps", "fs", "as"])
async def test_subnanosecond_update_does_not_round_async(tmp_path, unit):
db = await lancedb.connect_async(tmp_path)
data = pa.table({"time": pa.array([datetime(1970, 1, 1)], type=pa.timestamp("ns"))})
table = await db.create_table("precise", data)
with pytest.raises(ValueError, match="cannot preserve sub-nanosecond precision"):
await table.update({"time": np.datetime64(1, unit)})
assert (await table.to_arrow())["time"].cast(pa.int64()).to_pylist() == [0]
@pytest.mark.parametrize("unit", ["Y", "M", "W", "D"])
def test_update_numpy_date_column(tmp_path, unit):
value = np.datetime64("2025-01-02", unit)
db = lancedb.connect(tmp_path)
table = db.create_table(
"dates", pa.table({"day": pa.array([date(1970, 1, 1)], type=pa.date32())})
)
table.update(values={"day": value})
assert table.to_arrow()["day"].to_pylist() == [value.astype("datetime64[D]").item()]
@pytest.mark.asyncio
@pytest.mark.parametrize("unit", ["Y", "M", "W", "D"])
async def test_update_numpy_date_column_async(tmp_path, unit):
value = np.datetime64("2025-01-02", unit)
db = await lancedb.connect_async(tmp_path)
table = await db.create_table(
"dates", pa.table({"day": pa.array([date(1970, 1, 1)], type=pa.date32())})
)
await table.update({"day": value})
assert (await table.to_arrow())["day"].to_pylist() == [
value.astype("datetime64[D]").item()
]
@pytest.mark.parametrize("timezone", [None, "UTC", "America/New_York", "Asia/Tokyo"])
@pytest.mark.parametrize(
"value",
[np.datetime64("2025-01-02T03:04:05.123456789", "ns"), np.datetime64("NaT", "ns")],
)
def test_update_numpy_timestamp_matches_arrow_ingestion(tmp_path, timezone, value):
dtype = pa.timestamp("ns", tz=timezone)
ingested = pa.array(np.array([value]), type=dtype)
db = lancedb.connect(tmp_path)
table = db.create_table("timestamps", pa.table({"time": ingested}))
expected = table.to_arrow()["time"].cast(pa.int64()).to_pylist()
table.update(values={"time": value})
assert table.to_arrow()["time"].cast(pa.int64()).to_pylist() == expected
@pytest.mark.asyncio
@pytest.mark.parametrize("timezone", [None, "UTC", "America/New_York", "Asia/Tokyo"])
@pytest.mark.parametrize(
"value",
[np.datetime64("2025-01-02T03:04:05.123456789", "ns"), np.datetime64("NaT", "ns")],
)
async def test_update_numpy_timestamp_matches_arrow_ingestion_async(
tmp_path, timezone, value
):
dtype = pa.timestamp("ns", tz=timezone)
ingested = pa.array(np.array([value]), type=dtype)
db = await lancedb.connect_async(tmp_path)
table = await db.create_table("timestamps", pa.table({"time": ingested}))
expected = (await table.to_arrow())["time"].cast(pa.int64()).to_pylist()
await table.update({"time": value})
assert (await table.to_arrow())["time"].cast(pa.int64()).to_pylist() == expected
python/python/tests/test_voyageai_embeddings.py
+57
View File
@@ -6,6 +6,8 @@
These tests verify model registration and configuration without requiring API calls.
"""
import numpy as np
import pyarrow as pa
import pytest
from unittest.mock import MagicMock, patch
@@ -122,3 +124,58 @@ class TestVoyageAIModelRegistration:
assert "voyage-4" not in func.multimodal_embedding_models
assert "voyage-4-lite" not in func.multimodal_embedding_models
assert "voyage-4-large" not in func.multimodal_embedding_models
@pytest.mark.parametrize(
"inputs",
[
["first", "second"],
np.array(["first", "second"]),
pa.array(["first", "second"]),
pa.chunked_array([["first"], ["second"]]),
],
ids=["list", "numpy", "arrow", "chunked-arrow"],
)
def test_text_source_embeddings_accept_batch_containers(inputs):
from lancedb.embeddings.voyageai import VoyageAIEmbeddingFunction
client = MagicMock()
client.tokenize.return_value = [[1], [2]]
client.embed.return_value.embeddings = [[0.1, 0.2], [0.3, 0.4]]
VoyageAIEmbeddingFunction.client = client
func = get_registry().get("voyageai").create(name="voyage-3")
assert func.compute_source_embeddings(inputs) == [[0.1, 0.2], [0.3, 0.4]]
client.embed.assert_called_once_with(
texts=["first", "second"], model="voyage-3", input_type="document"
)
@pytest.mark.parametrize("inputs", [[], np.array([], dtype=str)])
def test_text_source_embeddings_accept_empty_batch(inputs):
from lancedb.embeddings.voyageai import VoyageAIEmbeddingFunction
client = MagicMock()
VoyageAIEmbeddingFunction.client = client
func = get_registry().get("voyageai").create(name="voyage-3")
assert func.compute_source_embeddings(inputs) == []
client.tokenize.assert_not_called()
client.embed.assert_not_called()
@pytest.mark.parametrize(
"inputs",
[["text", 42], np.array(["text", None]), np.array([["first", "second"]])],
)
def test_text_source_embeddings_reject_non_text_batch_entries(inputs):
from lancedb.embeddings.voyageai import VoyageAIEmbeddingFunction
client = MagicMock()
VoyageAIEmbeddingFunction.client = client
func = get_registry().get("voyageai").create(name="voyage-3")
with pytest.raises(ValueError, match="Each input should be str"):
func.compute_source_embeddings(inputs)
client.tokenize.assert_not_called()
client.embed.assert_not_called()
python/python/tests/test_voyageai_image_inputs.py
+88
View File
@@ -0,0 +1,88 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import base64
from io import BytesIO
from unittest.mock import MagicMock
import pytest
from lancedb.embeddings import get_registry
from lancedb.embeddings.voyageai import VoyageAIEmbeddingFunction
@pytest.mark.parametrize("model_name", ["voyage-multimodal-3", "voyage-multimodal-3.5"])
@pytest.mark.parametrize("input_kind", ["pil", "bytes", "path"])
@pytest.mark.parametrize(
"mode,color,expected_rgb",
[
("RGBA", (64, 128, 192, 128), (64, 128, 192)),
("LA", (96, 128), (96, 96, 96)),
("P", 0, (64, 128, 192)),
("RGB", (64, 128, 192), (64, 128, 192)),
],
)
def test_multimodal_source_embeddings_encode_image_modes(
model_name, input_kind, mode, color, expected_rgb, tmp_path, monkeypatch
):
Image = pytest.importorskip("PIL.Image")
image = Image.new(mode, (8, 6), color)
if mode == "P":
image.putpalette([64, 128, 192] + [0, 0, 0] * 255)
image.info["transparency"] = 0
original_pixels = image.tobytes()
image_input = image
if input_kind != "pil":
png = BytesIO()
image.save(png, format="PNG")
if input_kind == "bytes":
image_input = png.getvalue()
else:
image_input = tmp_path / "image.png"
image_input.write_bytes(png.getvalue())
client = MagicMock()
client.multimodal_embed.return_value.embeddings = [[0.1, 0.2]]
monkeypatch.setattr(VoyageAIEmbeddingFunction, "client", client)
func = get_registry().get("voyageai").create(name=model_name)
assert func.compute_source_embeddings(image_input) == [[0.1, 0.2]]
client.multimodal_embed.assert_called_once()
request = client.multimodal_embed.call_args.kwargs
assert request["model"] == model_name
assert request["input_type"] == "document"
assert len(request["inputs"]) == 1
(content,) = request["inputs"][0]["content"]
assert content["type"] == "image_base64"
prefix, encoded = content["image_base64"].split(",", 1)
assert prefix == "data:image/jpeg;base64"
with Image.open(BytesIO(base64.b64decode(encoded, validate=True))) as decoded:
assert decoded.format == "JPEG"
assert decoded.mode == "RGB"
assert decoded.size == (8, 6)
assert decoded.getpixel((0, 0)) == pytest.approx(expected_rgb, abs=3)
assert image.mode == mode
assert image.tobytes() == original_pixels
client.embed.assert_not_called()
@pytest.mark.parametrize("input_kind", ["bytes", "path"])
def test_multimodal_source_embeddings_reject_invalid_image(
input_kind, tmp_path, monkeypatch
):
PIL = pytest.importorskip("PIL")
image_input = b"not an image"
if input_kind == "path":
image_input = tmp_path / "invalid.png"
image_input.write_bytes(b"not an image")
client = MagicMock()
monkeypatch.setattr(VoyageAIEmbeddingFunction, "client", client)
func = get_registry().get("voyageai").create(name="voyage-multimodal-3")
with pytest.raises(PIL.UnidentifiedImageError):
func.compute_source_embeddings(image_input)
client.multimodal_embed.assert_not_called()
client.embed.assert_not_called()
python/python/tests/test_wal_hybrid.py
+10 -8
View File
@@ -14,7 +14,7 @@ from lancedb.rerankers import RRFReranker
def _wal_table(schema: pa.Schema) -> mock.Mock:
"""A MemWAL-backed table a previous query has already been refused
"""An RTE-backed table a previous query has already been refused
`_rowid` on, so the fusion goes straight to the primary key."""
table = mock.Mock()
table._hybrid_pk_fusion_learned.return_value = True
@@ -51,7 +51,7 @@ def test_wal_hybrid_projects_the_key_instead_of_asking_for_row_ids():
assert builder._pk_fusion.injected == ["id"]
for leg in (builder._vector_query, builder._fts_query):
assert leg._columns == ["text", "id"]
# A MemWAL table rejects with_row_id before it plans.
# An RTE table rejects with_row_id before it plans.
assert leg._with_row_id is not True
@@ -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_rte_false_joins_on_row_ids():
builder = _wal_hybrid().use_rte(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_rte"):
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_rte"):
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_rte") as info:
with_pk_fallback(
run, mock.Mock(), fused=lambda: False, caller_requested_row_id=True
)
python/python/tests/utils.py
+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_rte_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 RTE write path.
return self._json(request, {"rte_spec": None})
if path == f"/v1/table/{self.name}/count_rows/":
self._read_body(request)
return self._json(request, self.num_rows)
python/src/authz.rs
+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)
}
}
python/src/catalog.rs
+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()
python/src/connection.rs
+37 -5
View File
@@ -49,6 +49,18 @@ impl Connection {
}
}
fn parse_read_consistency_interval(interval: Option<f64>) -> PyResult<Option<Duration>> {
interval
.map(|seconds| {
Duration::try_from_secs_f64(seconds).map_err(|_| {
PyValueError::new_err(
"read_consistency_interval must be finite, non-negative, and within the supported range",
)
})
})
.transpose()
}
fn parse_namespace_client_pushdown_operations(
operations: Option<Vec<String>>,
) -> PyResult<HashSet<NamespaceClientPushdownOperation>> {
@@ -412,7 +424,8 @@ impl Connection {
})
}
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))]
#[allow(clippy::too_many_arguments)]
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false, auto_refresh_enabled=None, auto_refresh_interval_seconds=None))]
pub fn create_materialized_view(
self_: PyRef<'_, Self>,
name: String,
@@ -421,10 +434,18 @@ impl Connection {
filter: Option<String>,
limit: Option<u64>,
with_no_data: bool,
auto_refresh_enabled: Option<bool>,
auto_refresh_interval_seconds: Option<f64>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let mut builder = inner.create_materialized_view(name, source);
if let Some(enabled) = auto_refresh_enabled {
builder = builder.auto_refresh_enabled(enabled);
}
if let Some(seconds) = auto_refresh_interval_seconds {
builder = builder.auto_refresh_interval_seconds(seconds);
}
if let Some(projections) = projections {
builder = builder.select(projections);
}
@@ -440,7 +461,8 @@ impl Connection {
})
}
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))]
#[allow(clippy::too_many_arguments)]
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false, auto_refresh_enabled=None, auto_refresh_interval_seconds=None))]
pub fn create_materialized_view_async(
self_: PyRef<'_, Self>,
name: String,
@@ -449,10 +471,18 @@ impl Connection {
filter: Option<String>,
limit: Option<u64>,
with_no_data: bool,
auto_refresh_enabled: Option<bool>,
auto_refresh_interval_seconds: Option<f64>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let mut builder = inner.create_materialized_view(name, source);
if let Some(enabled) = auto_refresh_enabled {
builder = builder.auto_refresh_enabled(enabled);
}
if let Some(seconds) = auto_refresh_interval_seconds {
builder = builder.auto_refresh_interval_seconds(seconds);
}
if let Some(projections) = projections {
builder = builder.select(projections);
}
@@ -1059,6 +1089,7 @@ pub fn connect(
namespace_client_properties: Option<HashMap<String, String>>,
oauth_config: Option<crate::oauth::PyOAuthConfig>,
) -> PyResult<Bound<'_, PyAny>> {
let read_consistency_interval = parse_read_consistency_interval(read_consistency_interval)?;
future_into_py(py, async move {
let mut builder = lancedb::connect(&uri);
if let Some(api_key) = api_key {
@@ -1077,7 +1108,6 @@ pub fn connect(
#[cfg(not(feature = "remote"))]
let _ = sql_host_override;
if let Some(read_consistency_interval) = read_consistency_interval {
let read_consistency_interval = Duration::from_secs_f64(read_consistency_interval);
builder = builder.read_consistency_interval(read_consistency_interval);
}
if let Some(storage_options) = storage_options {
@@ -1126,7 +1156,7 @@ pub fn connect_namespace_client(
namespace_client_impl: Option<String>,
namespace_client_properties: Option<HashMap<String, String>>,
) -> PyResult<Connection> {
let read_consistency_interval = read_consistency_interval.map(Duration::from_secs_f64);
let read_consistency_interval = parse_read_consistency_interval(read_consistency_interval)?;
let namespace_client_pushdown_operations =
parse_namespace_client_pushdown_operations(namespace_client_pushdown_operations)?;
let ns_properties = namespace_client_properties.unwrap_or_default();
@@ -1183,7 +1213,7 @@ pub fn connect_namespace(
session: Option<crate::session::Session>,
namespace_client_pushdown_operations: Option<Vec<String>>,
) -> PyResult<Connection> {
let read_consistency_interval = read_consistency_interval.map(Duration::from_secs_f64);
let read_consistency_interval = parse_read_consistency_interval(read_consistency_interval)?;
let namespace_client_pushdown_operations =
parse_namespace_client_pushdown_operations(namespace_client_pushdown_operations)?;
@@ -1220,6 +1250,7 @@ pub struct PyClientConfig {
user_agent: String,
retry_config: Option<PyClientRetryConfig>,
timeout_config: Option<PyClientTimeoutConfig>,
blob_request_concurrency: Option<usize>,
extra_headers: Option<HashMap<String, String>>,
id_delimiter: Option<String>,
tls_config: Option<PyClientTlsConfig>,
@@ -1305,6 +1336,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),
python/src/expr.rs
+48 -10
View File
@@ -240,22 +240,60 @@ pub fn expr_lit(value: Bound<'_, PyAny>) -> PyResult<PyExpr> {
fn parse_decimal(s: &str) -> PyResult<(i128, u8, i8)> {
let s = s.trim();
let dot_pos = s.find('.');
let scale = if let Some(pos) = dot_pos {
(s.len() - pos - 1) as i8
} else {
0
let (mantissa, exponent) = match s.split_once(['e', 'E']) {
Some((mantissa, exponent)) => (
mantissa,
exponent.parse::<i64>().map_err(|e| {
PyValueError::new_err(format!("failed to parse decimal exponent: {}", e))
})?,
),
None => (s, 0),
};
let fractional_digits = mantissa.find('.').map_or(0, |pos| mantissa.len() - pos - 1);
// Keep the coefficient intact instead of expanding an exponent into zeros.
// Arrow Decimal128 has at most 38 digits and an i8 scale (at most 38).
let mut scale = i64::try_from(fractional_digits)
.ok()
.and_then(|fractional_digits| fractional_digits.checked_sub(exponent))
.and_then(|scale| i8::try_from(scale).ok())
.filter(|scale| *scale <= 38)
.ok_or_else(|| PyValueError::new_err("decimal scale is outside Decimal128 range"))?;
let digits = s.replace('.', "");
let val = digits
let digits = mantissa.replace('.', "");
let mut val = digits
.parse::<i128>()
.map_err(|e| PyValueError::new_err(format!("failed to parse decimal digits: {}", e)))?;
// Precision is total number of digits
let precision = digits.trim_start_matches('-').len() as u8;
// Positive scale cannot exceed precision, including for tiny values such as 1E-38.
let mut precision = digits
.trim_start_matches('-')
.len()
.max(usize::try_from(scale).unwrap_or(0));
if scale < 0 {
// Arrow cannot cast SQL text to a decimal with negative scale. Normalize
// positive exponents within Decimal128's bounded coefficient range.
if val != 0 {
let shift = -i16::from(scale) as u32;
precision += shift as usize;
if precision > 38 {
return Err(PyValueError::new_err(
"decimal precision exceeds Decimal128 maximum of 38",
));
}
val = 10_i128
.checked_pow(shift)
.and_then(|factor| val.checked_mul(factor))
.ok_or_else(|| PyValueError::new_err("decimal coefficient exceeds i128 range"))?;
}
scale = 0;
}
if precision > 38 {
return Err(PyValueError::new_err(
"decimal precision exceeds Decimal128 maximum of 38",
));
}
Ok((val, precision, scale))
Ok((val, precision as u8, scale))
}
/// Call an arbitrary registered SQL function by name.
python/src/lib.rs
+8 -3
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,
PyBlobFile, RefreshColumnResult, RefreshMaterializedViewResult, Table,
FunctionErrorFragment, FunctionErrorRecord, FunctionErrors, MergeResult, PyBlobFile,
RTEWriteSpec, 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::<RTEWriteSpec>()?;
m.add_class::<DeleteResult>()?;
m.add_class::<DropColumnsResult>()?;
m.add_class::<UpdateResult>()?;
python/src/permutation.rs
+44
View File
@@ -178,6 +178,50 @@ impl PyAsyncPermutationBuilder {
})
}
#[pyo3(signature = (num_samples, *, ratios=None, counts=None, fixed=None, seed=None, split_names=None))]
pub fn bootstrap_sample(
slf: PyRefMut<'_, Self>,
num_samples: u64,
ratios: Option<Vec<f64>>,
counts: Option<Vec<u64>>,
fixed: Option<u64>,
seed: Option<u64>,
split_names: Option<Vec<String>>,
) -> PyResult<Self> {
// Check that exactly one split type is provided
let split_args_count = [ratios.is_some(), counts.is_some(), fixed.is_some()]
.iter()
.filter(|&&x| x)
.count();
if split_args_count != 1 {
return Err(pyo3::exceptions::PyValueError::new_err(
"Exactly one of 'ratios', 'counts', or 'fixed' must be provided",
));
}
let sizes = if let Some(ratios) = ratios {
SplitSizes::Percentages(ratios)
} else if let Some(counts) = counts {
SplitSizes::Counts(counts)
} else if let Some(fixed) = fixed {
SplitSizes::Fixed(fixed)
} else {
unreachable!("One of the split arguments must be provided");
};
slf.modify(|builder| {
builder.with_split_strategy(
SplitStrategy::Bootstrap {
num_samples,
sizes,
seed,
},
split_names,
)
})
}
pub fn split_calculated(
slf: PyRefMut<'_, Self>,
calculation: String,
python/src/query.rs
+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_rte: 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_rte: query_request.use_rte,
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_rte: vector_query.base.use_rte,
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_rte(&mut self, enable: bool) {
self.inner = self.inner.clone().use_rte(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_rte(&mut self, enable: bool) {
self.inner = self.inner.clone().use_rte(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_rte(&mut self, enable: bool) {
self.inner = self.inner.clone().use_rte(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_rte(&mut self, enable: bool) {
self.inner = self.inner.clone().use_rte(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_rte(&mut self, enable: bool) {
self.inner_vec.use_rte(enable);
self.inner_fts.use_rte(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;
python/src/table.rs
+163 -91
View File
@@ -1,5 +1,6 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use lancedb::table::computed_columns::ComputedColumnOptions;
use std::{
collections::HashMap,
sync::{
@@ -39,16 +40,16 @@ use pyo3::{
mod scannable;
/// Convert `LsmStats` to a Python dict, preserving the per-bucket list.
/// Convert `RTEStats` 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 rte_stats_to_py(py: Python<'_>, stats: &lancedb::table::RTEStats) -> 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)?;
@@ -64,15 +65,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(
@@ -94,9 +95,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())
}
@@ -261,8 +262,8 @@ fn fmt_maintained(maintained: &Option<Vec<String>>) -> String {
}
}
/// Specification selecting Lance's MemWAL LSM-style write path for
/// `merge_insert`.
/// Specification enabling the realtime engine (RTE), which routes
/// `merge_insert` through Lance's MemWAL LSM-style write path.
///
/// Constructed via the `bucket(...)`, `identity(...)`, or `unsharded()`
/// classmethods, then optionally chain `with_maintained_indexes(...)` and
@@ -270,17 +271,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 RTEWriteSpec {
inner: lancedb::table::RTEWriteSpec,
}
#[pymethods]
impl LsmWriteSpec {
impl RTEWriteSpec {
/// 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::RTEWriteSpec::bucket(column, num_buckets),
}
}
@@ -294,20 +295,20 @@ impl LsmWriteSpec {
#[staticmethod]
pub fn identity(column: String) -> Self {
Self {
inner: lancedb::table::LsmWriteSpec::identity(column),
inner: lancedb::table::RTEWriteSpec::identity(column),
}
}
/// No sharding — every `merge_insert` call writes to a single
/// MemWAL shard.
/// RTE shard.
#[staticmethod]
pub fn unsharded() -> Self {
Self {
inner: lancedb::table::LsmWriteSpec::unsharded(),
inner: lancedb::table::RTEWriteSpec::unsharded(),
}
}
/// Set which indexes the MemWAL maintains. `None` (the default) is
/// Set which indexes the RTE maintains. `None` (the default) is
/// every index the table has, including ones created later; a list is
/// verbatim, and an empty list maintains nothing.
#[pyo3(signature = (indexes))]
@@ -327,33 +328,33 @@ impl LsmWriteSpec {
pub fn __repr__(&self) -> String {
match &self.inner {
lancedb::table::LsmWriteSpec::Bucket {
lancedb::table::RTEWriteSpec::Bucket {
column,
num_buckets,
maintained_indexes,
writer_config_defaults,
} => format!(
"LsmWriteSpec.bucket(column={:?}, num_buckets={}, maintained_indexes={}, writer_config_defaults={:?})",
"RTEWriteSpec.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::RTEWriteSpec::Identity {
column,
maintained_indexes,
writer_config_defaults,
} => format!(
"LsmWriteSpec.identity(column={:?}, maintained_indexes={}, writer_config_defaults={:?})",
"RTEWriteSpec.identity(column={:?}, maintained_indexes={}, writer_config_defaults={:?})",
column,
fmt_maintained(maintained_indexes),
writer_config_defaults,
),
lancedb::table::LsmWriteSpec::Unsharded {
lancedb::table::RTEWriteSpec::Unsharded {
maintained_indexes,
writer_config_defaults,
} => format!(
"LsmWriteSpec.unsharded(maintained_indexes={}, writer_config_defaults={:?})",
"RTEWriteSpec.unsharded(maintained_indexes={}, writer_config_defaults={:?})",
fmt_maintained(maintained_indexes),
writer_config_defaults,
),
@@ -365,9 +366,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::RTEWriteSpec::Bucket { .. } => "bucket",
lancedb::table::RTEWriteSpec::Identity { .. } => "identity",
lancedb::table::RTEWriteSpec::Unsharded { .. } => "unsharded",
}
}
@@ -375,9 +376,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::RTEWriteSpec::Bucket { column, .. }
| lancedb::table::RTEWriteSpec::Identity { column, .. } => Some(column.clone()),
lancedb::table::RTEWriteSpec::Unsharded { .. } => None,
}
}
@@ -385,12 +386,12 @@ impl LsmWriteSpec {
#[getter]
pub fn num_buckets(&self) -> Option<u32> {
match &self.inner {
lancedb::table::LsmWriteSpec::Bucket { num_buckets, .. } => Some(*num_buckets),
lancedb::table::RTEWriteSpec::Bucket { num_buckets, .. } => Some(*num_buckets),
_ => None,
}
}
/// Indexes the MemWAL keeps up to date, or `None` for every supported one.
/// Indexes the RTE keeps up to date, or `None` for every supported one.
#[getter]
pub fn maintained_indexes(&self) -> Option<Vec<String>> {
self.inner.maintained_indexes().map(<[String]>::to_vec)
@@ -403,14 +404,14 @@ impl LsmWriteSpec {
}
}
impl From<LsmWriteSpec> for lancedb::table::LsmWriteSpec {
fn from(spec: LsmWriteSpec) -> Self {
impl From<RTEWriteSpec> for lancedb::table::RTEWriteSpec {
fn from(spec: RTEWriteSpec) -> Self {
spec.inner
}
}
impl From<lancedb::table::LsmWriteSpec> for LsmWriteSpec {
fn from(inner: lancedb::table::LsmWriteSpec) -> Self {
impl From<lancedb::table::RTEWriteSpec> for RTEWriteSpec {
fn from(inner: lancedb::table::RTEWriteSpec) -> Self {
Self { inner }
}
}
@@ -697,6 +698,14 @@ impl PyBlobFile {
}
}
fn blob_file_error(operation: &str, error: lance_core::Error) -> pyo3::PyErr {
let message = format!("blob {operation} failed: {error}");
match error {
lance_core::Error::InvalidInput { .. } => PyValueError::new_err(message),
_ => PyRuntimeError::new_err(message),
}
}
#[pymethods]
impl PyBlobFile {
fn read_bytes(self_: PyRef<'_, Self>) -> PyResult<Py<PyBytes>> {
@@ -705,7 +714,7 @@ impl PyBlobFile {
let py = self_.py();
let bytes = py
.detach(move || block_on(async move { inner.read().await }))
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
.map_err(|e| blob_file_error("read", e))?;
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
}
@@ -713,10 +722,7 @@ impl PyBlobFile {
self_.ensure_open()?;
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
let bytes = inner
.read()
.await
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
let bytes = inner.read().await.map_err(|e| blob_file_error("read", e))?;
Python::attach(|py| Ok(PyBytes::new(py, bytes.as_ref()).unbind()))
})
}
@@ -739,7 +745,7 @@ impl PyBlobFile {
self_
.py()
.detach(move || block_on(async move { inner.close().await }))
.map_err(|e| PyRuntimeError::new_err(format!("blob close failed: {e}")))
.map_err(|e| blob_file_error("close", e))
}
fn is_closed(self_: PyRef<'_, Self>) -> bool {
@@ -752,7 +758,7 @@ impl PyBlobFile {
self_
.py()
.detach(move || block_on(async move { inner.seek(position).await }))
.map_err(|e| PyRuntimeError::new_err(format!("blob seek failed: {e}")))
.map_err(|e| blob_file_error("seek", e))
}
fn tell(self_: PyRef<'_, Self>) -> PyResult<u64> {
@@ -761,7 +767,7 @@ impl PyBlobFile {
self_
.py()
.detach(move || block_on(async move { inner.tell().await }))
.map_err(|e| PyRuntimeError::new_err(format!("blob tell failed: {e}")))
.map_err(|e| blob_file_error("tell", e))
}
fn size(self_: PyRef<'_, Self>) -> u64 {
@@ -778,17 +784,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())
}
}
@@ -1106,7 +1134,9 @@ impl Table {
train: Option<bool>,
) -> PyResult<Bound<'a, PyAny>> {
let index = extract_index_params(&index)?;
let timeout = wait_timeout.map(|t| t.extract::<std::time::Duration>().unwrap());
let timeout = wait_timeout
.map(|t| t.extract::<std::time::Duration>())
.transpose()?;
let mut op = self_
.inner_ref()?
.create_index_with_timeout(&[column], index, timeout);
@@ -1137,7 +1167,9 @@ impl Table {
train: Option<bool>,
) -> PyResult<Bound<'a, PyAny>> {
let index = extract_index_params(&index)?;
let timeout = wait_timeout.map(|t| t.extract::<std::time::Duration>().unwrap());
let timeout = wait_timeout
.map(|t| t.extract::<std::time::Duration>())
.transpose()?;
let mut op = self_
.inner_ref()?
.create_index_with_timeout(&[column], index, timeout);
@@ -1400,8 +1432,8 @@ impl Table {
})
}
pub fn query(&self) -> Query {
Query::new(self.inner_ref().unwrap().query())
pub fn query(&self) -> PyResult<Query> {
Ok(Query::new(self.inner_ref()?.query()))
}
#[getter]
@@ -1462,14 +1494,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
@@ -1593,8 +1641,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_rte) = parameters.use_rte {
builder.use_rte(use_rte);
}
if let Some(validate_single_shard) = parameters.validate_single_shard {
builder.validate_single_shard(validate_single_shard);
@@ -1619,29 +1667,29 @@ impl Table {
})
}
pub fn set_lsm_write_spec<'a>(
pub fn enable_rte<'a>(
self_: PyRef<'a, Self>,
spec: LsmWriteSpec,
spec: RTEWriteSpec,
) -> PyResult<Bound<'a, PyAny>> {
let inner = self_.inner_ref()?.clone();
let native_spec = lancedb::table::LsmWriteSpec::from(spec);
let native_spec = lancedb::table::RTEWriteSpec::from(spec);
future_into_py(self_.py(), async move {
inner.set_lsm_write_spec(native_spec).await.infer_error()
inner.enable_rte(native_spec).await.infer_error()
})
}
pub fn unset_lsm_write_spec(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
pub fn disable_rte(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_rte().await.infer_error()
})
}
pub fn get_lsm_write_spec(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
pub fn get_rte_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_rte_spec().await.infer_error()?;
Ok(spec.map(RTEWriteSpec::from))
})
}
@@ -1660,55 +1708,55 @@ impl Table {
Ok(())
}
/// Converge the table's LSM write path into its base table.
/// Converge the table's RTE 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_rte(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_rte().await.infer_error()
})
}
/// Seal every bucket's active memtable into L0.
pub fn flush_lsm(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
/// Freeze every table shard's active memtable into an SSTable.
pub fn flush_rte(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() },
async move { inner.flush_rte().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>> {
/// Trigger a background SSTable compaction pass per table shard. Returns once the
/// passes are dispatched, not once they finish — watch `get_rte_stats`.
pub fn compact_rte(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.compact_rte().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(
/// Live RTE state, or `None` when the RTE write path is not enabled.
#[pyo3(signature = (include_sstable_rows=false))]
pub fn get_rte_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_rte_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| rte_stats_to_py(py, &s)).transpose())
})
}
pub fn close_lsm_writers(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
pub fn close_rte_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_rte_writers().await.infer_error()
})
}
@@ -1760,35 +1808,59 @@ impl Table {
})
}
#[pyo3(signature = (columns, auto_refresh_enabled=None, auto_refresh_interval_seconds=None))]
pub fn add_computed_columns(
self_: PyRef<'_, Self>,
columns: Vec<(String, String)>,
auto_refresh_enabled: Option<bool>,
auto_refresh_interval_seconds: Option<u64>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
let options = ComputedColumnOptions {
auto_refresh_enabled: auto_refresh_enabled.unwrap_or(true),
auto_refresh_interval_seconds,
};
options.validate().infer_error()?;
future_into_py(self_.py(), async move {
let mut builder = inner.add_columns();
for (name, expression) in columns {
builder = builder.computed(name, expression);
}
let result = builder.execute().await.infer_error()?;
let result = builder
.computed_options(options)
.execute()
.await
.infer_error()?;
Ok(AddColumnsResult::from(result))
})
}
#[pyo3(signature = (application_json, output_name, auto_refresh_enabled=None, auto_refresh_interval_seconds=None))]
pub fn add_function_columns(
self_: PyRef<'_, Self>,
application_json: String,
output_name: Option<String>,
auto_refresh_enabled: Option<bool>,
auto_refresh_interval_seconds: Option<u64>,
) -> PyResult<Bound<'_, PyAny>> {
let application =
lancedb::function::FunctionApplication::from_json(&application_json).infer_error()?;
let inner = self_.inner_ref()?.clone();
let options = ComputedColumnOptions {
auto_refresh_enabled: auto_refresh_enabled.unwrap_or(true),
auto_refresh_interval_seconds,
};
options.validate().infer_error()?;
future_into_py(self_.py(), async move {
let builder = match output_name {
Some(name) => inner.add_columns().function_as(name, application),
None => inner.add_columns().function(application),
};
let result = builder.execute().await.infer_error()?;
let result = builder
.computed_options(options)
.execute()
.await
.infer_error()?;
Ok(AddColumnsResult::from(result))
})
}
@@ -2040,7 +2112,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_rte: Option<bool>,
validate_single_shard: Option<bool>,
}
python/uv.lock
Generated
+37 -11
View File
@@ -443,11 +443,11 @@ wheels = [
[[package]]
name = "certifi"
version = "2026.6.17"
version = "2026.10.19"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/c9/c7/424b75da314c1045981bd9777432fad05a9e0c69daa4ed7e308bbaffe405/certifi-2026.6.17.tar.gz", hash = "sha256:024c88eeec92ca068db80f02b8b07c9cef7b9fe261d1d535abfd5abd6f6af432", size = 134594, upload-time = "2026-06-17T10:31:07.894Z" }
sdist = { url = "https://files.pythonhosted.org/packages/a6/bd/1436c50cc46699de0230fbaacc1d1de7bfec413176099949e9d5112e88f2/certifi-2026.10.19.tar.gz", hash = "sha256:035c1e85572b141bf1c3a3eff0e33a8eb74917522f2559075ff8d2fea780a1ae", size = 137598, upload-time = "2026-10-20T02:24:16.604Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/ef/2f/c5464532e965badff2f4c4c1a3a83f5697f0d7c407ed0cda44aaa99bb451/certifi-2026.6.17-py3-none-any.whl", hash = "sha256:2227dcbaafe0d2f59279d1762ddddc37783ed4354594f194ffc31d20f41fc3db", size = 133289, upload-time = "2026-06-17T10:31:06.348Z" },
{ url = "https://files.pythonhosted.org/packages/c4/51/29046e48f693c54e5b7412216d447a1d8362c9653148d98c02cdf65fa733/certifi-2026.10.19-py3-none-any.whl", hash = "sha256:42d35340cc5b9ba78e818ddbff0c5b1c2b779de448a340fd1b989f5fa8670e1b", size = 136464, upload-time = "2026-10-20T02:24:15.072Z" },
]
[[package]]
@@ -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]]
python/vendor/lance-arrow/Cargo.toml
Vendored
+33
View File
@@ -0,0 +1,33 @@
[package]
name = "lance-arrow"
version = "15.0.0-beta.1"
edition = "2024"
authors = ["Lance Devs <dev@lance.org>"]
license = "Apache-2.0"
repository = "https://github.com/lance-format/lance"
readme = "README.md"
description = "Arrow Extension for Lance"
keywords = ["data-format", "data-science", "machine-learning", "apache-arrow", "data-analytics"]
categories = ["database-implementations"]
[lib]
crate-type = ["cdylib", "rlib"]
[dependencies]
arrow-array.workspace = true
arrow-buffer.workspace = true
arrow-data.workspace = true
arrow-ipc = { workspace = true, features = ["zstd"] }
arrow-ord.workspace = true
arrow-schema.workspace = true
arrow-select.workspace = true
bytemuck = { version = "1", default-features = false, features = ["extern_crate_alloc"] }
bytes.workspace = true
futures.workspace = true
half = { workspace = true, features = ["std", "bytemuck"] }
lance-jsonb.workspace = true
num-traits.workspace = true
rand = { workspace = true, features = ["small_rng"] }
[target.'cfg(target_arch = "wasm32")'.dependencies]
getrandom = { version = "0.2", features = ["js"] }
python/vendor/lance-arrow/README.lancedb.md
Vendored
+7
View File
@@ -0,0 +1,7 @@
This directory contains `rust/lance-arrow` from Lance tag `v15.0.0-beta.1`
(commit `46a7cfe1fba5efa1f90deea4c7b14d414b991409`), licensed under Apache-2.0.
LanceDB carries a local fix in `src/lib.rs` for sliced child validity bitmaps
until a Lance release includes the fix. The patch is selected in the root
`Cargo.toml`; the rest of Lance stays on the pinned release. An explicit unsafe
block in `src/bfloat16.rs` lets this copy pass LanceDB's Rust 2024 warning gate.
python/vendor/lance-arrow/README.md
Vendored
+6
View File
@@ -0,0 +1,6 @@
# lance-arrow
`lance-arrow` is an internal sub-crate, containing [Apache-Arrow](https://github.com/apache/arrow-rs)
extensions used by [Lance](https://github.com/lance-format/lance).
**Important Note**: This crate is **not intended for external usage**.
python/vendor/lance-arrow/src/bfloat16.rs
Vendored
+406
View File
@@ -0,0 +1,406 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
// SPDX-FileCopyrightText: Copyright The Lance Authors
//! bfloat16 support for Apache Arrow.
use std::fmt::Formatter;
use std::slice;
use arrow_array::{Array, FixedSizeBinaryArray, builder::BooleanBufferBuilder};
use arrow_buffer::{Buffer, MutableBuffer};
use arrow_data::ArrayData;
use arrow_schema::{ArrowError, DataType, Field as ArrowField};
use half::bf16;
use crate::{ARROW_EXT_NAME_KEY, FloatArray};
/// The name of the bfloat16 extension in Arrow metadata
pub const BFLOAT16_EXT_NAME: &str = "lance.bfloat16";
/// Check whether the given field is a bfloat16 field
///
/// A field is a bfloat16 field if it has a data type of `FixedSizeBinary(2)` and the metadata
/// contains the bfloat16 extension name.
pub fn is_bfloat16_field(field: &ArrowField) -> bool {
field.data_type() == &DataType::FixedSizeBinary(2)
&& field
.metadata()
.get(ARROW_EXT_NAME_KEY)
.map(|name| name == BFLOAT16_EXT_NAME)
.unwrap_or_default()
}
/// The bfloat16 data type
///
/// This implements the [`ArrowFloatType`](crate::floats::ArrowFloatType) trait for bfloat16 values.
#[derive(Debug)]
pub struct BFloat16Type {}
/// An array of bfloat16 values
///
/// Note that bfloat16 is not the same thing as fp16 which is supported natively by arrow-rs.
#[derive(Clone)]
pub struct BFloat16Array {
inner: FixedSizeBinaryArray,
}
impl std::fmt::Debug for BFloat16Array {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "BFloat16Array\n[\n")?;
from_arrow::print_long_array(&self.inner, f, |array, i, f| {
if array.is_null(i) {
write!(f, "null")
} else {
let binary_values = array.value(i);
let value =
bf16::from_bits(u16::from_le_bytes([binary_values[0], binary_values[1]]));
write!(f, "{:?}", value)
}
})?;
write!(f, "]")
}
}
impl BFloat16Array {
pub fn from_iter_values(iter: impl IntoIterator<Item = bf16>) -> Self {
let values: Vec<bf16> = iter.into_iter().collect();
values.into()
}
pub fn len(&self) -> usize {
self.inner.len()
}
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
pub fn is_null(&self, i: usize) -> bool {
self.inner.is_null(i)
}
pub fn null_count(&self) -> usize {
self.inner.null_count()
}
pub fn iter(&self) -> BFloat16Iter<'_> {
BFloat16Iter {
array: self,
index: 0,
}
}
pub fn value(&self, i: usize) -> bf16 {
assert!(
i < self.len(),
"Trying to access an element at index {} from a BFloat16Array of length {}",
i,
self.len()
);
// Safety:
// `i < self.len()
unsafe { self.value_unchecked(i) }
}
/// # Safety
/// Caller must ensure that `i < self.len()`
pub unsafe fn value_unchecked(&self, i: usize) -> bf16 {
// SAFETY: the caller guarantees that `i` is within this array's bounds.
let binary_value = unsafe { self.inner.value_unchecked(i) };
bf16::from_bits(u16::from_le_bytes([binary_value[0], binary_value[1]]))
}
pub fn into_inner(self) -> FixedSizeBinaryArray {
self.inner
}
}
impl FromIterator<Option<bf16>> for BFloat16Array {
fn from_iter<I: IntoIterator<Item = Option<bf16>>>(iter: I) -> Self {
let mut buffer = MutableBuffer::new(10);
// No null buffer builder :(
let mut nulls = BooleanBufferBuilder::new(10);
let mut len = 0;
for maybe_value in iter {
if let Some(value) = maybe_value {
let bytes = value.to_le_bytes();
buffer.extend(bytes);
} else {
buffer.extend([0u8, 0u8]);
}
nulls.append(maybe_value.is_some());
len += 1;
}
let null_buffer = nulls.finish();
let num_valid = null_buffer.count_set_bits();
let null_buffer = if num_valid == len {
None
} else {
Some(null_buffer.into_inner())
};
let array_data = ArrayData::builder(DataType::FixedSizeBinary(2))
.len(len)
.add_buffer(buffer.into())
.null_bit_buffer(null_buffer);
// SAFETY: the value buffer contains exactly `2 * len` bytes (two bytes
// pushed per iteration of the loop above, including the zero-fill for
// null slots), which matches the `FixedSizeBinary(2)` storage layout.
// The null bit buffer, when present, has `len` bits appended above, so
// its length covers the array's logical range.
let array_data = unsafe { array_data.build_unchecked() };
Self {
inner: FixedSizeBinaryArray::from(array_data),
}
}
}
impl FromIterator<bf16> for BFloat16Array {
fn from_iter<I: IntoIterator<Item = bf16>>(iter: I) -> Self {
Self::from_iter_values(iter)
}
}
impl From<Vec<bf16>> for BFloat16Array {
fn from(data: Vec<bf16>) -> Self {
let len = data.len();
// Zero-copy: `bf16` is `#[repr(transparent)]` over `u16` and derives
// `bytemuck::Pod`, so `cast_vec` reinterprets the allocation in place —
// no per-element copy or heap alloc. The crate-root `compile_error!`
// pins `target_endian = "little"`, so the resulting bytes match the
// `FixedSizeBinary(2)` on-disk order Lance writes elsewhere.
let raw: Vec<u16> = bytemuck::cast_vec(data);
let array_data = ArrayData::builder(DataType::FixedSizeBinary(2))
.len(len)
.add_buffer(Buffer::from_vec(raw));
// SAFETY: the value buffer contains exactly `2 * len` bytes — one
// `u16` per element after the layout-compatible cast — matching the
// `FixedSizeBinary(2)` storage layout. No null buffer is attached, so
// every element is logically valid.
let array_data = unsafe { array_data.build_unchecked() };
Self {
inner: FixedSizeBinaryArray::from(array_data),
}
}
}
impl TryFrom<FixedSizeBinaryArray> for BFloat16Array {
type Error = ArrowError;
fn try_from(value: FixedSizeBinaryArray) -> Result<Self, Self::Error> {
if value.value_length() == 2 {
Ok(Self { inner: value })
} else {
Err(ArrowError::InvalidArgumentError(
"FixedSizeBinaryArray must have a value length of 2".to_string(),
))
}
}
}
impl PartialEq<Self> for BFloat16Array {
fn eq(&self, other: &Self) -> bool {
self.inner.eq(&other.inner)
}
}
pub struct BFloat16Iter<'a> {
array: &'a BFloat16Array,
index: usize,
}
impl<'a> Iterator for BFloat16Iter<'a> {
type Item = Option<bf16>;
fn next(&mut self) -> Option<Self::Item> {
if self.index >= self.array.len() {
return None;
}
let i = self.index;
self.index += 1;
if self.array.is_null(i) {
Some(None)
} else {
Some(Some(self.array.value(i)))
}
}
}
/// Methods that are lifted from arrow-rs temporarily until they are made public.
mod from_arrow {
use arrow_array::Array;
/// Helper function for printing potentially long arrays.
pub(super) fn print_long_array<A, F>(
array: &A,
f: &mut std::fmt::Formatter,
print_item: F,
) -> std::fmt::Result
where
A: Array,
F: Fn(&A, usize, &mut std::fmt::Formatter) -> std::fmt::Result,
{
let head = std::cmp::min(10, array.len());
for i in 0..head {
if array.is_null(i) {
writeln!(f, " null,")?;
} else {
write!(f, " ")?;
print_item(array, i, f)?;
writeln!(f, ",")?;
}
}
if array.len() > 10 {
if array.len() > 20 {
writeln!(f, " ...{} elements...,", array.len() - 20)?;
}
let tail = std::cmp::max(head, array.len() - 10);
for i in tail..array.len() {
if array.is_null(i) {
writeln!(f, " null,")?;
} else {
write!(f, " ")?;
print_item(array, i, f)?;
writeln!(f, ",")?;
}
}
}
Ok(())
}
}
impl FloatArray<BFloat16Type> for FixedSizeBinaryArray {
type FloatType = BFloat16Type;
/// Returns the underlying `bf16` values as a borrowed slice.
///
/// # Preconditions
///
/// - `value_length()` must be 2 (the `FixedSizeBinary(2)` storage shape
/// used by [`BFloat16Array`]). Asserted at entry.
/// - The value buffer must be at least 2-byte aligned. Lance's in-tree
/// constructors always satisfy this: value buffers are built either via
/// `MutableBuffer` (aligned to arrow-buffer's `ALIGNMENT` constant, ≥32
/// bytes) or via `Buffer::from_vec::<u16>` (aligned to `align_of::<u16>()`
/// == 2); both meet `bf16`'s 2-byte requirement. Externally-built
/// `FixedSizeBinaryArray`s arriving via FFI, IPC, or
/// `Buffer::from_custom_allocation` are not required by arrow-rs to be
/// aligned beyond a single byte; passing one to this method violates the
/// precondition. A `debug_assert` below catches such inputs in debug and
/// test builds.
///
/// # Endianness
///
/// `lance-arrow` is gated on `target_endian = "little"` at the crate root,
/// so this method always returns values in the same byte order Lance writes
/// (see [`BFloat16Array::value`] and the [`FromIterator`] impls).
fn as_slice(&self) -> &[bf16] {
assert_eq!(
self.value_length(),
2,
"BFloat16 arrays must use FixedSizeBinary(2) storage"
);
debug_assert_eq!(
(self.value_data().as_ptr() as usize) % std::mem::align_of::<bf16>(),
0,
"BFloat16 value buffer must be at least 2-byte aligned"
);
// SAFETY:
// - The assert above pins `value_size == 2`, so `value_data().len() / 2`
// equals the array's logical element count.
// `FixedSizeBinaryArray::From<ArrayData>` constructs its value buffer
// as `buffers[0].slice_with_length(offset * 2, len * 2)` (arrow-array
// `fixed_size_binary_array.rs`), so `value_data()` already returns
// the offset-adjusted slice. Do not replace `value_data()` with an
// accessor that returns the un-sliced backing buffer.
// - `bf16` is `#[repr(transparent)]` over `u16` (size 2, alignment 2);
// every `u16` bit pattern is a valid `bf16`, so any byte content
// yields a defined value — never UB.
// - Alignment is the caller's responsibility per the precondition
// documented above. The `debug_assert_eq!` immediately preceding this
// block catches violations in debug and test builds only — release
// builds rely on callers honoring the precondition. arrow-rs
// declares `FixedSizeBinary(n)`'s
// `BufferSpec::FixedWidth { alignment: align_of::<u8>() == 1 }`
// (arrow-data `data.rs`), so arrow-rs alone does not guarantee
// 2-byte alignment. Lance's in-tree construction paths build value
// buffers via `MutableBuffer` (arrow-buffer `ALIGNMENT` constant,
// ≥32 bytes) or `Buffer::from_vec::<u16>` (2-byte aligned), both of
// which satisfy `bf16`'s 2-byte requirement.
// - The returned slice borrows from `self`; the underlying ref-counted,
// immutable Arrow buffer cannot be mutated or freed for the slice's
// lifetime.
unsafe {
slice::from_raw_parts(
self.value_data().as_ptr() as *const bf16,
self.value_data().len() / 2,
)
}
}
fn from_values(values: Vec<bf16>) -> Self {
BFloat16Array::from(values).into_inner()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_basics() {
let values: Vec<f32> = vec![1.0, 2.0, 3.0];
let values: Vec<bf16> = values.iter().map(|v| bf16::from_f32(*v)).collect();
let array = BFloat16Array::from_iter_values(values.clone());
let array2 = BFloat16Array::from(values.clone());
assert_eq!(array, array2);
assert_eq!(array.len(), 3);
// Pin the raw little-endian bytes emitted by `From<Vec<bf16>>` (rewritten to
// reinterpret the Vec via `bytemuck::cast_vec`), so a layout/byte-order
// regression is caught directly rather than only through Debug formatting.
// bf16 is the high 16 bits of the f32: 1.0->0x3F80, 2.0->0x4000, 3.0->0x4040.
let inner = array2.clone().into_inner();
let raw_bytes: Vec<u8> = (0..inner.len())
.flat_map(|i| inner.value(i).to_vec())
.collect();
assert_eq!(raw_bytes, vec![0x80, 0x3F, 0x00, 0x40, 0x40, 0x40]);
let expected_fmt = "BFloat16Array\n[\n 1.0,\n 2.0,\n 3.0,\n]";
assert_eq!(expected_fmt, format!("{:?}", array));
for (expected, value) in values.iter().zip(array.iter()) {
assert_eq!(Some(*expected), value);
}
for (expected, value) in values.as_slice().iter().zip(array2.iter()) {
assert_eq!(Some(*expected), value);
}
let arrow_array = array.into_inner();
assert_eq!(arrow_array.as_slice(), values.as_slice());
}
#[test]
fn test_nulls() {
let values: Vec<Option<bf16>> =
vec![Some(bf16::from_f32(1.0)), None, Some(bf16::from_f32(3.0))];
let array = BFloat16Array::from_iter(values.clone());
assert_eq!(array.len(), 3);
assert_eq!(array.null_count(), 1);
let expected_fmt = "BFloat16Array\n[\n 1.0,\n null,\n 3.0,\n]";
assert_eq!(expected_fmt, format!("{:?}", array));
for (expected, value) in values.iter().zip(array.iter()) {
assert_eq!(*expected, value);
}
}
}
python/vendor/lance-arrow/src/deepcopy.rs
Vendored
+356
View File
@@ -0,0 +1,356 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
// SPDX-FileCopyrightText: Copyright The Lance Authors
use std::sync::Arc;
use arrow_array::{Array, RecordBatch, make_array};
use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer};
use arrow_data::{ArrayData, ArrayDataBuilder, transform::MutableArrayData};
use arrow_schema::DataType;
pub fn deep_copy_buffer(buffer: &Buffer) -> Buffer {
Buffer::from(buffer.as_slice())
}
pub fn deep_copy_nulls(nulls: Option<&NullBuffer>) -> Option<NullBuffer> {
let nulls = nulls?;
let bit_buffer = deep_copy_buffer(nulls.inner().inner());
// SAFETY: `null_count` is taken from the source `NullBuffer`, which already
// upheld `NullBuffer::new_unchecked`'s invariant — the unset-bit count over
// the logical bit slice `[bit_offset, bit_offset + bit_len)`. `NullBuffer::slice`
// adjusts only `BooleanBuffer::bit_offset` / `bit_len` and never byte-advances
// the inner `Buffer`, so `deep_copy_buffer` (which copies the source `Buffer`'s
// `as_slice()` view from byte 0) reproduces the exact bit pattern at the same
// bit offsets; the unset-bit count is therefore preserved. `BooleanBuffer::new`
// panics (does not UB) if `bit_offset + bit_len > 8 * buffer.len()`, and the
// copy has the same length, so that check still passes.
Some(unsafe {
NullBuffer::new_unchecked(
BooleanBuffer::new(bit_buffer, nulls.offset(), nulls.len()),
nulls.null_count(),
)
})
}
pub fn deep_copy_array_data(data: &ArrayData) -> ArrayData {
let data_type = data.data_type().clone();
let len = data.len();
let nulls = deep_copy_nulls(data.nulls());
let offset = data.offset();
let buffers = data
.buffers()
.iter()
.map(deep_copy_buffer)
.collect::<Vec<_>>();
let child_data = data
.child_data()
.iter()
.map(deep_copy_array_data)
.collect::<Vec<_>>();
// SAFETY: `build_unchecked` inherits `ArrayData::new_unchecked`'s contract —
// `(data_type, len, offset, nulls, buffers, child_data)` must form a valid
// Arrow array. This call reproduces `data` structurally: `data_type`, `len`,
// and `offset` are forwarded unchanged; each buffer is replaced by a byte-
// identical copy of its offset-applied `as_slice()` view (the output buffer
// is `MutableBuffer`-allocated, at least as aligned as the source); `nulls`
// is deep-copied with the same bit offset/length and unset-bit count (see
// `deep_copy_nulls`); `child_data` is recursively cloned with the same
// guarantee. Every value-level invariant the source upheld — UTF-8 validity,
// monotonic offsets, in-bounds dictionary indices, run-end monotonicity,
// struct child-length matching — therefore transfers to the copy. If the
// source `ArrayData` was itself constructed via `new_unchecked` with an
// invalid payload, this function faithfully reproduces that invalidity.
unsafe {
ArrayDataBuilder::new(data_type)
.len(len)
.nulls(nulls)
.offset(offset)
.buffers(buffers)
.child_data(child_data)
.build_unchecked()
}
}
pub fn deep_copy_array(array: &dyn Array) -> Arc<dyn Array> {
let data = array.to_data();
let data = deep_copy_array_data(&data);
make_array(data)
}
pub fn deep_copy_batch(batch: &RecordBatch) -> crate::Result<RecordBatch> {
let arrays = batch
.columns()
.iter()
.map(|array| deep_copy_array(array))
.collect::<Vec<_>>();
RecordBatch::try_new(batch.schema(), arrays)
}
/// Deep copy array data, extracting only the sliced portion using MutableArrayData
/// This is the most efficient and correct way to copy just the sliced data
pub fn deep_copy_array_data_sliced(data: &ArrayData) -> ArrayData {
// Use MutableArrayData to efficiently copy just the slice
let mut mutable = MutableArrayData::new(vec![data], false, data.len());
// Which index this takes depends on the layout, because arrow's extenders
// disagree about who applies `ArrayData::offset()`. The bit-packed one adds
// it to the raw values buffer itself, so an offset-applied index there
// reads from twice the offset and runs off the end of the buffer. Every
// other layout either reads through an already-offset-applied view or
// forwards the index to children carrying their own offsets, and takes the
// absolute index it has always been given.
let start = match data.data_type() {
DataType::Boolean => 0,
_ => data.offset(),
};
mutable.extend(0, start, start + data.len());
// Freeze into immutable ArrayData
mutable.freeze()
}
/// Deep copy an array, extracting only the sliced portion using MutableArrayData
pub fn deep_copy_array_sliced(array: &dyn Array) -> Arc<dyn Array> {
let data = array.to_data();
let data = deep_copy_array_data_sliced(&data);
make_array(data)
}
/// Deep copy a RecordBatch, extracting only the sliced portion using MutableArrayData
pub fn deep_copy_batch_sliced(batch: &RecordBatch) -> crate::Result<RecordBatch> {
let arrays = batch
.columns()
.iter()
.map(|array| deep_copy_array_sliced(array))
.collect::<Vec<_>>();
RecordBatch::try_new(batch.schema(), arrays)
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use arrow_array::{Array, BooleanArray, Int32Array, RecordBatch, StringArray};
use arrow_data::ArrayDataBuilder;
use arrow_schema::{DataType, Field, Schema};
#[test]
fn raw_sliced_fixed_size_list_data_keeps_its_child_offset() {
// `ArrayData::slice` records the offset on the parent and leaves the
// child whole -- a shape `FixedSizeListArray::slice` never produces,
// but a legal one this public helper can be handed directly. Arrow's
// extender forwards `start * size` to that unsliced child and adds
// nothing, so the index it gets has to be the absolute one; this is
// here to stop the boolean fix from being generalised over it.
let child = Int32Array::from(vec![10, 11, 20, 21]).to_data();
let field = Arc::new(Field::new_list_field(DataType::Int32, false));
let data = ArrayDataBuilder::new(DataType::FixedSizeList(field, 2))
.len(2)
.add_child_data(child)
.build()
.unwrap();
let sliced = data.slice(1, 1);
let copied = super::deep_copy_array_data_sliced(&sliced);
let copied_child = Int32Array::from(copied.child_data()[0].clone());
assert_eq!(copied_child.values().as_ref(), &[20, 21]);
}
#[test]
fn raw_sliced_struct_of_fixed_size_list_reaches_the_right_values() {
// Slicing raw struct data pushes the window into the struct's children,
// and a fixed-size-list child takes it as its own parent offset with
// its values left whole. Two levels of offset, and the absolute index
// has to remain right through both of them.
let values = Int32Array::from(vec![10, 11, 20, 21, 30, 31]).to_data();
let item = Arc::new(Field::new_list_field(DataType::Int32, false));
let list = ArrayDataBuilder::new(DataType::FixedSizeList(item, 2))
.len(3)
.add_child_data(values)
.build()
.unwrap();
let field = Arc::new(Field::new("l", list.data_type().clone(), false));
let data = ArrayDataBuilder::new(DataType::Struct(vec![field].into()))
.len(3)
.add_child_data(list)
.build()
.unwrap();
let sliced = data.slice(2, 1);
let copied = super::deep_copy_array_data_sliced(&sliced);
let copied_values = Int32Array::from(copied.child_data()[0].child_data()[0].clone());
assert_eq!(copied_values.values().as_ref(), &[30, 31]);
}
#[test]
fn raw_parent_offset_struct_keeps_the_selected_child_row() {
// A struct whose children are whole and whose window is the parent
// offset: arrow's struct extender forwards the index to those children
// untouched, so it has to be the absolute one.
let field = Arc::new(Field::new("a", DataType::Int32, false));
let data = ArrayDataBuilder::new(DataType::Struct(vec![field].into()))
.len(1)
.offset(1)
.add_child_data(Int32Array::from(vec![10, 20]).to_data())
.build()
.unwrap();
let copied = super::deep_copy_array_data_sliced(&data);
let copied_child = Int32Array::from(copied.child_data()[0].clone());
assert_eq!(copied_child.values().as_ref(), &[20]);
}
#[test]
fn sliced_boolean_deep_copy_reads_from_the_slice() {
// A boolean slice keeps its `ArrayData::offset()` -- the buffer cannot
// advance by a fraction of a byte -- so a copy that adds the offset on
// top of what `MutableArrayData` already applies reads past the end of
// the values buffer and panics.
let array = BooleanArray::from(vec![true, false, true, false, true, false, true, false]);
for (offset, len) in [(0usize, 8usize), (1, 7), (3, 5), (7, 1)] {
let sliced = array.slice(offset, len);
let copied = super::deep_copy_array_sliced(&sliced);
let copied = copied.as_any().downcast_ref::<BooleanArray>().unwrap();
let expected: Vec<bool> = (0..len).map(|i| sliced.value(i)).collect();
let actual: Vec<bool> = (0..len).map(|i| copied.value(i)).collect();
assert_eq!(actual, expected, "offset={offset} len={len}");
}
}
#[test]
fn sliced_boolean_deep_copy_keeps_its_nulls() {
let array = BooleanArray::from(vec![
Some(true),
None,
Some(false),
Some(true),
None,
Some(false),
]);
let sliced = array.slice(1, 4);
let copied = super::deep_copy_array_sliced(&sliced);
let copied = copied.as_any().downcast_ref::<BooleanArray>().unwrap();
let expected: Vec<Option<bool>> = (0..sliced.len())
.map(|i| (!sliced.is_null(i)).then(|| sliced.value(i)))
.collect();
let actual: Vec<Option<bool>> = (0..copied.len())
.map(|i| (!copied.is_null(i)).then(|| copied.value(i)))
.collect();
assert_eq!(actual, expected);
}
#[test]
fn test_deep_copy_sliced_array_with_nulls() {
let array = Arc::new(Int32Array::from(vec![
Some(1),
None,
Some(3),
None,
Some(5),
]));
let sliced_array = array.slice(1, 3);
let copied_array = super::deep_copy_array(&sliced_array);
assert_eq!(sliced_array.len(), copied_array.len());
assert_eq!(sliced_array.nulls(), copied_array.nulls());
}
#[test]
fn test_deep_copy_array_data_sliced() {
let array = Int32Array::from((0..1000).collect::<Vec<i32>>());
let sliced = array.slice(100, 10);
let sliced_data = sliced.to_data();
let copied_data = super::deep_copy_array_data_sliced(&sliced_data);
assert_eq!(copied_data.len(), 10);
assert_eq!(copied_data.offset(), 0);
// Verify data correctness
let copied_array = Int32Array::from(copied_data);
for i in 0..10 {
assert_eq!(copied_array.value(i), 100 + i as i32);
}
}
#[test]
fn test_deep_copy_array_sliced() {
let array = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]));
let sliced = array.slice(1, 3);
let copied = super::deep_copy_array_sliced(&sliced);
assert_eq!(copied.len(), 3);
let copied_int = copied.as_any().downcast_ref::<Int32Array>().unwrap();
assert_eq!(copied_int.value(0), 2);
assert_eq!(copied_int.value(1), 3);
assert_eq!(copied_int.value(2), 4);
}
#[test]
fn test_deep_copy_batch_sliced() {
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("name", DataType::Utf8, false),
]));
let id_array = Arc::new(Int32Array::from((0..100).collect::<Vec<i32>>()));
let name_array = Arc::new(StringArray::from(
(0..100)
.map(|i| format!("name_{}", i))
.collect::<Vec<String>>(),
));
let batch = RecordBatch::try_new(
schema,
vec![id_array as Arc<dyn Array>, name_array as Arc<dyn Array>],
)
.unwrap();
let sliced = batch.slice(10, 5);
let copied = super::deep_copy_batch_sliced(&sliced).unwrap();
assert_eq!(copied.num_rows(), 5);
assert_eq!(copied.num_columns(), 2);
// Verify data correctness
let id_col = copied
.column(0)
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
let name_col = copied
.column(1)
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
for i in 0..5 {
assert_eq!(id_col.value(i), 10 + i as i32);
assert_eq!(name_col.value(i), format!("name_{}", 10 + i));
}
}
#[test]
fn test_deep_copy_array_sliced_with_nulls() {
let array = Arc::new(Int32Array::from(vec![
Some(1),
None,
Some(3),
None,
Some(5),
]));
let sliced = array.slice(1, 3); // [None, Some(3), None]
let copied = super::deep_copy_array_sliced(&sliced);
assert_eq!(copied.len(), 3);
assert_eq!(copied.null_count(), 2); // Two nulls in the slice
let copied_int = copied.as_any().downcast_ref::<Int32Array>().unwrap();
assert!(!copied_int.is_valid(0)); // None
assert!(copied_int.is_valid(1)); // Some(3)
assert!(!copied_int.is_valid(2)); // None
assert_eq!(copied_int.value(1), 3);
}
}
python/vendor/lance-arrow/src/floats.rs
Vendored
+292
View File
@@ -0,0 +1,292 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
// SPDX-FileCopyrightText: Copyright The Lance Authors
//! Floats Array
use std::fmt::{Debug, Display};
use std::iter::Sum;
use std::sync::Arc;
use std::{
fmt::Formatter,
ops::{AddAssign, DivAssign},
};
use arrow_array::{
Array, FixedSizeBinaryArray, Float16Array, Float32Array, Float64Array,
types::{Float16Type, Float32Type, Float64Type},
};
use arrow_schema::{DataType, Field};
use half::{bf16, f16};
use num_traits::{AsPrimitive, Bounded, Float, FromPrimitive};
use super::bfloat16::{BFloat16Array, BFloat16Type};
use crate::Result;
use crate::bfloat16::is_bfloat16_field;
/// Float data type.
///
/// This helps differentiate between the different float types,
/// because bf16 is not officially supported [DataType] in arrow-rs.
#[derive(Debug)]
pub enum FloatType {
BFloat16,
Float16,
Float32,
Float64,
}
impl std::fmt::Display for FloatType {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Self::BFloat16 => write!(f, "bfloat16"),
Self::Float16 => write!(f, "float16"),
Self::Float32 => write!(f, "float32"),
Self::Float64 => write!(f, "float64"),
}
}
}
/// Try to convert a [DataType] to a [FloatType]. To support bfloat16, always
/// prefer using the `TryFrom<&Field>` implementation.
impl TryFrom<&DataType> for FloatType {
type Error = crate::ArrowError;
fn try_from(value: &DataType) -> Result<Self> {
match *value {
DataType::Float16 => Ok(Self::Float16),
DataType::Float32 => Ok(Self::Float32),
DataType::Float64 => Ok(Self::Float64),
_ => Err(crate::ArrowError::InvalidArgumentError(format!(
"{:?} is not a floating type",
value
))),
}
}
}
impl TryFrom<&Field> for FloatType {
type Error = crate::ArrowError;
fn try_from(field: &Field) -> Result<Self> {
match field.data_type() {
DataType::FixedSizeBinary(2) if is_bfloat16_field(field) => Ok(Self::BFloat16),
_ => Self::try_from(field.data_type()),
}
}
}
/// Trait for float types used in Lance indexes
///
/// This mimics the utilities provided by [`arrow_array::ArrowPrimitiveType`]
/// but applies to all float types (including bfloat16)
pub trait ArrowFloatType: Debug {
type Native: FromPrimitive
+ FloatToArrayType<ArrowType = Self>
+ AsPrimitive<f32>
+ Debug
+ Display;
const FLOAT_TYPE: FloatType;
const MIN: Self::Native;
const MAX: Self::Native;
/// Arrow Float Array Type.
type ArrayType: FloatArray<Self>;
/// Returns empty array of this type.
fn empty_array() -> Self::ArrayType {
<Self::ArrayType as FloatArray<Self>>::from_values(Vec::new())
}
}
/// Trait to be implemented by native types that have a corresponding [`ArrowFloatType`]
/// implementation.
///
/// This helps define what operations are supported by native floats and also helps convert
/// from a native type back to the corresponding Arrow float type.
pub trait FloatToArrayType:
Float
+ Bounded
+ Sum
+ AddAssign<Self>
+ AsPrimitive<f64>
+ AsPrimitive<f32>
+ DivAssign
+ Send
+ Sync
+ Copy
{
/// The corresponding [`ArrowFloatType`] implementation for this native type
type ArrowType: ArrowFloatType<Native = Self>;
}
impl FloatToArrayType for bf16 {
type ArrowType = BFloat16Type;
}
impl FloatToArrayType for f16 {
type ArrowType = Float16Type;
}
impl FloatToArrayType for f32 {
type ArrowType = Float32Type;
}
impl FloatToArrayType for f64 {
type ArrowType = Float64Type;
}
impl ArrowFloatType for BFloat16Type {
type Native = bf16;
const FLOAT_TYPE: FloatType = FloatType::BFloat16;
const MIN: Self::Native = bf16::MIN;
const MAX: Self::Native = bf16::MAX;
type ArrayType = FixedSizeBinaryArray;
}
impl ArrowFloatType for Float16Type {
type Native = f16;
const FLOAT_TYPE: FloatType = FloatType::Float16;
const MIN: Self::Native = f16::MIN;
const MAX: Self::Native = f16::MAX;
type ArrayType = Float16Array;
}
impl ArrowFloatType for Float32Type {
type Native = f32;
const FLOAT_TYPE: FloatType = FloatType::Float32;
const MIN: Self::Native = f32::MIN;
const MAX: Self::Native = f32::MAX;
type ArrayType = Float32Array;
}
impl ArrowFloatType for Float64Type {
type Native = f64;
const FLOAT_TYPE: FloatType = FloatType::Float64;
const MIN: Self::Native = f64::MIN;
const MAX: Self::Native = f64::MAX;
type ArrayType = Float64Array;
}
/// [FloatArray] is a trait that is implemented by all float type arrays
///
/// This is similar to [`arrow_array::PrimitiveArray`] but applies to all float types (including bfloat16)
/// and is implemented as a trait and not a struct
pub trait FloatArray<T: ArrowFloatType + ?Sized>: Array + Clone + 'static {
type FloatType: ArrowFloatType;
/// Returns a reference to the underlying data as a slice.
///
/// # Panics
///
/// Implementations may panic if the array's storage shape does not match
/// the expected element layout. In particular, the `bf16` impl panics if
/// `value_length() != 2` (the `FixedSizeBinary(2)` shape required by
/// `BFloat16Array`).
///
/// # Preconditions
///
/// Implementations may impose additional invariants on the underlying
/// buffer. The `bf16` impl requires the value buffer to be at least
/// 2-byte aligned — satisfied automatically by every in-tree Lance
/// constructor, but external callers passing externally-built arrays
/// (FFI, IPC, `Buffer::from_custom_allocation`) must ensure alignment.
/// See the impl's docstring for details.
fn as_slice(&self) -> &[T::Native];
/// Construct an array from a vector of values.
fn from_values(values: Vec<T::Native>) -> Self;
/// Construct an array from an iterator of values.
fn from_iter_values(values: impl IntoIterator<Item = T::Native>) -> Self
where
Self: Sized,
{
Self::from_values(values.into_iter().collect())
}
}
impl FloatArray<Float16Type> for Float16Array {
type FloatType = Float16Type;
fn as_slice(&self) -> &[<Float16Type as ArrowFloatType>::Native] {
self.values()
}
fn from_values(values: Vec<<Float16Type as ArrowFloatType>::Native>) -> Self {
Self::from(values)
}
}
impl FloatArray<Float32Type> for Float32Array {
type FloatType = Float32Type;
fn as_slice(&self) -> &[<Float32Type as ArrowFloatType>::Native] {
self.values()
}
fn from_values(values: Vec<<Float32Type as ArrowFloatType>::Native>) -> Self {
Self::from(values)
}
}
impl FloatArray<Float64Type> for Float64Array {
type FloatType = Float64Type;
fn as_slice(&self) -> &[<Float64Type as ArrowFloatType>::Native] {
self.values()
}
fn from_values(values: Vec<<Float64Type as ArrowFloatType>::Native>) -> Self {
Self::from(values)
}
}
/// Convert a float32 array to another float array
///
/// This is used during queries as query vectors are always provided as float32 arrays
/// and need to be converted to the appropriate float type for the index.
pub fn coerce_float_vector(input: &Float32Array, float_type: FloatType) -> Result<Arc<dyn Array>> {
match float_type {
FloatType::BFloat16 => Ok(Arc::new(
BFloat16Array::from_iter_values(input.values().iter().map(|v| bf16::from_f32(*v)))
.into_inner(),
)),
FloatType::Float16 => Ok(Arc::new(Float16Array::from_iter_values(
input.values().iter().map(|v| f16::from_f32(*v)),
))),
FloatType::Float32 => Ok(Arc::new(input.clone())),
FloatType::Float64 => Ok(Arc::new(Float64Array::from_iter_values(
input.values().iter().map(|v| *v as f64),
))),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_coerce_float_vector_bfloat16() {
let input = Float32Array::from(vec![1.0f32, 2.0, 3.0]);
let array = coerce_float_vector(&input, FloatType::BFloat16).unwrap();
assert_eq!(array.data_type(), &DataType::FixedSizeBinary(2));
let fixed = array
.as_any()
.downcast_ref::<FixedSizeBinaryArray>()
.unwrap();
let expected: Vec<bf16> = input.values().iter().map(|v| bf16::from_f32(*v)).collect();
assert_eq!(fixed.as_slice(), expected.as_slice());
}
}
python/vendor/lance-arrow/src/ipc.rs
Vendored
+621
View File
@@ -0,0 +1,621 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
// SPDX-FileCopyrightText: Copyright The Lance Authors
//! Zero-copy Arrow IPC stream read/write utilities.
//!
//! Provides helpers for serializing and deserializing [`RecordBatch`]es as
//! self-delimiting Arrow IPC streams using synchronous [`Read`]/[`Write`] I/O.
//!
//! These are designed for embedding IPC streams inside larger binary formats
//! (e.g. a cache entry that contains multiple IPC sections). Each stream is
//! self-delimiting (schema + batches + EOS marker) and can be read back
//! independently.
//!
//! # Zero-copy reads
//!
//! [`read_ipc_stream`] and [`read_ipc_stream_single`] take `&Bytes` and use
//! [`Bytes::slice`] to produce each message buffer. Because `Bytes::slice`
//! increments a reference count rather than copying, the resulting
//! [`Buffer`]s — and the array data decoded from them by [`FileDecoder`] —
//! are all backed by the same allocation as the input.
use std::io::{Read, Write};
use std::sync::Arc;
use arrow_array::RecordBatch;
use arrow_buffer::Buffer;
use arrow_ipc::convert::fb_to_schema;
use arrow_ipc::reader::FileDecoder;
use arrow_ipc::root_as_message;
use arrow_ipc::writer::StreamWriter;
use arrow_schema::ArrowError;
use bytes::Bytes;
// ---------------------------------------------------------------------------
// Length-prefixed byte utilities
// ---------------------------------------------------------------------------
/// Write `data` prefixed by its length as a little-endian `u64`.
///
/// Paired with [`read_len_prefixed_bytes`].
pub fn write_len_prefixed_bytes(writer: &mut dyn Write, data: &[u8]) -> Result<(), ArrowError> {
writer
.write_all(&(data.len() as u64).to_le_bytes())
.map_err(|e| ArrowError::IoError(e.to_string(), e))?;
writer
.write_all(data)
.map_err(|e| ArrowError::IoError(e.to_string(), e))
}
/// Read a byte slice written by [`write_len_prefixed_bytes`].
///
/// Reads an 8-byte little-endian length then exactly that many bytes.
pub fn read_len_prefixed_bytes(reader: &mut dyn Read) -> Result<Vec<u8>, ArrowError> {
let mut len_buf = [0u8; 8];
reader
.read_exact(&mut len_buf)
.map_err(|e| ArrowError::IoError(e.to_string(), e))?;
let len = u64::from_le_bytes(len_buf) as usize;
let mut buf = vec![0u8; len];
reader
.read_exact(&mut buf)
.map_err(|e| ArrowError::IoError(e.to_string(), e))?;
Ok(buf)
}
// ---------------------------------------------------------------------------
// IPC stream utilities
// ---------------------------------------------------------------------------
// 4-byte continuation marker used by modern Arrow IPC streams.
const IPC_CONTINUATION: [u8; 4] = [0xff; 4];
/// Write `batch` as a single-batch Arrow IPC stream to `writer`.
pub fn write_ipc_stream(batch: &RecordBatch, writer: &mut dyn Write) -> Result<(), ArrowError> {
let mut sw = StreamWriter::try_new(&mut *writer, batch.schema_ref())?;
sw.write(batch)?;
sw.finish()
}
/// Write all batches from `iter` as a single Arrow IPC stream to `writer`.
///
/// `iter` must yield at least one batch; the schema is inferred from the first
/// batch. Returns `ArrowError::InvalidArgumentError` if the iterator is empty.
/// If you need to write an empty stream (schema only, no rows), construct a
/// `StreamWriter` directly.
pub fn write_ipc_stream_batches<I>(iter: I, writer: &mut dyn Write) -> Result<(), ArrowError>
where
I: IntoIterator<Item = RecordBatch>,
{
let mut iter = iter.into_iter();
let first = iter
.next()
.ok_or_else(|| ArrowError::InvalidArgumentError("no batches to serialize".into()))?;
let mut sw = StreamWriter::try_new(&mut *writer, first.schema_ref())?;
sw.write(&first)?;
for batch in iter {
sw.write(&batch)?;
}
sw.finish()
}
/// Read one complete Arrow IPC stream message from `data` as a zero-copy [`Buffer`].
///
/// Parses the first message starting at byte 0 of `data`. Returns `None` on
/// EOS (size field == 0) or empty input. The returned [`Buffer`] is backed by
/// `data`'s allocation — no bytes are copied.
///
/// The caller should advance its position by `buf.len()` after each call.
fn read_one_ipc_message(data: &Bytes) -> Result<Option<Buffer>, ArrowError> {
let bytes = data.as_ref();
if bytes.is_empty() {
return Ok(None);
}
if bytes.len() < 4 {
return Err(ArrowError::IoError(
"IPC: truncated header".into(),
std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated IPC header"),
));
}
let has_continuation = bytes[..4] == IPC_CONTINUATION;
let (size_bytes, prefix_len): ([u8; 4], usize) = if has_continuation {
if bytes.len() < 8 {
return Err(ArrowError::IoError(
"IPC: truncated header after continuation".into(),
std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"truncated after continuation",
),
));
}
(bytes[4..8].try_into().unwrap(), 8)
} else {
(bytes[..4].try_into().unwrap(), 4)
};
let meta_size = u32::from_le_bytes(size_bytes) as usize;
if meta_size == 0 {
return Ok(None); // EOS
}
let meta_end = prefix_len + meta_size;
if bytes.len() < meta_end {
return Err(ArrowError::IoError(
"IPC: truncated metadata".into(),
std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated IPC metadata"),
));
}
let msg = root_as_message(&bytes[prefix_len..meta_end])
.map_err(|e| ArrowError::ParseError(format!("IPC message parse error: {e}")))?;
let body_len = msg.bodyLength() as usize;
let total = meta_end + body_len;
if bytes.len() < total {
return Err(ArrowError::IoError(
"IPC: truncated body".into(),
std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated IPC body"),
));
}
// Zero-copy: Bytes::slice shares the backing allocation; Buffer::from
// wraps it without copying.
Ok(Some(Buffer::from(data.slice(0..total))))
}
/// Read a length-prefixed byte slice at `offset` in `data`, advancing `offset`.
///
/// Reads an 8-byte little-endian length, then slices exactly that many bytes
/// from `data`. The returned [`Bytes`] is zero-copy (shares `data`'s allocation).
pub fn read_len_prefixed_bytes_at(data: &Bytes, offset: &mut usize) -> Result<Bytes, ArrowError> {
let bytes = data.as_ref();
let len_end = offset
.checked_add(8)
.filter(|&e| e <= bytes.len())
.ok_or_else(|| {
ArrowError::IoError(
"length-prefixed bytes: truncated length field".into(),
std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated length"),
)
})?;
let len = u64::from_le_bytes(bytes[*offset..len_end].try_into().unwrap()) as usize;
*offset = len_end;
let data_end = offset
.checked_add(len)
.filter(|&e| e <= bytes.len())
.ok_or_else(|| {
ArrowError::IoError(
"length-prefixed bytes: truncated data".into(),
std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated data"),
)
})?;
let result = data.slice(*offset..data_end);
*offset = data_end;
Ok(result)
}
/// Read all [`RecordBatch`]es from one Arrow IPC stream starting at `offset`,
/// advancing `offset` past the stream (including the EOS marker).
///
/// Zero-copy: array buffers borrow from `data`'s allocation.
pub fn read_ipc_stream_at(
data: &Bytes,
offset: &mut usize,
) -> Result<Vec<RecordBatch>, ArrowError> {
let batches = read_ipc_stream(&data.slice(*offset..))?;
// Recompute how many bytes were consumed by re-parsing message sizes.
// We can't get this from read_ipc_stream directly, so we re-walk the
// message headers (metadata only, no body re-read) to sum up lengths.
let slice = &data.as_ref()[*offset..];
let mut consumed = 0usize;
loop {
let rem = &slice[consumed..];
if rem.is_empty() {
break;
}
let has_cont = rem.len() >= 4 && rem[..4] == IPC_CONTINUATION;
let (size_bytes, prefix_len): ([u8; 4], usize) = if has_cont {
if rem.len() < 8 {
break;
}
(rem[4..8].try_into().unwrap(), 8)
} else {
if rem.len() < 4 {
break;
}
(rem[..4].try_into().unwrap(), 4)
};
let meta_size = u32::from_le_bytes(size_bytes) as usize;
if meta_size == 0 {
// EOS — consume it and stop.
consumed += prefix_len;
break;
}
let meta_end = prefix_len + meta_size;
if rem.len() < meta_end {
break;
}
let msg = root_as_message(&rem[prefix_len..meta_end])
.map_err(|e| ArrowError::ParseError(format!("IPC message parse error: {e}")))?;
let body_len = msg.bodyLength() as usize;
consumed += meta_end + body_len;
}
*offset += consumed;
Ok(batches)
}
/// Read exactly one [`RecordBatch`] from one Arrow IPC stream starting at `offset`,
/// advancing `offset` past the stream (including the EOS marker).
///
/// Zero-copy: array buffers borrow from `data`'s allocation.
pub fn read_ipc_stream_single_at(
data: &Bytes,
offset: &mut usize,
) -> Result<RecordBatch, ArrowError> {
let mut batches = read_ipc_stream_at(data, offset)?;
match batches.len() {
1 => Ok(batches.remove(0)),
n => Err(ArrowError::ParseError(format!(
"expected exactly 1 IPC record batch, got {n}"
))),
}
}
/// Extract the prefix length and metadata size from a raw IPC message buffer.
///
/// Modern IPC streams have an 8-byte prefix `[continuation: 4][size: 4]`.
/// Legacy streams have a 4-byte prefix `[size: 4]`. Returns `(prefix_len, meta_size)`.
fn parse_ipc_message_prefix(buf: &Buffer) -> Result<(usize, usize), ArrowError> {
let has_continuation = buf.len() >= 4 && buf[..4] == IPC_CONTINUATION;
if has_continuation {
if buf.len() < 8 {
return Err(ArrowError::ParseError(
"IPC message buffer too short".into(),
));
}
let meta_size = u32::from_le_bytes(buf[4..8].try_into().unwrap()) as usize;
Ok((8, meta_size))
} else {
if buf.len() < 4 {
return Err(ArrowError::ParseError(
"IPC message buffer too short".into(),
));
}
let meta_size = u32::from_le_bytes(buf[..4].try_into().unwrap()) as usize;
Ok((4, meta_size))
}
}
/// Read all [`RecordBatch`]es from one Arrow IPC stream.
///
/// Zero-copy: each batch's array data buffers are borrowed from the input
/// message buffer(s) and not copied during decoding.
///
/// Uses [`FileDecoder`] directly (rather than `StreamDecoder`) to avoid a
/// known edge case where `StreamDecoder` does not produce a batch for messages
/// with a zero-length body when the message exactly fills the decode buffer.
pub fn read_ipc_stream(data: &Bytes) -> Result<Vec<RecordBatch>, ArrowError> {
let mut offset = 0usize;
let schema_buf = read_one_ipc_message(&data.slice(offset..))?.ok_or_else(|| {
ArrowError::ParseError("IPC stream: expected schema message, got EOS".into())
})?;
offset += schema_buf.len();
let (prefix_len, meta_size) = parse_ipc_message_prefix(&schema_buf)?;
let schema_msg = root_as_message(&schema_buf[prefix_len..prefix_len + meta_size])
.map_err(|e| ArrowError::ParseError(format!("IPC schema parse error: {e}")))?;
let schema = Arc::new(fb_to_schema(schema_msg.header_as_schema().ok_or_else(
|| ArrowError::ParseError("IPC stream: first message is not a schema".into()),
)?));
let mut decoder = FileDecoder::new(schema, schema_msg.version());
let mut batches = Vec::new();
loop {
let Some(buf) = read_one_ipc_message(&data.slice(offset..))? else {
break;
};
offset += buf.len();
let (prefix_len, meta_size) = parse_ipc_message_prefix(&buf)?;
let msg = root_as_message(&buf[prefix_len..prefix_len + meta_size])
.map_err(|e| ArrowError::ParseError(format!("IPC message parse error: {e}")))?;
let body_len = msg.bodyLength() as usize;
// Block offset = 0 since the buffer starts at the message boundary.
// metaDataLength = prefix_len + meta_size (prefix + flatbuf + padding).
let block = arrow_ipc::Block::new(0, (prefix_len + meta_size) as i32, body_len as i64);
match msg.header_type() {
arrow_ipc::MessageHeader::RecordBatch => {
if let Some(batch) = decoder.read_record_batch(&block, &buf)? {
batches.push(batch);
}
}
arrow_ipc::MessageHeader::DictionaryBatch => {
decoder.read_dictionary(&block, &buf)?;
}
_ => break,
}
}
Ok(batches)
}
/// Read exactly one [`RecordBatch`] from one Arrow IPC stream.
pub fn read_ipc_stream_single(data: &Bytes) -> Result<RecordBatch, ArrowError> {
let mut batches = read_ipc_stream(data)?;
match batches.len() {
1 => Ok(batches.remove(0)),
n => Err(ArrowError::ParseError(format!(
"expected exactly 1 IPC record batch, got {n}"
))),
}
}
// ---------------------------------------------------------------------------
// Aligned IPC sections
// ---------------------------------------------------------------------------
/// Byte alignment that each IPC section's stream start is padded to.
///
/// When several IPC streams are concatenated into one larger blob (e.g. a
/// cache entry), a section that starts at an arbitrary offset would leave its
/// array data misaligned. [`FileDecoder`] with `require_alignment = false`
/// then silently copies each buffer into a freshly aligned allocation on
/// every read, defeating zero-copy. Padding each section start to a 64-byte
/// boundary keeps the decoded buffers borrowed directly from the input.
pub const IPC_SECTION_ALIGNMENT: usize = 64;
/// Number of zero-padding bytes needed to advance `pos` to the next
/// [`IPC_SECTION_ALIGNMENT`] boundary.
fn section_padding(pos: usize) -> usize {
(IPC_SECTION_ALIGNMENT - (pos % IPC_SECTION_ALIGNMENT)) % IPC_SECTION_ALIGNMENT
}
/// A [`Write`] adapter that counts the bytes written through it.
struct CountingWriter<'a> {
inner: &'a mut dyn Write,
count: usize,
}
impl Write for CountingWriter<'_> {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
let n = self.inner.write(buf)?;
self.count += n;
Ok(n)
}
fn flush(&mut self) -> std::io::Result<()> {
self.inner.flush()
}
}
/// Write zero padding so the next byte lands on an [`IPC_SECTION_ALIGNMENT`]
/// boundary, advancing `pos` past it.
fn write_section_padding(writer: &mut dyn Write, pos: &mut usize) -> Result<(), ArrowError> {
let pad = section_padding(*pos);
if pad > 0 {
const ZEROS: [u8; IPC_SECTION_ALIGNMENT] = [0u8; IPC_SECTION_ALIGNMENT];
writer
.write_all(&ZEROS[..pad])
.map_err(|e| ArrowError::IoError(e.to_string(), e))?;
*pos += pad;
}
Ok(())
}
/// Write `batch` as a 64-byte-aligned single-batch Arrow IPC section.
///
/// `pos` is the absolute byte offset of `writer` within the enclosing blob.
/// Zero padding is written first so the IPC stream begins on an
/// [`IPC_SECTION_ALIGNMENT`] boundary, then the stream itself. `pos` is
/// advanced past both the padding and the stream so the caller can write
/// further aligned sections.
///
/// Paired with [`read_ipc_section_at`]. For the decoded buffers to be borrowed
/// zero-copy, the blob must ultimately be read back from a buffer whose base
/// address is at least 64-byte aligned.
pub fn write_ipc_section(
writer: &mut dyn Write,
pos: &mut usize,
batch: &RecordBatch,
) -> Result<(), ArrowError> {
write_section_padding(writer, pos)?;
let mut counting = CountingWriter {
inner: writer,
count: 0,
};
write_ipc_stream(batch, &mut counting)?;
*pos += counting.count;
Ok(())
}
/// Read a single [`RecordBatch`] from an aligned IPC section at `offset`.
///
/// Skips the alignment padding written by [`write_ipc_section`], then reads
/// the stream, advancing `offset` past the section (padding + stream + EOS).
///
/// Zero-copy: array buffers borrow from `data`'s allocation when `data`'s base
/// address is at least 64-byte aligned (see [`write_ipc_section`]).
pub fn read_ipc_section_at(data: &Bytes, offset: &mut usize) -> Result<RecordBatch, ArrowError> {
*offset += section_padding(*offset);
read_ipc_stream_single_at(data, offset)
}
/// Write `batches` as a single 64-byte-aligned multi-batch Arrow IPC section.
///
/// Like [`write_ipc_section`] but emits every batch from `iter` into one IPC
/// stream (schema + N batches + EOS). `iter` must yield at least one batch.
/// Paired with [`read_ipc_section_batches_at`].
pub fn write_ipc_section_batches<I>(
writer: &mut dyn Write,
pos: &mut usize,
iter: I,
) -> Result<(), ArrowError>
where
I: IntoIterator<Item = RecordBatch>,
{
write_section_padding(writer, pos)?;
let mut counting = CountingWriter {
inner: writer,
count: 0,
};
write_ipc_stream_batches(iter, &mut counting)?;
*pos += counting.count;
Ok(())
}
/// Read all [`RecordBatch`]es from an aligned multi-batch IPC section at
/// `offset`, advancing `offset` past the section (padding + stream + EOS).
///
/// Zero-copy: array buffers borrow from `data`'s allocation when `data`'s base
/// address is at least 64-byte aligned (see [`write_ipc_section_batches`]).
pub fn read_ipc_section_batches_at(
data: &Bytes,
offset: &mut usize,
) -> Result<Vec<RecordBatch>, ArrowError> {
*offset += section_padding(*offset);
read_ipc_stream_at(data, offset)
}
#[cfg(test)]
mod tests {
use arrow_array::{ArrayRef, record_batch};
use super::*;
#[test]
fn test_ipc_roundtrip() {
let batch1 = record_batch!(
("int", Int32, [1, 2, 3]),
("str", Utf8, ["foo", "bar", "baz"])
)
.unwrap();
let batch2 = record_batch!(("int", Int32, [4, 5]), ("str", Utf8, ["qux", "quux"])).unwrap();
let batches = vec![batch1.clone(), batch2.clone()];
let mut buf = Vec::new();
write_ipc_stream_batches(batches, &mut buf).unwrap();
let data = Bytes::from(buf);
let batches = read_ipc_stream(&data).unwrap();
assert_eq!(batches.len(), 2);
assert_eq!(batches[0], batch1);
assert_eq!(batches[1], batch2);
let data_base = data.as_ptr() as usize;
let data_end = data_base + data.len();
let assert_col_zero_copy = |array: &ArrayRef| {
for buffer in array.to_data().buffers() {
let ptr = buffer.as_ptr() as usize;
assert!(
ptr >= data_base && ptr < data_end,
"buffer at {ptr:#x} is not backed by the input Bytes allocation \
[{data_base:#x}..{data_end:#x})"
);
}
};
for batch in &batches {
assert_eq!(batch.schema(), batch1.schema());
assert_col_zero_copy(batch.column(0));
assert_col_zero_copy(batch.column(1));
}
}
/// Allocate a [`Bytes`] whose base address is 64-byte aligned, modelling a
/// backend that reads cache entries into an aligned buffer. A plain
/// `Bytes::from(vec)` only guarantees the allocator's alignment for `u8`.
fn aligned_bytes(payload: &[u8]) -> Bytes {
let mut v = vec![0u8; payload.len() + IPC_SECTION_ALIGNMENT];
let pad = section_padding(v.as_ptr() as usize);
v[pad..pad + payload.len()].copy_from_slice(payload);
Bytes::from(v).slice(pad..pad + payload.len())
}
#[test]
fn test_aligned_ipc_sections_are_zero_copy() {
// A LargeBinary column exercises the i64-offset buffer whose 8-byte
// alignment requirement triggers a realigning memcpy when misaligned.
let blocks = arrow_array::LargeBinaryArray::from_vec(vec![&b"hello"[..], b"world"]);
let section_a = RecordBatch::try_from_iter([("a", Arc::new(blocks) as ArrayRef)]).unwrap();
let section_b = record_batch!(("b", Int64, [10i64, 20, 30, 40, 50])).unwrap();
let mut buf = Vec::new();
// Arbitrary, deliberately non-64-aligned preamble so the first section
// must be padded rather than landing at offset 0 by luck.
buf.extend_from_slice(&[0xABu8; 7]);
let mut pos = buf.len();
// The first section's stream begins after padding the 7-byte preamble
// up to the next 64-byte boundary.
assert_eq!(7 + section_padding(7), IPC_SECTION_ALIGNMENT);
write_ipc_section(&mut buf, &mut pos, &section_a).unwrap();
write_ipc_section(&mut buf, &mut pos, &section_b).unwrap();
let data = aligned_bytes(&buf);
assert_eq!(
section_padding(data.as_ptr() as usize),
0,
"base not aligned"
);
let mut offset = 7;
let read_a = read_ipc_section_at(&data, &mut offset).unwrap();
let read_b = read_ipc_section_at(&data, &mut offset).unwrap();
assert_eq!(read_a, section_a);
assert_eq!(read_b, section_b);
let data_base = data.as_ptr() as usize;
let data_end = data_base + data.len();
for batch in [&read_a, &read_b] {
for buffer in batch.column(0).to_data().buffers() {
let ptr = buffer.as_ptr() as usize;
assert!(
ptr >= data_base && ptr < data_end,
"section buffer at {ptr:#x} was realigned out of the input \
[{data_base:#x}..{data_end:#x}) — misaligned section",
);
}
}
}
#[test]
fn test_aligned_multi_batch_section_roundtrip_zero_copy() {
// A multi-batch section (e.g. IVF SQ storage chunks) must round-trip
// every batch and decode the first batch's buffers zero-copy.
let b1 = record_batch!(("v", Int64, [1i64, 2, 3])).unwrap();
let b2 = record_batch!(("v", Int64, [4i64, 5])).unwrap();
let b3 = record_batch!(("v", Int64, [6i64])).unwrap();
let mut buf = vec![0xCDu8; 5];
let mut pos = buf.len();
write_ipc_section_batches(&mut buf, &mut pos, [b1.clone(), b2.clone(), b3.clone()])
.unwrap();
let data = aligned_bytes(&buf);
let mut offset = 5;
let read = read_ipc_section_batches_at(&data, &mut offset).unwrap();
assert_eq!(read, vec![b1, b2, b3]);
assert_eq!(offset, buf.len(), "offset should land at section end");
let data_base = data.as_ptr() as usize;
let data_end = data_base + data.len();
for buffer in read[0].column(0).to_data().buffers() {
let ptr = buffer.as_ptr() as usize;
assert!(
ptr >= data_base && ptr < data_end,
"first batch buffer at {ptr:#x} was realigned out of the input",
);
}
}
}
python/vendor/lance-arrow/src/json.rs
Vendored
+1538
View File
File diff suppressed because it is too large Load Diff
python/vendor/lance-arrow/src/lib.rs
Vendored
+3102
View File
File diff suppressed because it is too large Load Diff

Some files were not shown because too many files have changed in this diff Show More