mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-19 02:02:27 +00:00
Merge main into gatekeeper/fix-2325-1
This commit is contained in:
+2
-2
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb-python"
|
||||
version = "0.40.0-beta.11"
|
||||
version = "0.41.0-beta.7"
|
||||
publish = false
|
||||
edition.workspace = true
|
||||
description = "Python bindings for LanceDB"
|
||||
@@ -8,7 +8,7 @@ license.workspace = true
|
||||
repository.workspace = true
|
||||
keywords.workspace = true
|
||||
categories.workspace = true
|
||||
rust-version = "1.91.0"
|
||||
rust-version = "1.94.0"
|
||||
|
||||
[lib]
|
||||
name = "_lancedb"
|
||||
|
||||
@@ -10,7 +10,7 @@ dependencies = [
|
||||
"pyarrow>=16",
|
||||
"pydantic>=2.7.4,<3",
|
||||
"tqdm>=4.27.0",
|
||||
"lance-namespace>=0.3.2"
|
||||
"lance-namespace>=0.13.0"
|
||||
]
|
||||
description = "lancedb"
|
||||
authors = [{ name = "LanceDB Devs", email = "dev@lancedb.com" }]
|
||||
@@ -63,7 +63,7 @@ tests = [
|
||||
"polars>=0.19, <=1.32.3",
|
||||
"pyarrow<25",
|
||||
"pyarrow-stubs>=16.0",
|
||||
"pylance==9.0.0",
|
||||
"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"]
|
||||
|
||||
@@ -15,7 +15,7 @@ from ._lancedb import FtsToken
|
||||
from ._lancedb import FunctionErrorFragment as FunctionErrorFragment
|
||||
from ._lancedb import FunctionErrorRecord as FunctionErrorRecord
|
||||
from ._lancedb import FunctionErrors as FunctionErrors
|
||||
from ._lancedb import LsmWriteSpec
|
||||
from ._lancedb import MemWalWriteSpec
|
||||
from ._lancedb import tokenize as _tokenize
|
||||
from .common import URI, sanitize_uri
|
||||
from urllib.parse import urlparse
|
||||
@@ -46,7 +46,6 @@ from .materialized_view import (
|
||||
AsyncMaterializedView,
|
||||
MaterializedView,
|
||||
MaterializedViewDefinition,
|
||||
vector_dedup as vector_dedup,
|
||||
)
|
||||
from .view import ViewDescription as ViewDescription
|
||||
from .table import AsyncTable, CompactionOptions, Table
|
||||
@@ -62,7 +61,6 @@ from .namespace import (
|
||||
from .catalog import (
|
||||
AsyncCatalog,
|
||||
Catalog,
|
||||
ListDatabasesResponse,
|
||||
connect_catalog,
|
||||
connect_catalog_async,
|
||||
)
|
||||
@@ -323,8 +321,8 @@ def tokenize(
|
||||
language: str = "English",
|
||||
max_token_length: Optional[int] = 40,
|
||||
lower_case: bool = True,
|
||||
stem: bool = True,
|
||||
remove_stop_words: bool = True,
|
||||
stem: Optional[bool] = None,
|
||||
remove_stop_words: Optional[bool] = None,
|
||||
custom_stop_words: Optional[List[str]] = None,
|
||||
ascii_folding: bool = True,
|
||||
ngram_min_length: int = 3,
|
||||
@@ -335,6 +333,8 @@ def tokenize(
|
||||
|
||||
This does not require an FTS index. The tokenizer options match
|
||||
:class:`lancedb.index.FTS`. ``custom_stop_words`` accepts a list of strings.
|
||||
When ``stem`` or ``remove_stop_words`` is None, the base tokenizer's default
|
||||
applies: False for ``code`` and ``ngram``, True otherwise.
|
||||
"""
|
||||
|
||||
return _tokenize(
|
||||
@@ -578,7 +578,6 @@ async def connect_async(
|
||||
__all__ = [
|
||||
"Catalog",
|
||||
"AsyncCatalog",
|
||||
"ListDatabasesResponse",
|
||||
"connect_catalog",
|
||||
"connect_catalog_async",
|
||||
"AsyncMaterializedView",
|
||||
@@ -612,7 +611,7 @@ __all__ = [
|
||||
"SqlQuery",
|
||||
"LanceDBConnection",
|
||||
"LanceNamespaceDBConnection",
|
||||
"LsmWriteSpec",
|
||||
"MemWalWriteSpec",
|
||||
"RemoteDBConnection",
|
||||
"Session",
|
||||
"Table",
|
||||
|
||||
+250
-53
@@ -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]]:
|
||||
|
||||
@@ -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
|
||||
@@ -58,8 +58,8 @@ def tokenize(
|
||||
language: str = "English",
|
||||
max_token_length: Optional[int] = 40,
|
||||
lower_case: bool = True,
|
||||
stem: bool = True,
|
||||
remove_stop_words: bool = True,
|
||||
stem: Optional[bool] = None,
|
||||
remove_stop_words: Optional[bool] = None,
|
||||
custom_stop_words: Optional[List[str]] = None,
|
||||
ascii_folding: bool = True,
|
||||
ngram_min_length: int = 3,
|
||||
@@ -149,12 +149,24 @@ class Connection(object):
|
||||
limit: Optional[int],
|
||||
) -> list[str]: ... # Deprecated: Use list_tables instead
|
||||
async def open_job(self, job_id: str) -> Job: ...
|
||||
async def create_function_async(self, request_json: str) -> Job: ...
|
||||
async def get_function(self, name: str, version: str) -> str: ...
|
||||
async def list_functions(self) -> List[str]: ...
|
||||
async def drop_function(self, name: str, version: str) -> bool: ...
|
||||
async def create_function_async(
|
||||
self, request_json: str, namespace_path: Optional[List[str]] = None
|
||||
) -> Job: ...
|
||||
async def get_function(
|
||||
self, name: str, version: str, namespace_path: Optional[List[str]] = None
|
||||
) -> str: ...
|
||||
def list_functions(
|
||||
self,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
*,
|
||||
page_token: Optional[str] = None,
|
||||
page_limit: Optional[int] = None,
|
||||
) -> FunctionListing: ...
|
||||
async def drop_function(
|
||||
self, name: str, version: str, namespace_path: Optional[List[str]] = None
|
||||
) -> bool: ...
|
||||
async def drop_function_async(
|
||||
self, name: str, version: str
|
||||
self, name: str, version: str, namespace_path: Optional[List[str]] = None
|
||||
) -> Tuple[bool, Job]: ...
|
||||
async def create_secret(
|
||||
self, name: str, value: str, namespace_path: Optional[List[str]] = None
|
||||
@@ -162,9 +174,13 @@ class Connection(object):
|
||||
async def alter_secret(
|
||||
self, name: str, value: str, namespace_path: Optional[List[str]] = None
|
||||
) -> None: ...
|
||||
async def list_secrets(
|
||||
self, namespace_path: Optional[List[str]] = None
|
||||
) -> List[str]: ...
|
||||
def list_secrets(
|
||||
self,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
*,
|
||||
page_token: Optional[str] = None,
|
||||
page_limit: Optional[int] = None,
|
||||
) -> NameListing: ...
|
||||
async def drop_secret(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
) -> None: ...
|
||||
@@ -183,10 +199,16 @@ class Connection(object):
|
||||
async def drop_view_async(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
) -> Job: ...
|
||||
async def list_views(
|
||||
self, namespace_path: Optional[List[str]] = None
|
||||
) -> List[str]: ...
|
||||
async def list_jobs(self) -> List[JobInfo]: ...
|
||||
def list_views(
|
||||
self,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
*,
|
||||
page_token: Optional[str] = None,
|
||||
page_limit: Optional[int] = None,
|
||||
) -> NameListing: ...
|
||||
def list_jobs(
|
||||
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||
) -> JobListing: ...
|
||||
async def cancel_job(self, job_id: str) -> bool: ...
|
||||
async def pause_job(self, job_id: str) -> str: ...
|
||||
async def resume_job(self, job_id: str) -> str: ...
|
||||
@@ -195,6 +217,7 @@ class Connection(object):
|
||||
query: str,
|
||||
*,
|
||||
default_namespace_path: Optional[List[str]] = None,
|
||||
parameters: Optional[pa.RecordBatch] = None,
|
||||
) -> SqlQuery: ...
|
||||
async def describe_query(self, query_id: UUID) -> QueryDescription: ...
|
||||
async def create_table(
|
||||
@@ -247,8 +270,6 @@ class Connection(object):
|
||||
filter: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
with_no_data: bool = False,
|
||||
*,
|
||||
vector_source_json: Optional[str] = None,
|
||||
) -> Table: ...
|
||||
async def create_materialized_view_async(
|
||||
self,
|
||||
@@ -258,10 +279,10 @@ class Connection(object):
|
||||
filter: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
with_no_data: bool = False,
|
||||
*,
|
||||
vector_source_json: Optional[str] = None,
|
||||
) -> Job: ...
|
||||
async def list_materialized_views(self) -> List[str]: ...
|
||||
def list_materialized_views(
|
||||
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||
) -> NameListing: ...
|
||||
async def drop_materialized_view(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
) -> None: ...
|
||||
@@ -290,6 +311,7 @@ class BlobFile:
|
||||
def tell(self) -> int: ...
|
||||
def size(self) -> int: ...
|
||||
def read_range(self, offset: int, length: int) -> bytes: ...
|
||||
def read_ranges(self, ranges: list[tuple[int, int]]) -> list[bytes]: ...
|
||||
def read_up_to(self, length: int) -> bytes: ...
|
||||
|
||||
class Job:
|
||||
@@ -402,6 +424,7 @@ class Table:
|
||||
def name(self) -> str: ...
|
||||
def __repr__(self) -> str: ...
|
||||
def is_open(self) -> bool: ...
|
||||
def _is_native(self) -> bool: ...
|
||||
def close(self) -> None: ...
|
||||
async def schema(self) -> pa.Schema: ...
|
||||
async def add(
|
||||
@@ -492,10 +515,10 @@ class Table:
|
||||
limit: Optional[int] = None,
|
||||
) -> FunctionErrors: ...
|
||||
async def refresh_materialized_view(
|
||||
self, full: bool = False, source_version: Optional[int] = None
|
||||
self, source_version: Optional[int] = None
|
||||
) -> RefreshMaterializedViewResult: ...
|
||||
async def refresh_materialized_view_async(
|
||||
self, full: bool = False, source_version: Optional[int] = None
|
||||
self, source_version: Optional[int] = None
|
||||
) -> Job: ...
|
||||
async def materialized_view_definition(self) -> str: ...
|
||||
async def add_columns_with_schema(self, schema: pa.Schema) -> AddColumnsResult: ...
|
||||
@@ -516,17 +539,17 @@ class Table:
|
||||
async def initial_storage_options(self) -> Optional[Dict[str, str]]: ...
|
||||
async def latest_storage_options(self) -> Optional[Dict[str, str]]: ...
|
||||
async def set_unenforced_primary_key(self, columns: List[str]) -> None: ...
|
||||
async def set_lsm_write_spec(self, spec: LsmWriteSpec) -> None: ...
|
||||
async def unset_lsm_write_spec(self) -> None: ...
|
||||
async def get_lsm_write_spec(self) -> Optional[LsmWriteSpec]: ...
|
||||
async def enable_memwal(self, spec: MemWalWriteSpec) -> None: ...
|
||||
async def disable_memwal(self) -> None: ...
|
||||
async def get_memwal_spec(self) -> Optional[MemWalWriteSpec]: ...
|
||||
# WAL-PK-FUSION: delete both.
|
||||
def hybrid_pk_fusion_learned(self) -> bool: ...
|
||||
def note_hybrid_pk_fusion(self) -> None: ...
|
||||
async def checkpoint_lsm(self) -> None: ...
|
||||
async def flush_lsm(self) -> None: ...
|
||||
async def compact_lsm(self) -> None: ...
|
||||
async def get_lsm_stats(self, include_generation_rows: bool) -> Optional[dict]: ...
|
||||
async def close_lsm_writers(self) -> None: ...
|
||||
async def checkpoint_memwal(self) -> None: ...
|
||||
async def flush_memwal(self) -> None: ...
|
||||
async def compact_memwal(self) -> None: ...
|
||||
async def get_memwal_stats(self, include_sstable_rows: bool) -> Optional[dict]: ...
|
||||
async def close_memwal_writers(self) -> None: ...
|
||||
@property
|
||||
def tags(self) -> Tags: ...
|
||||
@property
|
||||
@@ -542,7 +565,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]
|
||||
@@ -639,7 +662,7 @@ class Query:
|
||||
def fast_search(self): ...
|
||||
def with_row_id(self): ...
|
||||
def postfilter(self): ...
|
||||
def use_lsm(self, enable: bool): ...
|
||||
def use_memwal(self, enable: bool): ...
|
||||
def nearest_to(self, query_vec: pa.Array) -> VectorQuery: ...
|
||||
def nearest_to_text(self, query: dict) -> FTSQuery: ...
|
||||
def order_by(self, ordering: Optional[List[ColumnOrdering]]): ...
|
||||
@@ -656,7 +679,7 @@ class Query:
|
||||
class TakeQuery:
|
||||
def select(self, columns: List[str]): ...
|
||||
def with_row_id(self): ...
|
||||
def use_lsm(self, enable: bool): ...
|
||||
def use_memwal(self, enable: bool): ...
|
||||
async def output_schema(self) -> pa.Schema: ...
|
||||
async def execute(self) -> RecordBatchStream: ...
|
||||
async def explain_plan(self, verbose: Optional[bool]) -> str: ...
|
||||
@@ -675,7 +698,7 @@ class FTSQuery:
|
||||
def fast_search(self): ...
|
||||
def with_row_id(self): ...
|
||||
def postfilter(self): ...
|
||||
def use_lsm(self, enable: bool): ...
|
||||
def use_memwal(self, enable: bool): ...
|
||||
def get_query(self) -> str: ...
|
||||
def add_query_vector(self, query_vec: pa.Array) -> None: ...
|
||||
def nearest_to(self, query_vec: pa.Array) -> HybridQuery: ...
|
||||
@@ -703,9 +726,10 @@ class VectorQuery:
|
||||
def column(self, column: str): ...
|
||||
def distance_type(self, distance_type: str): ...
|
||||
def postfilter(self): ...
|
||||
def use_lsm(self, enable: bool): ...
|
||||
def use_memwal(self, enable: bool): ...
|
||||
def refine_factor(self, refine_factor: int): ...
|
||||
def nprobes(self, nprobes: int): ...
|
||||
def search_effort(self, search_effort: Optional[float]): ...
|
||||
def minimum_nprobes(self, minimum_nprobes: int): ...
|
||||
def maximum_nprobes(self, maximum_nprobes: int): ...
|
||||
def bypass_vector_index(self): ...
|
||||
@@ -727,10 +751,11 @@ class HybridQuery:
|
||||
def fast_search(self): ...
|
||||
def with_row_id(self): ...
|
||||
def postfilter(self): ...
|
||||
def use_lsm(self, enable: bool): ...
|
||||
def use_memwal(self, enable: bool): ...
|
||||
def distance_type(self, distance_type: str): ...
|
||||
def refine_factor(self, refine_factor: int): ...
|
||||
def nprobes(self, nprobes: int): ...
|
||||
def search_effort(self, search_effort: Optional[float]): ...
|
||||
def minimum_nprobes(self, minimum_nprobes: int): ...
|
||||
def maximum_nprobes(self, maximum_nprobes: int): ...
|
||||
def bypass_vector_index(self): ...
|
||||
@@ -754,11 +779,13 @@ class PyQueryRequest:
|
||||
select_source_columns: Optional[Dict[str, str]]
|
||||
fast_search: Optional[bool]
|
||||
with_row_id: Optional[bool]
|
||||
use_lsm: Optional[bool]
|
||||
use_memwal: Optional[bool]
|
||||
column: Optional[str]
|
||||
query_vector: Optional[List[pa.Array]]
|
||||
nprobes: Optional[int]
|
||||
minimum_nprobes: Optional[int]
|
||||
maximum_nprobes: Optional[int]
|
||||
search_effort: Optional[float]
|
||||
lower_bound: Optional[float]
|
||||
upper_bound: Optional[float]
|
||||
ef: Optional[int]
|
||||
@@ -813,22 +840,26 @@ class FtsToken:
|
||||
text: str
|
||||
position: int
|
||||
|
||||
class LsmWriteSpec:
|
||||
class MemWalWriteSpec:
|
||||
"""Specification selecting Lance's MemWAL LSM-style write path for
|
||||
`merge_insert`."""
|
||||
|
||||
@staticmethod
|
||||
def bucket(column: str, num_buckets: int) -> "LsmWriteSpec": ...
|
||||
def bucket(column: str, num_buckets: int) -> "MemWalWriteSpec": ...
|
||||
@staticmethod
|
||||
def identity(column: str) -> "LsmWriteSpec": ...
|
||||
def identity(column: str) -> "MemWalWriteSpec": ...
|
||||
@staticmethod
|
||||
def unsharded() -> "LsmWriteSpec": ...
|
||||
def with_maintained_indexes(self, indexes: Optional[List[str]]) -> "LsmWriteSpec":
|
||||
def unsharded() -> "MemWalWriteSpec": ...
|
||||
def with_maintained_indexes(
|
||||
self, indexes: Optional[List[str]]
|
||||
) -> "MemWalWriteSpec":
|
||||
"""Set which indexes the MemWAL keeps up to date. None resolves every
|
||||
index on the table at install, failing if one cannot be maintained;
|
||||
a list is verbatim, empty means none."""
|
||||
...
|
||||
def with_writer_config_defaults(self, defaults: Dict[str, str]) -> "LsmWriteSpec":
|
||||
def with_writer_config_defaults(
|
||||
self, defaults: Dict[str, str]
|
||||
) -> "MemWalWriteSpec":
|
||||
"""Return a copy of this spec recording the given default
|
||||
`ShardWriter` configuration in the MemWAL index."""
|
||||
...
|
||||
@@ -930,6 +961,8 @@ class PermutationReader:
|
||||
def __init__(self, base_table: Table, permutation_table: Table): ...
|
||||
|
||||
class Catalog:
|
||||
@property
|
||||
def authz(self) -> AuthorizationClient: ...
|
||||
@property
|
||||
def uri(self) -> str: ...
|
||||
async def create_database(
|
||||
@@ -939,9 +972,15 @@ class Catalog:
|
||||
async def drop_database(
|
||||
self, name: str, *, ignore_missing: bool = False
|
||||
) -> None: ...
|
||||
async def list_databases(
|
||||
self, *, limit: Optional[int] = None, page_token: Optional[str] = None
|
||||
) -> tuple[list[str], Optional[str]]: ...
|
||||
def list_databases(
|
||||
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||
) -> DatabaseNames: ...
|
||||
|
||||
class DatabaseNames:
|
||||
def num_page_results(self) -> int: ...
|
||||
def page_token(self) -> Optional[str]: ...
|
||||
def __aiter__(self) -> DatabaseNames: ...
|
||||
async def __anext__(self) -> str: ...
|
||||
|
||||
async def connect_catalog(
|
||||
endpoint: str,
|
||||
@@ -952,3 +991,78 @@ async def connect_catalog(
|
||||
read_consistency_interval: Optional[float] = None,
|
||||
oauth_config: Optional[Any] = None,
|
||||
) -> Catalog: ...
|
||||
|
||||
class NameListing:
|
||||
def num_page_results(self) -> int: ...
|
||||
def page_token(self) -> Optional[str]: ...
|
||||
def __aiter__(self) -> NameListing: ...
|
||||
async def __anext__(self) -> str: ...
|
||||
|
||||
class FunctionListing:
|
||||
def num_page_results(self) -> int: ...
|
||||
def page_token(self) -> Optional[str]: ...
|
||||
def __aiter__(self) -> FunctionListing: ...
|
||||
async def __anext__(self) -> str: ...
|
||||
|
||||
class JobListing:
|
||||
def num_page_results(self) -> int: ...
|
||||
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: ...
|
||||
|
||||
@@ -76,7 +76,7 @@ _WAL_ROW_ID_REFUSAL = "does not support with_row_id"
|
||||
WAL_ROW_ID_UNSUPPORTED = (
|
||||
"hybrid search on a MemWAL table cannot return _rowid: the fresh tier has "
|
||||
"no stable row id, and the ids the fusion joins on are synthesized from "
|
||||
"the primary key. Set use_lsm(False) to read the base table only (results "
|
||||
"the primary key. Set use_memwal(False) to read the base table only (results "
|
||||
"will exclude un-compacted MemWAL data)"
|
||||
)
|
||||
|
||||
@@ -151,13 +151,13 @@ def _stamp_surrogate_row_ids(
|
||||
return stamped
|
||||
|
||||
|
||||
def pk_fusion_learned(table, use_lsm: Optional[bool]) -> bool:
|
||||
def pk_fusion_learned(table, use_memwal: Optional[bool]) -> bool:
|
||||
"""Whether a previous hybrid query on `table` was refused `_rowid`, so this
|
||||
one should go straight to the primary key.
|
||||
|
||||
Learned from a refusal, never probed, so asking is free.
|
||||
"""
|
||||
if table is None or use_lsm is False:
|
||||
if table is None or use_memwal is False:
|
||||
return False
|
||||
# Only an unambiguous True forks to the key: a table type that answers the
|
||||
# question loosely keeps the `_rowid` path.
|
||||
@@ -192,7 +192,7 @@ class PkFusion:
|
||||
raise NotImplementedError(
|
||||
"hybrid search cannot project a blob column on a MemWAL table: "
|
||||
"fetching blobs needs a real _rowid, and the fresh tier has no "
|
||||
"stable one. Set use_lsm(False) to read the base table only "
|
||||
"stable one. Set use_memwal(False) to read the base table only "
|
||||
"(results will exclude un-compacted MemWAL data)"
|
||||
)
|
||||
return cls(_pk_columns(schema))
|
||||
|
||||
@@ -0,0 +1,588 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Deployment-wide authorization through ``catalog.authz``.
|
||||
|
||||
ACLs grant one privilege on an object to a principal, group, or role. Role
|
||||
bindings grant a role to a principal or group. Identity-provider group membership
|
||||
is read-only. Name resolution and permission checks are performed by the server.
|
||||
"""
|
||||
|
||||
from collections.abc import AsyncIterator, Iterator, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
import json
|
||||
from typing import Literal, Optional, TypeVar, Union, cast
|
||||
|
||||
from . import _lancedb
|
||||
from .background_loop import LOOP
|
||||
|
||||
__all__ = [
|
||||
"Authorization",
|
||||
"AsyncAuthorization",
|
||||
"Subject",
|
||||
"Object",
|
||||
"Privilege",
|
||||
"AccessControlEntry",
|
||||
"RoleBinding",
|
||||
"IdentityInfo",
|
||||
"PrincipalInfo",
|
||||
"GroupInfo",
|
||||
]
|
||||
|
||||
_PAGE_SIZE = 100
|
||||
_T = TypeVar("_T")
|
||||
_SubjectKind = Literal[
|
||||
"principal_id",
|
||||
"principal_name",
|
||||
"group_id",
|
||||
"group_name",
|
||||
"role",
|
||||
"principal_api_key",
|
||||
"unknown",
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Subject:
|
||||
"""An immutable identity selector. Prefer the named factory methods.
|
||||
|
||||
IDs are stable; names must resolve unambiguously on the server. API-key
|
||||
selectors identify the subject of an operation, not its authenticated caller.
|
||||
Secret values are omitted from representations. Unknown canonical subject
|
||||
kinds returned by newer servers use ``kind="unknown"`` and preserve the full
|
||||
wire string in ``value``; they can be passed back to ACL methods.
|
||||
"""
|
||||
|
||||
kind: _SubjectKind
|
||||
_value: str = field(repr=False)
|
||||
|
||||
_inner: _lancedb.AuthzSubject = field(init=False, repr=False, compare=False)
|
||||
|
||||
def __post_init__(self):
|
||||
object.__setattr__(
|
||||
self, "_inner", _lancedb.AuthzSubject(self.kind, self._value)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def principal_id(cls, id: str) -> "Subject":
|
||||
"""Select a principal by canonical identity-provider ID."""
|
||||
return cls("principal_id", id)
|
||||
|
||||
@classmethod
|
||||
def principal_name(cls, name: str) -> "Subject":
|
||||
"""Select a principal by an unambiguous provider name."""
|
||||
return cls("principal_name", name)
|
||||
|
||||
@classmethod
|
||||
def group_id(cls, id: str) -> "Subject":
|
||||
"""Select a group by canonical identity-provider ID."""
|
||||
return cls("group_id", id)
|
||||
|
||||
@classmethod
|
||||
def group_name(cls, name: str) -> "Subject":
|
||||
"""Select a group by an unambiguous provider name."""
|
||||
return cls("group_name", name)
|
||||
|
||||
@classmethod
|
||||
def role(cls, name: str) -> "Subject":
|
||||
"""Select a LanceDB-managed role, without an ``r:`` prefix."""
|
||||
return cls("role", name)
|
||||
|
||||
@classmethod
|
||||
def principal_api_key(cls, api_key: str) -> "Subject":
|
||||
"""Select the principal holding this secret API key."""
|
||||
return cls("principal_api_key", api_key)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
value = self._inner.display_value
|
||||
return f"Subject.{self.kind}({value!r})"
|
||||
|
||||
@property
|
||||
def value(self) -> str:
|
||||
"""The selected value, or full wire string for unknown subject kinds.
|
||||
|
||||
API-key values are secrets.
|
||||
"""
|
||||
return self._value
|
||||
|
||||
def __reduce__(self):
|
||||
return type(self), (self.kind, self._value)
|
||||
|
||||
def _wire(self) -> str:
|
||||
return self._inner.to_wire()
|
||||
|
||||
@classmethod
|
||||
def _canonical(cls, value: str) -> "Subject":
|
||||
inner = _lancedb.AuthzSubject.from_canonical(value)
|
||||
return cls(cast(_SubjectKind, inner.kind), inner.value)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Object:
|
||||
"""An immutable resource name in the authorization syntax.
|
||||
|
||||
Use the builders to avoid interpreting separators as part of a name.
|
||||
Methods also accept raw object strings, validated by the same Rust parser.
|
||||
Unknown resource types are preserved for compatibility with newer servers.
|
||||
Omitted or empty namespace paths select the database's public namespace.
|
||||
"""
|
||||
|
||||
_value: str
|
||||
_inner: _lancedb.AuthzObject = field(init=False, repr=False, compare=False)
|
||||
|
||||
def __post_init__(self):
|
||||
object.__setattr__(self, "_inner", _lancedb.AuthzObject(self._value))
|
||||
|
||||
@classmethod
|
||||
def system(cls) -> "Object":
|
||||
"""The deployment-wide system object."""
|
||||
return cls(str(_lancedb.AuthzObject.system()))
|
||||
|
||||
@classmethod
|
||||
def database(cls, name: str) -> "Object":
|
||||
"""A database; slash-containing names remain one logical database."""
|
||||
return cls(str(_lancedb.AuthzObject.database(name)))
|
||||
|
||||
@classmethod
|
||||
def namespace(
|
||||
cls, *, database: str, namespace_path: Sequence[str] = ("public",)
|
||||
) -> "Object":
|
||||
"""A namespace path, defaulting to the database's public namespace."""
|
||||
return cls(str(_lancedb.AuthzObject.namespace(database, namespace_path)))
|
||||
|
||||
@classmethod
|
||||
def table(
|
||||
cls, *, database: str, name: str, namespace_path: Sequence[str] = ("public",)
|
||||
) -> "Object":
|
||||
"""A table, defaulting to the database's public namespace."""
|
||||
return cls(str(_lancedb.AuthzObject.table(database, namespace_path, name)))
|
||||
|
||||
@classmethod
|
||||
def view(
|
||||
cls, *, database: str, name: str, namespace_path: Sequence[str] = ("public",)
|
||||
) -> "Object":
|
||||
"""A view, defaulting to the database's public namespace."""
|
||||
return cls(str(_lancedb.AuthzObject.view(database, namespace_path, name)))
|
||||
|
||||
@classmethod
|
||||
def secret(
|
||||
cls, *, database: str, name: str, namespace_path: Sequence[str] = ("public",)
|
||||
) -> "Object":
|
||||
"""A secret, defaulting to the database's public namespace."""
|
||||
return cls(str(_lancedb.AuthzObject.secret(database, namespace_path, name)))
|
||||
|
||||
@classmethod
|
||||
def function(
|
||||
cls, *, database: str, name: str, namespace_path: Sequence[str] = ("public",)
|
||||
) -> "Object":
|
||||
"""A function, defaulting to the database's public namespace."""
|
||||
return cls(str(_lancedb.AuthzObject.function(database, namespace_path, name)))
|
||||
|
||||
def __str__(self) -> str:
|
||||
return str(self._inner)
|
||||
|
||||
def __reduce__(self):
|
||||
return type(self), (self._value,)
|
||||
|
||||
|
||||
class Privilege(str, Enum):
|
||||
"""Known privileges, validated and serialized by Rust.
|
||||
|
||||
ACL methods also accept strings such as ``"SELECT"`` and future privilege
|
||||
names such as ``"PRIVILEGE20"``. Future names are preserved verbatim;
|
||||
their meaning and validity are determined by the server.
|
||||
|
||||
Grant ``CREATE_SECRET`` or ``CREATE_FUNCTION`` on a namespace to delegate
|
||||
creation of those resources without transferring namespace ownership.
|
||||
"""
|
||||
|
||||
def __new__(cls, value: str):
|
||||
canonical = str(_lancedb.AuthzPrivilege(value))
|
||||
member = str.__new__(cls, canonical)
|
||||
member._value_ = canonical
|
||||
return member
|
||||
|
||||
OWNERSHIP = "OWNERSHIP"
|
||||
USAGE = "USAGE"
|
||||
SELECT = "SELECT"
|
||||
INSERT = "INSERT"
|
||||
UPDATE = "UPDATE"
|
||||
DELETE = "DELETE"
|
||||
CREATE_DATABASE = "CREATE_DATABASE"
|
||||
CREATE_NAMESPACE = "CREATE_NAMESPACE"
|
||||
CREATE_TABLE = "CREATE_TABLE"
|
||||
CREATE_MATERIALIZED_VIEW = "CREATE_MATERIALIZED_VIEW"
|
||||
CREATE_VIEW = "CREATE_VIEW"
|
||||
CREATE_SECRET = "CREATE_SECRET"
|
||||
CREATE_FUNCTION = "CREATE_FUNCTION"
|
||||
OPERATE = "OPERATE"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AccessControlEntry:
|
||||
"""One object/subject/privilege grant; subject contains a canonical ID or role."""
|
||||
|
||||
object: str
|
||||
subject: Subject
|
||||
subject_name: Optional[str]
|
||||
privilege: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RoleBinding:
|
||||
"""A direct principal/group role binding, not an expanded effective role."""
|
||||
|
||||
subject: Subject
|
||||
subject_name: Optional[str]
|
||||
role: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class IdentityInfo:
|
||||
"""A canonical identity-provider ID and its optional display name."""
|
||||
|
||||
id: str
|
||||
name: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PrincipalInfo:
|
||||
"""A principal's groups and effective roles, including group-inherited roles."""
|
||||
|
||||
principal_id: str
|
||||
principal_name: Optional[str]
|
||||
groups: list[IdentityInfo]
|
||||
roles: list[str]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GroupInfo:
|
||||
"""A group's provider-managed members and LanceDB-managed direct roles."""
|
||||
|
||||
group_id: str
|
||||
group_name: Optional[str]
|
||||
principals: list[IdentityInfo]
|
||||
roles: list[str]
|
||||
|
||||
|
||||
def _subject(subject: Subject, context: Optional[str] = None) -> str:
|
||||
if not isinstance(subject, Subject):
|
||||
raise ValueError("subject must be a Subject selector")
|
||||
return subject._inner.to_wire(context)
|
||||
|
||||
|
||||
def _object(value: Union[str, Object]) -> str:
|
||||
if not isinstance(value, (str, Object)) or not str(value):
|
||||
raise ValueError("object must be a nonempty string or Object")
|
||||
return str(value) if isinstance(value, Object) else str(_lancedb.AuthzObject(value))
|
||||
|
||||
|
||||
def _limit(limit: Optional[int]) -> None:
|
||||
if limit is not None and (
|
||||
isinstance(limit, bool) or not isinstance(limit, int) or limit < 0
|
||||
):
|
||||
raise ValueError("limit must be a nonnegative integer or None")
|
||||
|
||||
|
||||
def _privilege(value: Union[str, Privilege]) -> str:
|
||||
if not isinstance(value, str):
|
||||
raise ValueError("privilege must be a string or Privilege")
|
||||
return str(_lancedb.AuthzPrivilege(value))
|
||||
|
||||
|
||||
def _entry(value: dict) -> AccessControlEntry:
|
||||
return AccessControlEntry(
|
||||
value["object"],
|
||||
Subject._canonical(value["subject_id"]),
|
||||
value.get("subject_name"),
|
||||
value["privilege"],
|
||||
)
|
||||
|
||||
|
||||
def _binding(value: dict) -> RoleBinding:
|
||||
return RoleBinding(
|
||||
Subject._canonical(value["subject_id"]),
|
||||
value.get("subject_name"),
|
||||
value["role"],
|
||||
)
|
||||
|
||||
|
||||
class AsyncAuthorization:
|
||||
"""Async authorization returned by ``AsyncCatalog.authz``.
|
||||
|
||||
HTTP failures raise [HttpError][lancedb.remote.errors.HttpError], preserving
|
||||
status and request ID. Invalid local arguments raise ValueError. Lists are
|
||||
lazy async iterators: failures can occur during iteration after some results
|
||||
have been yielded. Iteration does not provide a consistent snapshot.
|
||||
"""
|
||||
|
||||
def __init__(self, inner: _lancedb.AuthorizationClient):
|
||||
self._inner = inner
|
||||
|
||||
async def _pages(self, method, request: dict, key: str, limit: Optional[int]):
|
||||
yielded = 0
|
||||
while limit is None or yielded < limit:
|
||||
request["max_results"] = (
|
||||
_PAGE_SIZE if limit is None else min(_PAGE_SIZE, limit - yielded)
|
||||
)
|
||||
response = json.loads(await method(json.dumps(request)))
|
||||
for value in response[key]:
|
||||
yield value
|
||||
yielded += 1
|
||||
if limit is not None and yielded >= limit:
|
||||
return
|
||||
token = response.get("continuation_token")
|
||||
if not token:
|
||||
return
|
||||
if token == request.get("continuation_token"):
|
||||
raise RuntimeError(
|
||||
"Authorization listing returned a repeated continuation token"
|
||||
)
|
||||
request["continuation_token"] = token
|
||||
|
||||
def list_acls(
|
||||
self,
|
||||
*,
|
||||
object: Optional[Union[str, Object]] = None,
|
||||
subject: Optional[Subject] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> AsyncIterator[AccessControlEntry]:
|
||||
"""Iterate grants, optionally filtered by exact object and/or subject.
|
||||
|
||||
Requires USAGE on the supplied object, or on system if object is omitted.
|
||||
``limit`` bounds yielded entries across all pages; None lists all and zero
|
||||
makes no requests. Missing matches produce an empty iterator.
|
||||
"""
|
||||
_limit(limit)
|
||||
request = {}
|
||||
if object is not None:
|
||||
request["object"] = _object(object)
|
||||
if subject is not None:
|
||||
request["subject"] = _subject(subject)
|
||||
|
||||
async def entries():
|
||||
async for value in self._pages(
|
||||
self._inner.list_acls, request, "entries", limit
|
||||
):
|
||||
yield _entry(value)
|
||||
|
||||
return entries()
|
||||
|
||||
async def add_acl(
|
||||
self,
|
||||
*,
|
||||
object: Union[str, Object],
|
||||
subject: Subject,
|
||||
privilege: Union[str, Privilege],
|
||||
) -> None:
|
||||
"""Grant one privilege; requires OWNERSHIP on the object.
|
||||
|
||||
A duplicate grant raises HTTP 409. Granting OWNERSHIP transfers ownership.
|
||||
Writes are not automatically retried after ambiguous failures.
|
||||
"""
|
||||
await self._inner.add_acl(
|
||||
json.dumps(
|
||||
{
|
||||
"object": _object(object),
|
||||
"subject": _subject(subject),
|
||||
"privilege": _privilege(privilege),
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
async def delete_acl(
|
||||
self,
|
||||
*,
|
||||
object: Union[str, Object],
|
||||
subject: Subject,
|
||||
privilege: Union[str, Privilege],
|
||||
) -> None:
|
||||
"""Revoke one privilege; requires OWNERSHIP on the object.
|
||||
|
||||
A missing grant raises HTTP 404. OWNERSHIP itself cannot be deleted.
|
||||
"""
|
||||
await self._inner.delete_acl(
|
||||
json.dumps(
|
||||
{
|
||||
"object": _object(object),
|
||||
"subject": _subject(subject),
|
||||
"privilege": _privilege(privilege),
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
def list_role_bindings(
|
||||
self,
|
||||
*,
|
||||
subject: Optional[Subject] = None,
|
||||
role: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> AsyncIterator[RoleBinding]:
|
||||
"""Iterate direct bindings, optionally filtered by principal/group and role.
|
||||
|
||||
Requires USAGE on system. Limit semantics match ``list_acls``. Principals'
|
||||
roles inherited through groups are available through ``get_principal``.
|
||||
"""
|
||||
_limit(limit)
|
||||
request = {}
|
||||
if subject is not None:
|
||||
request["subject"] = _subject(subject, "role_member")
|
||||
if role is not None:
|
||||
request["role"] = role
|
||||
|
||||
async def bindings():
|
||||
async for value in self._pages(
|
||||
self._inner.list_role_bindings, request, "bindings", limit
|
||||
):
|
||||
yield _binding(value)
|
||||
|
||||
return bindings()
|
||||
|
||||
async def add_role_binding(self, *, subject: Subject, role: str) -> None:
|
||||
"""Grant a role idempotently; requires OPERATE on system."""
|
||||
await self._inner.add_role_binding(
|
||||
json.dumps({"subject": _subject(subject, "role_member"), "role": role})
|
||||
)
|
||||
|
||||
async def delete_role_binding(self, *, subject: Subject, role: str) -> None:
|
||||
"""Idempotently revoke a role; requires OPERATE on system."""
|
||||
await self._inner.delete_role_binding(
|
||||
json.dumps({"subject": _subject(subject, "role_member"), "role": role})
|
||||
)
|
||||
|
||||
async def get_principal(self, subject: Subject) -> PrincipalInfo:
|
||||
"""Look up groups and effective roles; requires USAGE on system.
|
||||
|
||||
Select by ID, name, or API key. Missing/ambiguous identities raise HTTP
|
||||
404/409; unavailable or unconfigured resolution raises HTTP 503/501.
|
||||
"""
|
||||
value = json.loads(
|
||||
await self._inner.get_principal(
|
||||
json.dumps({"principal": _subject(subject, "principal")})
|
||||
)
|
||||
)
|
||||
return PrincipalInfo(
|
||||
value["principal_id"],
|
||||
value.get("principal_name"),
|
||||
[IdentityInfo(**group) for group in value["groups"]],
|
||||
value["roles"],
|
||||
)
|
||||
|
||||
async def get_group(self, subject: Subject) -> GroupInfo:
|
||||
"""Look up a group's members and roles by ID/name; requires USAGE on system.
|
||||
|
||||
Membership is read-only and comes from the identity provider. Lookup
|
||||
errors have the same HTTP semantics as ``get_principal``.
|
||||
"""
|
||||
value = json.loads(
|
||||
await self._inner.get_group(
|
||||
json.dumps({"group": _subject(subject, "group")})
|
||||
)
|
||||
)
|
||||
return GroupInfo(
|
||||
value["group_id"],
|
||||
value.get("group_name"),
|
||||
[IdentityInfo(**principal) for principal in value["principals"]],
|
||||
value["roles"],
|
||||
)
|
||||
|
||||
|
||||
async def _next(iterator: AsyncIterator[_T]) -> _T:
|
||||
return await iterator.__anext__()
|
||||
|
||||
|
||||
def _sync_iter(iterator: AsyncIterator[_T]) -> Iterator[_T]:
|
||||
try:
|
||||
while True:
|
||||
yield LOOP.run(_next(iterator))
|
||||
except StopAsyncIteration:
|
||||
return
|
||||
finally:
|
||||
close = getattr(iterator, "aclose", None)
|
||||
if close is not None:
|
||||
LOOP.run(close())
|
||||
|
||||
|
||||
class Authorization:
|
||||
"""Synchronous authorization returned by ``Catalog.authz``.
|
||||
|
||||
Shares the async API's permissions, errors, and lazy pagination behavior.
|
||||
Group membership comes from the identity provider and cannot be edited here.
|
||||
"""
|
||||
|
||||
def __init__(self, inner: AsyncAuthorization):
|
||||
self._inner = inner
|
||||
|
||||
def list_acls(
|
||||
self,
|
||||
*,
|
||||
object: Optional[Union[str, Object]] = None,
|
||||
subject: Optional[Subject] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> Iterator[AccessControlEntry]:
|
||||
"""Iterate grants.
|
||||
|
||||
See [list_acls][lancedb.authz.AsyncAuthorization.list_acls] for permissions.
|
||||
"""
|
||||
return _sync_iter(
|
||||
self._inner.list_acls(object=object, subject=subject, limit=limit)
|
||||
)
|
||||
|
||||
def add_acl(
|
||||
self,
|
||||
*,
|
||||
object: Union[str, Object],
|
||||
subject: Subject,
|
||||
privilege: Union[str, Privilege],
|
||||
) -> None:
|
||||
"""Grant a privilege; duplicates raise HTTP 409.
|
||||
|
||||
OWNERSHIP transfers ownership.
|
||||
"""
|
||||
LOOP.run(
|
||||
self._inner.add_acl(object=object, subject=subject, privilege=privilege)
|
||||
)
|
||||
|
||||
def delete_acl(
|
||||
self,
|
||||
*,
|
||||
object: Union[str, Object],
|
||||
subject: Subject,
|
||||
privilege: Union[str, Privilege],
|
||||
) -> None:
|
||||
"""Revoke a privilege. Missing grants raise HTTP 404.
|
||||
|
||||
OWNERSHIP cannot be deleted.
|
||||
"""
|
||||
LOOP.run(
|
||||
self._inner.delete_acl(object=object, subject=subject, privilege=privilege)
|
||||
)
|
||||
|
||||
def list_role_bindings(
|
||||
self,
|
||||
*,
|
||||
subject: Optional[Subject] = None,
|
||||
role: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> Iterator[RoleBinding]:
|
||||
"""Iterate direct role bindings, using the same limit semantics as list_acls."""
|
||||
return _sync_iter(
|
||||
self._inner.list_role_bindings(subject=subject, role=role, limit=limit)
|
||||
)
|
||||
|
||||
def add_role_binding(self, *, subject: Subject, role: str) -> None:
|
||||
"""Grant a role idempotently; requires OPERATE on system."""
|
||||
LOOP.run(self._inner.add_role_binding(subject=subject, role=role))
|
||||
|
||||
def delete_role_binding(self, *, subject: Subject, role: str) -> None:
|
||||
"""Idempotently revoke a role; requires OPERATE on system."""
|
||||
LOOP.run(self._inner.delete_role_binding(subject=subject, role=role))
|
||||
|
||||
def get_principal(self, subject: Subject) -> PrincipalInfo:
|
||||
"""Look up groups and effective roles; requires USAGE on system."""
|
||||
return LOOP.run(self._inner.get_principal(subject))
|
||||
|
||||
def get_group(self, subject: Subject) -> GroupInfo:
|
||||
"""Look up a group's members and roles; requires USAGE on system."""
|
||||
return LOOP.run(self._inner.get_group(subject))
|
||||
@@ -3,23 +3,78 @@
|
||||
|
||||
"""Remote catalogs manage databases through a server's root namespace."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from functools import cached_property
|
||||
from datetime import timedelta
|
||||
from typing import Any, Optional, Union
|
||||
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
|
||||
|
||||
|
||||
@dataclass
|
||||
class ListDatabasesResponse:
|
||||
"""A page of database names and an optional continuation token."""
|
||||
class AsyncDatabaseNames(AsyncIterator[str]):
|
||||
"""An async iterator of database names with pagination state.
|
||||
|
||||
databases: list[str]
|
||||
page_token: Optional[str] = None
|
||||
Returned by [list_databases][lancedb.catalog.AsyncCatalog.list_databases].
|
||||
"""
|
||||
|
||||
def __init__(self, inner: _lancedb.DatabaseNames):
|
||||
self._inner = inner
|
||||
|
||||
def __aiter__(self) -> "AsyncDatabaseNames":
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> str:
|
||||
return await self._inner.__anext__()
|
||||
|
||||
def num_page_results(self) -> int:
|
||||
"""Return the number of names available without another REST request."""
|
||||
return self._inner.num_page_results()
|
||||
|
||||
def page_token(self) -> Optional[str]:
|
||||
"""Return the token for the next REST request.
|
||||
|
||||
Before the first request, this is the supplied starting token (``None``
|
||||
starts at the beginning). After the final page is fetched, it is ``None``,
|
||||
even if names remain cached. Drain the cache before saving a token to
|
||||
avoid skipping those names when resuming. A failed request terminates
|
||||
iteration and leaves its token available for resuming a new iterator.
|
||||
"""
|
||||
return self._inner.page_token()
|
||||
|
||||
|
||||
class DatabaseNames(Iterator[str]):
|
||||
"""A synchronous iterator of database names with pagination state.
|
||||
|
||||
Returned by [Catalog.list_databases][lancedb.catalog.Catalog.list_databases].
|
||||
"""
|
||||
|
||||
def __init__(self, inner: AsyncDatabaseNames):
|
||||
self._inner = inner
|
||||
|
||||
def __iter__(self) -> "DatabaseNames":
|
||||
return self
|
||||
|
||||
def __next__(self) -> str:
|
||||
try:
|
||||
return LOOP.run(self._inner.__anext__())
|
||||
except StopAsyncIteration:
|
||||
raise StopIteration from None
|
||||
|
||||
def num_page_results(self) -> int:
|
||||
"""Return the number of names available without another REST request."""
|
||||
return self._inner.num_page_results()
|
||||
|
||||
def page_token(self) -> Optional[str]:
|
||||
"""Return the token for the next REST request.
|
||||
|
||||
See [page_token][lancedb.catalog.AsyncDatabaseNames.page_token]
|
||||
for initial, cached, final-page, and error behavior.
|
||||
"""
|
||||
return self._inner.page_token()
|
||||
|
||||
|
||||
class AsyncCatalog:
|
||||
@@ -33,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."""
|
||||
@@ -50,14 +110,25 @@ class AsyncCatalog:
|
||||
"""Connect to an existing database by its logical name."""
|
||||
return AsyncConnection(await self._inner.connect_database(name))
|
||||
|
||||
async def list_databases(
|
||||
self, *, limit: Optional[int] = None, page_token: Optional[str] = None
|
||||
) -> ListDatabasesResponse:
|
||||
"""List a page of databases. Pass the returned token for the next page."""
|
||||
names, token = await self._inner.list_databases(
|
||||
limit=limit, page_token=page_token
|
||||
def list_databases(
|
||||
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||
) -> AsyncDatabaseNames:
|
||||
"""Iterate lazily over all database names using ``async for``.
|
||||
|
||||
``page_token`` resumes from a saved token; ``None`` starts at the beginning.
|
||||
``page_limit`` sets the maximum names per REST response (not the total);
|
||||
``None`` uses the server default. Request errors are raised during iteration.
|
||||
|
||||
Examples
|
||||
--------
|
||||
```python
|
||||
async for name in catalog.list_databases():
|
||||
print(name)
|
||||
```
|
||||
"""
|
||||
return AsyncDatabaseNames(
|
||||
self._inner.list_databases(page_token=page_token, page_limit=page_limit)
|
||||
)
|
||||
return ListDatabasesResponse(names, token)
|
||||
|
||||
async def drop_database(self, name: str, *, ignore_missing: bool = False) -> None:
|
||||
"""Drop an empty database. A nonempty database is an error."""
|
||||
@@ -73,7 +144,8 @@ class Catalog:
|
||||
```python
|
||||
catalog = lancedb.connect_catalog("https://my-server.example", api_key="secret")
|
||||
db = catalog.create_database("analytics", exist_ok=True)
|
||||
page = catalog.list_databases(limit=20)
|
||||
for name in catalog.list_databases():
|
||||
print(name)
|
||||
```
|
||||
"""
|
||||
|
||||
@@ -92,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."""
|
||||
@@ -118,10 +195,17 @@ class Catalog:
|
||||
)
|
||||
|
||||
def list_databases(
|
||||
self, *, limit: Optional[int] = None, page_token: Optional[str] = None
|
||||
) -> ListDatabasesResponse:
|
||||
"""List a page of databases. Pass the returned token for the next page."""
|
||||
return LOOP.run(self._inner.list_databases(limit=limit, page_token=page_token))
|
||||
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||
) -> DatabaseNames:
|
||||
"""Iterate lazily over database names, fetching more results as needed.
|
||||
|
||||
``page_token`` resumes from a saved token; ``None`` starts at the beginning.
|
||||
``page_limit`` sets the maximum names per REST response (not the total);
|
||||
``None`` uses the server default. Request errors are raised during iteration.
|
||||
"""
|
||||
return DatabaseNames(
|
||||
self._inner.list_databases(page_token=page_token, page_limit=page_limit)
|
||||
)
|
||||
|
||||
def drop_database(self, name: str, *, ignore_missing: bool = False) -> None:
|
||||
"""Drop an empty database. A nonempty database is an error."""
|
||||
|
||||
+445
-121
@@ -31,6 +31,7 @@ if sys.version_info >= (3, 12):
|
||||
else:
|
||||
from overrides import EnforceOverrides, override # type: ignore
|
||||
|
||||
from lancedb.listing import AsyncListing, Listing
|
||||
from lancedb.embeddings.registry import EmbeddingFunctionRegistry
|
||||
|
||||
from lancedb.common import data_to_reader, sanitize_uri, validate_schema
|
||||
@@ -52,12 +53,11 @@ from .functions import FunctionVersion, UdfDefinition
|
||||
from .job import AsyncJob, Job, _typed_job
|
||||
from .sql import AsyncQuery as AsyncSqlQuery
|
||||
from .sql import Query as SqlQuery
|
||||
from .sql import QueryDescription
|
||||
from .sql import QueryDescription, QueryParameters, to_parameter_batch
|
||||
from .materialized_view import (
|
||||
AsyncMaterializedView,
|
||||
MaterializedView,
|
||||
MaterializedViewSource,
|
||||
VectorDedupSource,
|
||||
SelectArg,
|
||||
normalize_select,
|
||||
)
|
||||
@@ -122,7 +122,48 @@ def _view_description(
|
||||
|
||||
|
||||
class DBConnection(EnforceOverrides):
|
||||
"""An active LanceDB connection interface."""
|
||||
"""An active LanceDB connection interface.
|
||||
|
||||
Use [close][lancedb.db.DBConnection.close] to release the connection's
|
||||
underlying resources, or use the connection as a context manager to close it
|
||||
automatically when leaving the block, including when an exception is raised.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import lancedb
|
||||
>>> with lancedb.connect("memory://") as db:
|
||||
... assert db.is_open()
|
||||
>>> db.is_open()
|
||||
False
|
||||
"""
|
||||
|
||||
def __enter__(self) -> DBConnection:
|
||||
return self
|
||||
|
||||
def __exit__(self, *_) -> None:
|
||||
self.close()
|
||||
|
||||
@abstractmethod
|
||||
def is_open(self) -> bool:
|
||||
"""Return True if the connection is open."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def close(self) -> None:
|
||||
"""Close the connection, releasing any underlying resources.
|
||||
|
||||
It is safe to call this method multiple times. Database operations on a
|
||||
closed connection raise ``RuntimeError: Connection is closed``.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import lancedb
|
||||
>>> db = lancedb.connect("memory://")
|
||||
>>> db.close()
|
||||
>>> db.is_open()
|
||||
False
|
||||
"""
|
||||
pass
|
||||
|
||||
def list_namespaces(
|
||||
self,
|
||||
@@ -276,7 +317,7 @@ class DBConnection(EnforceOverrides):
|
||||
page_token: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> ListTablesResponse:
|
||||
"""List all tables in this database with pagination support.
|
||||
"""List a page of tables in this database.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
@@ -287,7 +328,8 @@ class DBConnection(EnforceOverrides):
|
||||
Token for pagination. Use the token from a previous response
|
||||
to get the next page of results.
|
||||
limit: int, optional
|
||||
The maximum number of results to return.
|
||||
The maximum number of results to return, default 100.
|
||||
Zero returns an empty page without a continuation token.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -302,7 +344,7 @@ class DBConnection(EnforceOverrides):
|
||||
def table_names(
|
||||
self,
|
||||
page_token: Optional[str] = None,
|
||||
limit: int = 10,
|
||||
limit: Optional[int] = None,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> Iterable[str]:
|
||||
@@ -316,8 +358,9 @@ class DBConnection(EnforceOverrides):
|
||||
page_token: str, optional
|
||||
The token to use for pagination. If not present, start from the beginning.
|
||||
Typically, this token is last table name from the previous page.
|
||||
limit: int, default 10
|
||||
The size of the page to return.
|
||||
limit: int, optional
|
||||
The maximum number of tables to return. None returns all tables;
|
||||
zero returns an empty list.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -368,9 +411,10 @@ class DBConnection(EnforceOverrides):
|
||||
- [LanceModel][lancedb.pydantic.LanceModel]
|
||||
mode: str; default "create"
|
||||
The mode to use when creating the table.
|
||||
Can be either "create" or "overwrite".
|
||||
Can be "create", "overwrite", or "exist_ok".
|
||||
By default, if the table already exists, an exception is raised.
|
||||
If you want to overwrite the table, use mode="overwrite".
|
||||
To open an existing table without adding data, use mode="exist_ok".
|
||||
exist_ok: bool, default False
|
||||
If a table by the same name already exists, then raise an exception
|
||||
if exist_ok=False. If exist_ok=True, then open the existing table;
|
||||
@@ -571,21 +615,12 @@ class DBConnection(EnforceOverrides):
|
||||
table. The view is a normal table: it can be queried, indexed and
|
||||
searched, and it appears in ``table_names``.
|
||||
|
||||
Ordinary source tables must have stable row ids (create them with the
|
||||
``new_table_enable_stable_row_ids`` storage option): they keep the
|
||||
view's provenance valid across source compactions, and cannot be
|
||||
enabled after a table exists. Native dedup sources also support physical
|
||||
row IDs interpreted against their pinned snapshot.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name: str
|
||||
The name of the view.
|
||||
source: str or VectorDedupSource
|
||||
The name of the source table, or an indexed dedup source declared
|
||||
with [vector_dedup][lancedb.vector_dedup]. Dedup sources capture a
|
||||
fixed snapshot and preserve the source table. They accept no
|
||||
additional select, where or limit options.
|
||||
source: str
|
||||
The name of the source table.
|
||||
select: list or dict, optional
|
||||
The view's columns: column names, ``(alias, SQL expression)``
|
||||
pairs, or a dict of the same. Omitting it selects every source
|
||||
@@ -636,11 +671,16 @@ class DBConnection(EnforceOverrides):
|
||||
"materialized views are not supported on this connection type"
|
||||
)
|
||||
|
||||
def list_materialized_views(self) -> List[str]:
|
||||
def list_materialized_views(
|
||||
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||
) -> Listing[str]:
|
||||
"""The names of the materialized views in this database.
|
||||
|
||||
Found by reading every table's schema, so this costs an open per
|
||||
table.
|
||||
Local listings inspect table schemas one page at a time.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"materialized views are not supported on this connection type"
|
||||
@@ -784,6 +824,7 @@ class DBConnection(EnforceOverrides):
|
||||
definition: UdfDefinition,
|
||||
*,
|
||||
secrets: Optional[Sequence[EnvVarSecret]] = None,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> FunctionVersion:
|
||||
"""Build and register a scalar Python UDF, then return its version.
|
||||
|
||||
@@ -800,6 +841,9 @@ class DBConnection(EnforceOverrides):
|
||||
Function needs, each naming a Secret and the environment variable
|
||||
its value arrives in. The Function's source is unchanged by this;
|
||||
it reads the variable the way it already did.
|
||||
namespace_path : list of str, optional
|
||||
The namespace holding the Function. None or an empty list is the
|
||||
root namespace.
|
||||
|
||||
Examples
|
||||
--------
|
||||
@@ -815,13 +859,16 @@ class DBConnection(EnforceOverrides):
|
||||
)
|
||||
```
|
||||
"""
|
||||
return self.create_function_async(definition, secrets=secrets).wait()
|
||||
return self.create_function_async(
|
||||
definition, secrets=secrets, namespace_path=namespace_path
|
||||
).wait()
|
||||
|
||||
def create_function_async(
|
||||
self,
|
||||
definition: UdfDefinition,
|
||||
*,
|
||||
secrets: Optional[Sequence[EnvVarSecret]] = None,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> Job[FunctionVersion]:
|
||||
"""Submit a scalar Python UDF for building and registration.
|
||||
|
||||
@@ -833,17 +880,30 @@ class DBConnection(EnforceOverrides):
|
||||
"Function catalog operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def get_function(self, name: str, *, version: str) -> FunctionVersion:
|
||||
def get_function(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
version: str,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> FunctionVersion:
|
||||
"""Open one exact immutable Function version from the remote catalog."""
|
||||
raise NotImplementedError(
|
||||
"Function catalog operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def list_functions(self) -> List[FunctionVersion]:
|
||||
"""List every published immutable Function version.
|
||||
def list_functions(
|
||||
self,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
page_limit: Optional[int] = None,
|
||||
) -> Listing[FunctionVersion]:
|
||||
"""List every published immutable Function version in a namespace.
|
||||
|
||||
Results are ordered by Function name then version. Local connections
|
||||
raise ``NotImplementedError``.
|
||||
Functions in child namespaces are not included. Results are ordered by
|
||||
Function name then version. Local connections raise
|
||||
``NotImplementedError``.
|
||||
|
||||
Examples
|
||||
--------
|
||||
@@ -852,12 +912,22 @@ class DBConnection(EnforceOverrides):
|
||||
```python
|
||||
[(function.name, function.version) for function in db.list_functions()]
|
||||
```
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Function catalog operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def drop_function(self, name: str, *, version: str) -> bool:
|
||||
def drop_function(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
version: str,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> bool:
|
||||
"""Drop a Function name and the object it was bound to.
|
||||
|
||||
The requested version must exist in the currently named object. Returns
|
||||
@@ -870,7 +940,13 @@ class DBConnection(EnforceOverrides):
|
||||
"Function catalog operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def drop_function_async(self, name: str, *, version: str) -> Tuple[bool, Job]:
|
||||
def drop_function_async(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
version: str,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> Tuple[bool, Job]:
|
||||
"""Drop a Function name and return its cleanup job.
|
||||
|
||||
The name is unbound before this returns; the object's content may still
|
||||
@@ -911,11 +987,21 @@ class DBConnection(EnforceOverrides):
|
||||
"Secret operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def list_secrets(self, *, namespace_path: Optional[List[str]] = None) -> List[str]:
|
||||
def list_secrets(
|
||||
self,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
page_limit: Optional[int] = None,
|
||||
) -> Listing[str]:
|
||||
"""The names of every Secret in this database.
|
||||
|
||||
Names only. No method returns a stored credential, by construction
|
||||
rather than by policy. Local connections raise ``NotImplementedError``.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Secret operations are not supported for this connection type"
|
||||
@@ -1013,12 +1099,22 @@ class DBConnection(EnforceOverrides):
|
||||
"View operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def list_views(self, *, namespace_path: Optional[List[str]] = None) -> List[str]:
|
||||
def list_views(
|
||||
self,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
page_limit: Optional[int] = None,
|
||||
) -> Listing[str]:
|
||||
"""The names of the views in one namespace.
|
||||
|
||||
Names only; a definition comes from
|
||||
[describe_view][lancedb.db.DBConnection.describe_view]. Local
|
||||
connections raise ``NotImplementedError``.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"View operations are not supported for this connection type"
|
||||
@@ -1037,8 +1133,15 @@ class DBConnection(EnforceOverrides):
|
||||
"""
|
||||
raise NotImplementedError("open_job is not supported for this connection type")
|
||||
|
||||
def list_jobs(self) -> List[JobInfo]:
|
||||
"""List server-side jobs across the database's tables."""
|
||||
def list_jobs(
|
||||
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||
) -> Listing[JobInfo]:
|
||||
"""List server-side jobs across the database's tables.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
raise NotImplementedError("list_jobs is not supported for this connection type")
|
||||
|
||||
def cancel_job(self, job_id: str) -> bool:
|
||||
@@ -1077,6 +1180,7 @@ class DBConnection(EnforceOverrides):
|
||||
query: str,
|
||||
*,
|
||||
default_namespace_path: Optional[List[str]] = None,
|
||||
parameters: Optional[QueryParameters] = None,
|
||||
) -> pa.RecordBatchReader:
|
||||
"""Execute SQL and return a blocking Arrow reader.
|
||||
|
||||
@@ -1087,6 +1191,7 @@ class DBConnection(EnforceOverrides):
|
||||
return self.execute_query_async(
|
||||
query,
|
||||
default_namespace_path=default_namespace_path,
|
||||
parameters=parameters,
|
||||
).reader()
|
||||
|
||||
def execute_query_async(
|
||||
@@ -1094,9 +1199,32 @@ class DBConnection(EnforceOverrides):
|
||||
query: str,
|
||||
*,
|
||||
default_namespace_path: Optional[List[str]] = None,
|
||||
parameters: Optional[QueryParameters] = None,
|
||||
) -> SqlQuery:
|
||||
"""Start executing SQL and return its query handle.
|
||||
|
||||
``parameters`` binds values to the statement's ``$1`` / ``$name``
|
||||
placeholders. They travel as Arrow rather than as SQL text, so a
|
||||
float keeps its exact value and type and a vector stays a compact
|
||||
fixed-size list. A list binds by position, a dict binds by name, and
|
||||
a one-row ``pyarrow.RecordBatch`` or ``pyarrow.Table`` binds both
|
||||
ways. A 1-D numpy or pyarrow array becomes a fixed-size list of its
|
||||
element type, the shape of a vector column:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
db.execute_query(
|
||||
"SELECT id FROM docs ORDER BY distance(vector, $1) LIMIT $2",
|
||||
parameters=[np.array([0.1, 0.2, 0.3], dtype=np.float32), 10],
|
||||
)
|
||||
|
||||
Parameters are supported in queries, not in DDL or DML. A
|
||||
parameterized query runs on the call that returns its rows:
|
||||
cancelling it, or dropping its handle before reading, stops it on the
|
||||
server. Its status follows that stream as well: ``describe()``
|
||||
reports ``running`` until the reader has received every row, so open
|
||||
``reader()`` directly rather than waiting for ``finished``.
|
||||
|
||||
Local connections do not support SQL.
|
||||
"""
|
||||
raise NotImplementedError("SQL is not supported for this connection type")
|
||||
@@ -1247,6 +1375,14 @@ class LanceDBConnection(DBConnection):
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}(uri={self._conn.uri!r})"
|
||||
|
||||
@override
|
||||
def is_open(self) -> bool:
|
||||
return self._conn.is_open()
|
||||
|
||||
@override
|
||||
def close(self) -> None:
|
||||
self._conn.close()
|
||||
|
||||
@override
|
||||
def serialize(self) -> str:
|
||||
import json
|
||||
@@ -1356,7 +1492,7 @@ class LanceDBConnection(DBConnection):
|
||||
page_token: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> ListTablesResponse:
|
||||
"""List all tables in this database with pagination support.
|
||||
"""List a page of tables in this database.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
@@ -1367,7 +1503,8 @@ class LanceDBConnection(DBConnection):
|
||||
Token for pagination. Use the token from a previous response
|
||||
to get the next page of results.
|
||||
limit: int, optional
|
||||
The maximum number of results to return.
|
||||
The maximum number of results to return, default 100.
|
||||
Zero returns an empty page without a continuation token.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -1386,7 +1523,7 @@ class LanceDBConnection(DBConnection):
|
||||
def table_names(
|
||||
self,
|
||||
page_token: Optional[str] = None,
|
||||
limit: int = 10,
|
||||
limit: Optional[int] = None,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> Iterable[str]:
|
||||
@@ -1401,8 +1538,9 @@ class LanceDBConnection(DBConnection):
|
||||
The namespace to list tables in.
|
||||
page_token: str, optional
|
||||
The token to use for pagination.
|
||||
limit: int, default 10
|
||||
The maximum number of tables to return.
|
||||
limit: int, optional
|
||||
The maximum number of tables to return. None returns all tables;
|
||||
zero returns an empty list.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -1469,8 +1607,6 @@ class LanceDBConnection(DBConnection):
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
if mode.lower() not in ["create", "overwrite"]:
|
||||
raise ValueError("mode must be either 'create' or 'overwrite'")
|
||||
validate_table_name(name)
|
||||
|
||||
tbl = LanceTable.create(
|
||||
@@ -1578,10 +1714,7 @@ class LanceDBConnection(DBConnection):
|
||||
Examples
|
||||
--------
|
||||
>>> import lancedb
|
||||
>>> db = lancedb.connect(
|
||||
... "./.lancedb",
|
||||
... storage_options={"new_table_enable_stable_row_ids": "true"},
|
||||
... )
|
||||
>>> db = lancedb.connect("./.lancedb")
|
||||
>>> data = [{"name": "ada", "age": 36}, {"name": "kid", "age": 7}]
|
||||
>>> table = db.create_table("people", data)
|
||||
>>> view = db.create_materialized_view(
|
||||
@@ -1636,9 +1769,20 @@ class LanceDBConnection(DBConnection):
|
||||
return view
|
||||
|
||||
@override
|
||||
def list_materialized_views(self) -> List[str]:
|
||||
"""The names of the materialized views in this database."""
|
||||
return LOOP.run(self._conn.list_materialized_views())
|
||||
def list_materialized_views(
|
||||
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||
) -> Listing[str]:
|
||||
"""The names of the materialized views in this database.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
return Listing(
|
||||
self._conn.list_materialized_views(
|
||||
page_token=page_token, page_limit=page_limit
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def drop_materialized_view(
|
||||
@@ -1808,25 +1952,78 @@ class LanceDBConnection(DBConnection):
|
||||
definition: UdfDefinition,
|
||||
*,
|
||||
secrets: Optional[Sequence[EnvVarSecret]] = None,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> Job[FunctionVersion]:
|
||||
job = LOOP.run(self._conn.create_function_async(definition, secrets=secrets))
|
||||
job = LOOP.run(
|
||||
self._conn.create_function_async(
|
||||
definition, secrets=secrets, namespace_path=namespace_path
|
||||
)
|
||||
)
|
||||
return Job(job)
|
||||
|
||||
@override
|
||||
def get_function(self, name: str, *, version: str) -> FunctionVersion:
|
||||
return LOOP.run(self._conn.get_function(name, version=version))
|
||||
def get_function(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
version: str,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> FunctionVersion:
|
||||
return LOOP.run(
|
||||
self._conn.get_function(
|
||||
name, version=version, namespace_path=namespace_path
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def list_functions(self) -> List[FunctionVersion]:
|
||||
return LOOP.run(self._conn.list_functions())
|
||||
def list_functions(
|
||||
self,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
page_limit: Optional[int] = None,
|
||||
) -> Listing[FunctionVersion]:
|
||||
"""List functions.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
return Listing(
|
||||
self._conn.list_functions(
|
||||
namespace_path=namespace_path,
|
||||
page_token=page_token,
|
||||
page_limit=page_limit,
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def drop_function(self, name: str, *, version: str) -> bool:
|
||||
return LOOP.run(self._conn.drop_function(name, version=version))
|
||||
def drop_function(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
version: str,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> bool:
|
||||
return LOOP.run(
|
||||
self._conn.drop_function(
|
||||
name, version=version, namespace_path=namespace_path
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def drop_function_async(self, name: str, *, version: str) -> Tuple[bool, Job]:
|
||||
dropped, job = LOOP.run(self._conn.drop_function_async(name, version=version))
|
||||
def drop_function_async(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
version: str,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> Tuple[bool, Job]:
|
||||
dropped, job = LOOP.run(
|
||||
self._conn.drop_function_async(
|
||||
name, version=version, namespace_path=namespace_path
|
||||
)
|
||||
)
|
||||
return dropped, Job(job)
|
||||
|
||||
@override
|
||||
@@ -1842,8 +2039,26 @@ class LanceDBConnection(DBConnection):
|
||||
LOOP.run(self._conn.alter_secret(name, value, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def list_secrets(self, *, namespace_path: Optional[List[str]] = None) -> List[str]:
|
||||
return LOOP.run(self._conn.list_secrets(namespace_path=namespace_path))
|
||||
def list_secrets(
|
||||
self,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
page_limit: Optional[int] = None,
|
||||
) -> Listing[str]:
|
||||
"""List secrets.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
return Listing(
|
||||
self._conn.list_secrets(
|
||||
namespace_path=namespace_path,
|
||||
page_token=page_token,
|
||||
page_limit=page_limit,
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def drop_secret(
|
||||
@@ -1886,13 +2101,40 @@ class LanceDBConnection(DBConnection):
|
||||
)
|
||||
|
||||
@override
|
||||
def list_views(self, *, namespace_path: Optional[List[str]] = None) -> List[str]:
|
||||
return LOOP.run(self._conn.list_views(namespace_path=namespace_path))
|
||||
def list_views(
|
||||
self,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
page_limit: Optional[int] = None,
|
||||
) -> Listing[str]:
|
||||
"""List views.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
return Listing(
|
||||
self._conn.list_views(
|
||||
namespace_path=namespace_path,
|
||||
page_token=page_token,
|
||||
page_limit=page_limit,
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def list_jobs(self) -> List[JobInfo]:
|
||||
"""List server-side jobs across the database's tables."""
|
||||
return LOOP.run(self._conn.list_jobs())
|
||||
def list_jobs(
|
||||
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||
) -> Listing[JobInfo]:
|
||||
"""List server-side jobs across the database's tables.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
return Listing(
|
||||
self._conn.list_jobs(page_token=page_token, page_limit=page_limit)
|
||||
)
|
||||
|
||||
@override
|
||||
def cancel_job(self, job_id: str) -> bool:
|
||||
@@ -2128,7 +2370,7 @@ class AsyncConnection(object):
|
||||
page_token: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> ListTablesResponse:
|
||||
"""List all tables in this database with pagination support.
|
||||
"""List a page of tables in this database.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
@@ -2139,7 +2381,8 @@ class AsyncConnection(object):
|
||||
Token for pagination. Use the token from a previous response
|
||||
to get the next page of results.
|
||||
limit: int, optional
|
||||
The maximum number of results to return.
|
||||
The maximum number of results to return, default 100.
|
||||
Zero returns an empty page without a continuation token.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -2176,8 +2419,9 @@ class AsyncConnection(object):
|
||||
|
||||
This can be combined with limit to implement pagination by setting this to
|
||||
the last table name from the previous page.
|
||||
limit: int, default 10
|
||||
The number of results to return.
|
||||
limit: int, optional
|
||||
The maximum number of tables to return. None returns all tables;
|
||||
zero returns an empty list.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -2236,11 +2480,12 @@ class AsyncConnection(object):
|
||||
- pyarrow.Schema
|
||||
|
||||
- [LanceModel][lancedb.pydantic.LanceModel]
|
||||
mode: Literal["create", "overwrite"]; default "create"
|
||||
mode: Literal["create", "overwrite", "exist_ok"]; default "create"
|
||||
The mode to use when creating the table.
|
||||
Can be either "create" or "overwrite".
|
||||
Can be "create", "overwrite", or "exist_ok".
|
||||
By default, if the table already exists, an exception is raised.
|
||||
If you want to overwrite the table, use mode="overwrite".
|
||||
To open an existing table without adding data, use mode="exist_ok".
|
||||
exist_ok: bool, default False
|
||||
If a table by the same name already exists, then raise an exception
|
||||
if exist_ok=False. If exist_ok=True, then open the existing table;
|
||||
@@ -2511,19 +2756,7 @@ class AsyncConnection(object):
|
||||
where: Optional[str],
|
||||
limit: Optional[int],
|
||||
) -> Tuple[str, Dict[str, str]]:
|
||||
if not isinstance(source, VectorDedupSource):
|
||||
return source, {}
|
||||
if select is not None or where is not None or limit is not None:
|
||||
raise ValueError(
|
||||
"vector_dedup cannot be combined with select, where or limit"
|
||||
)
|
||||
version = source.dataset_version
|
||||
if version is None:
|
||||
table = await self.open_table(source.source)
|
||||
version = await table.version()
|
||||
return source.source, {
|
||||
"vector_source_json": source._native_source_json(version)
|
||||
}
|
||||
return source, {}
|
||||
|
||||
async def create_materialized_view(
|
||||
self,
|
||||
@@ -2591,13 +2824,22 @@ class AsyncConnection(object):
|
||||
await view.definition()
|
||||
return view
|
||||
|
||||
async def list_materialized_views(self) -> List[str]:
|
||||
def list_materialized_views(
|
||||
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||
) -> AsyncListing[str]:
|
||||
"""The names of the materialized views in this database.
|
||||
|
||||
Found by reading every table's schema, so this costs an open per
|
||||
table.
|
||||
Local listings inspect table schemas one page at a time.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
return await self._inner.list_materialized_views()
|
||||
return AsyncListing(
|
||||
self._inner.list_materialized_views(
|
||||
page_token=page_token, page_limit=page_limit
|
||||
)
|
||||
)
|
||||
|
||||
async def drop_materialized_view(
|
||||
self,
|
||||
@@ -2787,6 +3029,7 @@ class AsyncConnection(object):
|
||||
definition: UdfDefinition,
|
||||
*,
|
||||
secrets: Optional[Sequence[EnvVarSecret]] = None,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> AsyncJob[FunctionVersion]:
|
||||
"""Submit a scalar Python UDF for building and registration.
|
||||
|
||||
@@ -2794,36 +3037,72 @@ class AsyncConnection(object):
|
||||
artifact. Waiting on the job returns the immutable Function version.
|
||||
``secrets`` is a sequence of
|
||||
[EnvVarSecret][lancedb.secrets.EnvVarSecret], each naming a Secret and
|
||||
the environment variable its value arrives in. Local connections raise
|
||||
``NotImplementedError``.
|
||||
the environment variable its value arrives in. ``namespace_path`` is
|
||||
the namespace the Function is registered in; None or an empty list is
|
||||
the root. Local connections raise ``NotImplementedError``.
|
||||
"""
|
||||
if not isinstance(definition, UdfDefinition):
|
||||
raise TypeError("create_function_async requires a @udf definition")
|
||||
request = definition.bind_secrets(secrets)
|
||||
inner = await self._inner.create_function_async(request.to_canonical_json())
|
||||
inner = await self._inner.create_function_async(
|
||||
request.to_canonical_json(), namespace_path
|
||||
)
|
||||
return _typed_job(inner, FunctionVersion.from_json)
|
||||
|
||||
async def get_function(self, name: str, *, version: str) -> FunctionVersion:
|
||||
async def get_function(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
version: str,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> FunctionVersion:
|
||||
"""Open one exact immutable Function version from the remote catalog."""
|
||||
return FunctionVersion.from_json(await self._inner.get_function(name, version))
|
||||
return FunctionVersion.from_json(
|
||||
await self._inner.get_function(name, version, namespace_path)
|
||||
)
|
||||
|
||||
async def list_functions(self) -> List[FunctionVersion]:
|
||||
"""List every published immutable Function version.
|
||||
def list_functions(
|
||||
self,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
page_limit: Optional[int] = None,
|
||||
) -> AsyncListing[FunctionVersion]:
|
||||
"""List every published immutable Function version in a namespace.
|
||||
|
||||
Results are ordered by Function name then version. Local connections
|
||||
raise ``NotImplementedError``.
|
||||
Functions in child namespaces are not included. Results are ordered by
|
||||
Function name then version. Local connections raise
|
||||
``NotImplementedError``.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
return [
|
||||
FunctionVersion.from_json(value)
|
||||
for value in await self._inner.list_functions()
|
||||
]
|
||||
return AsyncListing(
|
||||
self._inner.list_functions(
|
||||
namespace_path=namespace_path,
|
||||
page_token=page_token,
|
||||
page_limit=page_limit,
|
||||
),
|
||||
FunctionVersion.from_json,
|
||||
)
|
||||
|
||||
async def drop_function(self, name: str, *, version: str) -> bool:
|
||||
async def drop_function(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
version: str,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> bool:
|
||||
"""Drop a Function name and the object it was bound to."""
|
||||
return await self._inner.drop_function(name, version)
|
||||
return await self._inner.drop_function(name, version, namespace_path)
|
||||
|
||||
async def drop_function_async(
|
||||
self, name: str, *, version: str
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
version: str,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> Tuple[bool, AsyncJob]:
|
||||
"""Drop a Function name and return its cleanup job.
|
||||
|
||||
@@ -2831,7 +3110,9 @@ class AsyncConnection(object):
|
||||
be being deleted. Await :meth:`AsyncJob.wait` to wait for that to
|
||||
finish.
|
||||
"""
|
||||
dropped, job = await self._inner.drop_function_async(name, version)
|
||||
dropped, job = await self._inner.drop_function_async(
|
||||
name, version, namespace_path
|
||||
)
|
||||
return dropped, AsyncJob(job)
|
||||
|
||||
async def create_secret(
|
||||
@@ -2862,12 +3143,25 @@ class AsyncConnection(object):
|
||||
list(validate_namespace_path(namespace_path)),
|
||||
)
|
||||
|
||||
async def list_secrets(
|
||||
self, *, namespace_path: Optional[List[str]] = None
|
||||
) -> List[str]:
|
||||
"""The names of every Secret in this database. Names only."""
|
||||
return await self._inner.list_secrets(
|
||||
list(validate_namespace_path(namespace_path))
|
||||
def list_secrets(
|
||||
self,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
page_limit: Optional[int] = None,
|
||||
) -> AsyncListing[str]:
|
||||
"""The names of every Secret in this database. Names only.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
return AsyncListing(
|
||||
self._inner.list_secrets(
|
||||
namespace_path=namespace_path,
|
||||
page_token=page_token,
|
||||
page_limit=page_limit,
|
||||
)
|
||||
)
|
||||
|
||||
async def drop_secret(
|
||||
@@ -2939,15 +3233,39 @@ class AsyncConnection(object):
|
||||
await self._inner.drop_view_async(name, namespace_path=namespace_path)
|
||||
)
|
||||
|
||||
async def list_views(
|
||||
self, *, namespace_path: Optional[List[str]] = None
|
||||
) -> List[str]:
|
||||
"""The names of the views in one namespace."""
|
||||
return await self._inner.list_views(list(namespace_path or []))
|
||||
def list_views(
|
||||
self,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
page_limit: Optional[int] = None,
|
||||
) -> AsyncListing[str]:
|
||||
"""The names of the views in one namespace.
|
||||
|
||||
async def list_jobs(self) -> List[JobInfo]:
|
||||
"""List server-side jobs across the database's tables."""
|
||||
return await self._inner.list_jobs()
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
return AsyncListing(
|
||||
self._inner.list_views(
|
||||
namespace_path=namespace_path,
|
||||
page_token=page_token,
|
||||
page_limit=page_limit,
|
||||
)
|
||||
)
|
||||
|
||||
def list_jobs(
|
||||
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||
) -> AsyncListing[JobInfo]:
|
||||
"""List server-side jobs across the database's tables.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
return AsyncListing(
|
||||
self._inner.list_jobs(page_token=page_token, page_limit=page_limit)
|
||||
)
|
||||
|
||||
async def cancel_job(self, job_id: str) -> bool:
|
||||
"""Request cancellation of a server-side job by id.
|
||||
@@ -2980,6 +3298,7 @@ class AsyncConnection(object):
|
||||
query: str,
|
||||
*,
|
||||
default_namespace_path: Optional[List[str]] = None,
|
||||
parameters: Optional[QueryParameters] = None,
|
||||
) -> AsyncRecordBatchReader:
|
||||
"""Execute SQL and return an asynchronous Arrow reader.
|
||||
|
||||
@@ -2990,6 +3309,7 @@ class AsyncConnection(object):
|
||||
submitted = await self.execute_query_async(
|
||||
query,
|
||||
default_namespace_path=default_namespace_path,
|
||||
parameters=parameters,
|
||||
)
|
||||
return await submitted.reader()
|
||||
|
||||
@@ -2998,17 +3318,21 @@ class AsyncConnection(object):
|
||||
query: str,
|
||||
*,
|
||||
default_namespace_path: Optional[List[str]] = None,
|
||||
parameters: Optional[QueryParameters] = None,
|
||||
) -> AsyncSqlQuery:
|
||||
"""Start executing SQL and return its query handle.
|
||||
|
||||
The database from ``connect_async`` is used for unqualified database
|
||||
references. The namespace defaults to ``["public"]``. Local
|
||||
connections raise ``NotImplementedError``.
|
||||
references. The namespace defaults to ``["public"]``. ``parameters``
|
||||
binds values to ``$1`` / ``$name`` placeholders; see
|
||||
[DBConnection.execute_query_async][lancedb.db.DBConnection.execute_query_async].
|
||||
Local connections raise ``NotImplementedError``.
|
||||
"""
|
||||
return AsyncSqlQuery(
|
||||
await self._inner.execute_query_async(
|
||||
query,
|
||||
default_namespace_path=default_namespace_path,
|
||||
parameters=to_parameter_batch(parameters),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -182,7 +182,11 @@ class EmbeddingFunction(BaseModel, ABC):
|
||||
def VectorField(self, **kwargs):
|
||||
"""
|
||||
Creates a pydantic Field that can automatically annotate
|
||||
the target vector column for this embedding function
|
||||
the target vector column for this embedding function.
|
||||
|
||||
The field can be annotated as ``list[float]`` for compatibility with
|
||||
static type checkers. LanceDB will infer the fixed vector dimension from
|
||||
this embedding function.
|
||||
"""
|
||||
return Field(json_schema_extra={"vector_column_for": self}, **kwargs)
|
||||
|
||||
|
||||
@@ -356,16 +356,23 @@ _ObjectVersion = Annotated[
|
||||
|
||||
|
||||
class FunctionVersion(_RemoteValue):
|
||||
"""A pinned object revision, independent of its executable image digest."""
|
||||
"""A pinned object revision, independent of its executable image digest.
|
||||
|
||||
``name`` and ``namespace_path`` are the Function's parts, as a
|
||||
[SecretReference][lancedb.functions.SecretReference] carries a Secret's;
|
||||
they are what [get_function][lancedb.db.DBConnection.get_function] takes.
|
||||
An empty ``namespace_path`` is the root namespace.
|
||||
"""
|
||||
|
||||
name: str
|
||||
namespace_path: tuple[str, ...] = ()
|
||||
object_id: str
|
||||
location: str
|
||||
version: _ObjectVersion
|
||||
image: FunctionImage
|
||||
signature: FunctionSignature
|
||||
secret_bindings: tuple[SecretBinding, ...] = ()
|
||||
created_at: str
|
||||
created_at_millis: int
|
||||
metadata: Mapping[str, str]
|
||||
disabled: bool
|
||||
|
||||
@@ -444,6 +451,7 @@ class FunctionVersion(_RemoteValue):
|
||||
return FunctionApplication(
|
||||
function=FunctionVersionRef(
|
||||
name=self.name,
|
||||
namespace_path=self.namespace_path,
|
||||
object_id=self.object_id,
|
||||
location=self.location,
|
||||
version=self.version,
|
||||
@@ -502,6 +510,7 @@ class FunctionRegistrationRequest(_RemoteValue):
|
||||
|
||||
class FunctionVersionRef(_OpenRemoteValue):
|
||||
name: str
|
||||
namespace_path: tuple[str, ...] = ()
|
||||
object_id: str
|
||||
location: str
|
||||
version: _ObjectVersion
|
||||
|
||||
@@ -112,9 +112,42 @@ class Fm:
|
||||
class FTS:
|
||||
"""Describe a FTS index configuration.
|
||||
|
||||
`FTS` is a full-text search index that can be used on `String` columns
|
||||
`FTS` supports string columns, lists of strings, and JSON columns created
|
||||
from PyArrow ``pa.json_()``. JSON columns use Lance's JSON-aware tokenizer
|
||||
and are stored as ``LargeBinary`` with the ``lance.json`` logical type.
|
||||
Arbitrary binary columns are not supported.
|
||||
|
||||
For example, it works with `title`, `description`, `content`, etc.
|
||||
JSON queries use ``path,str,value`` triples. Paths are relative to the JSON
|
||||
column and use dots for nested object keys. For example,
|
||||
``meta.tag,str,nature`` matches the string value at ``$.meta.tag``;
|
||||
it does not search other paths or JSON key names. The value is processed
|
||||
with the configured text tokenizer. Use ``str``, not ``string``.
|
||||
|
||||
Create and query a JSON index through the modern API::
|
||||
|
||||
import pyarrow as pa
|
||||
from lancedb.index import FTS
|
||||
|
||||
table = db.create_table(
|
||||
"documents",
|
||||
pa.table({"doc": pa.array(
|
||||
['{"meta":{"tag":"nature"}}'], type=pa.json_()
|
||||
)}),
|
||||
)
|
||||
table.create_index("doc", config=FTS())
|
||||
table.search(
|
||||
"meta.tag,str,nature", query_type="fts", fts_columns="doc"
|
||||
).to_list()
|
||||
|
||||
Create the index on the JSON column itself, not a JSON subpath such as
|
||||
``doc.meta.tag``. JSON-path scalar indexes do not provide FTS path scoping.
|
||||
Plain text such as ``nature`` is not a valid JSON FTS query. The current
|
||||
Lance tokenizer can panic on malformed triples; see
|
||||
https://github.com/lance-format/lance/issues/8812.
|
||||
|
||||
Remote tables use the same ``FTS()`` configuration and require a server
|
||||
that supports JSON FTS. Wait for ``table.wait_for_index(["doc_idx"])``
|
||||
before searching a newly created remote index with its default name.
|
||||
|
||||
Examples
|
||||
--------
|
||||
@@ -138,6 +171,7 @@ class FTS:
|
||||
- "simple": Splits text by whitespace and punctuation.
|
||||
- "whitespace": Split text by whitespace, but not punctuation.
|
||||
- "raw": No tokenization. The entire text is treated as a single token.
|
||||
- "code": Tokenizes source code and identifiers.
|
||||
- "ngram": N-gram tokenizer for substring-style matching.
|
||||
- "icu": ICU dictionary-based word segmentation.
|
||||
- "icu/split": ICU segmentation with simple-style delimiter splitting.
|
||||
@@ -151,12 +185,16 @@ class FTS:
|
||||
ignored.
|
||||
lower_case : bool, default True
|
||||
Whether to convert the token to lower case. This makes queries case-insensitive.
|
||||
stem : bool, default True
|
||||
stem : bool, optional
|
||||
Whether to stem the token. Stemming reduces words to their root form.
|
||||
For example, in English "running" and "runs" would both be reduced to "run".
|
||||
remove_stop_words : bool, default True
|
||||
``None`` uses the base tokenizer's default: False for ``code`` and
|
||||
``ngram``, True otherwise.
|
||||
remove_stop_words : bool, optional
|
||||
Whether to remove stop words. Stop words are common words that are often
|
||||
removed from text before indexing. For example, in English "the" and "and".
|
||||
``None`` uses the base tokenizer's default: False for ``code`` and
|
||||
``ngram``, True otherwise.
|
||||
custom_stop_words : list of str, optional
|
||||
Custom words replace the built-in language stop words
|
||||
and only take effect when ``remove_stop_words`` is True. ``None`` uses
|
||||
@@ -197,8 +235,8 @@ class FTS:
|
||||
language: str = "English"
|
||||
max_token_length: Optional[int] = 40
|
||||
lower_case: bool = True
|
||||
stem: bool = True
|
||||
remove_stop_words: bool = True
|
||||
stem: Optional[bool] = None
|
||||
remove_stop_words: Optional[bool] = None
|
||||
ascii_folding: bool = True
|
||||
ngram_min_length: int = 3
|
||||
ngram_max_length: int = 3
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Lazy, resumable resource listings."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, AsyncIterator, Callable, Generic, Iterator, Optional, TypeVar
|
||||
|
||||
from .background_loop import LOOP
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class AsyncListing(AsyncIterator[T], Generic[T]):
|
||||
"""A lazy async iterator with inspectable pagination state.
|
||||
|
||||
Returned by connection resource-listing methods. For example::
|
||||
|
||||
async for name in connection.list_views(page_limit=20):
|
||||
print(name)
|
||||
|
||||
Requests happen during iteration; request errors terminate the iterator.
|
||||
"""
|
||||
|
||||
def __init__(self, inner: Any, convert: Callable[[Any], T] = lambda value: value):
|
||||
self._inner = inner
|
||||
self._convert = convert
|
||||
|
||||
def __aiter__(self) -> AsyncListing[T]:
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> T:
|
||||
return self._convert(await self._inner.__anext__())
|
||||
|
||||
def num_page_results(self) -> int:
|
||||
"""Return the number of cached items available without another request."""
|
||||
return self._inner.num_page_results()
|
||||
|
||||
def page_token(self) -> Optional[str]:
|
||||
"""Return the token for the next request.
|
||||
|
||||
Initially this is the supplied starting token (``None`` starts at the
|
||||
beginning). After fetching the final page it is ``None``, even if items
|
||||
remain cached. Drain the cache before saving a token to avoid skipping
|
||||
items when resuming. A failed request terminates iteration and retains
|
||||
its token for retrying with a new iterator.
|
||||
"""
|
||||
return self._inner.page_token()
|
||||
|
||||
|
||||
class Listing(Iterator[T], Generic[T]):
|
||||
"""A lazy synchronous iterator with inspectable pagination state.
|
||||
|
||||
For example::
|
||||
|
||||
for name in connection.list_views(page_limit=20):
|
||||
print(name)
|
||||
|
||||
Use ``list(connection.list_views())`` to collect all results.
|
||||
"""
|
||||
|
||||
def __init__(self, inner: AsyncListing[T]):
|
||||
self._inner = inner
|
||||
|
||||
def __iter__(self) -> Listing[T]:
|
||||
return self
|
||||
|
||||
def __next__(self) -> T:
|
||||
try:
|
||||
return LOOP.run(self._inner.__anext__())
|
||||
except StopAsyncIteration:
|
||||
raise StopIteration from None
|
||||
|
||||
def num_page_results(self) -> int:
|
||||
"""Return the number of cached items available without another request."""
|
||||
return self._inner.num_page_results()
|
||||
|
||||
def page_token(self) -> Optional[str]:
|
||||
"""Return the next request's token.
|
||||
|
||||
See [AsyncListing.page_token][lancedb.listing.AsyncListing.page_token]
|
||||
for initial, cached, final-page, and error behavior.
|
||||
"""
|
||||
return self._inner.page_token()
|
||||
@@ -6,11 +6,11 @@ maintained by refresh. See ``DBConnection.create_materialized_view``."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .background_loop import LOOP
|
||||
from .job import AsyncJob, Job, _typed_job
|
||||
|
||||
@@ -21,7 +21,6 @@ if TYPE_CHECKING:
|
||||
from .table import AsyncTable, LanceTable
|
||||
|
||||
DEFINITION_META_KEY = b"mv.definition"
|
||||
_MAX_FLOAT32 = float.fromhex("0x1.fffffep+127")
|
||||
|
||||
SelectArg = Union[
|
||||
str,
|
||||
@@ -31,104 +30,12 @@ SelectArg = Union[
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class VectorDedupSource:
|
||||
"""A declarative source of retained original rows from an indexed snapshot.
|
||||
|
||||
Construct with [vector_dedup][lancedb.vector_dedup] and pass it as the
|
||||
``source`` of ``create_materialized_view`` or
|
||||
``create_materialized_view_async``. An omitted ``dataset_version`` is
|
||||
captured once when creation is submitted, not when this object is made.
|
||||
Refreshes keep that captured version. The source table is never modified.
|
||||
"""
|
||||
|
||||
source: str
|
||||
column: str
|
||||
distance_threshold: float
|
||||
dataset_version: Optional[int] = None
|
||||
|
||||
def __post_init__(self):
|
||||
for name in ("source", "column"):
|
||||
value = getattr(self, name)
|
||||
if not isinstance(value, str) or not value:
|
||||
raise ValueError(f"{name} must be a non-empty string")
|
||||
if self.dataset_version is not None and (
|
||||
type(self.dataset_version) is not int
|
||||
or not 0 < self.dataset_version < 2**64
|
||||
):
|
||||
raise ValueError("dataset_version must be a positive uint64 integer")
|
||||
value = self.distance_threshold
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
||||
raise TypeError("distance_threshold must be a number")
|
||||
try:
|
||||
value = float(value)
|
||||
except OverflowError as exc:
|
||||
raise ValueError("distance_threshold must fit a finite float32") from exc
|
||||
if not math.isfinite(value) or abs(value) > _MAX_FLOAT32:
|
||||
raise ValueError("distance_threshold must fit a finite float32")
|
||||
object.__setattr__(self, "distance_threshold", value)
|
||||
|
||||
def _native_source_json(self, version: int) -> str:
|
||||
return json.dumps(
|
||||
{
|
||||
"kind": "Dedup",
|
||||
"dataset_version": version,
|
||||
"column": self.column,
|
||||
"distance_threshold": repr(self.distance_threshold),
|
||||
},
|
||||
allow_nan=False,
|
||||
)
|
||||
MaterializedViewSource = str
|
||||
|
||||
|
||||
def vector_dedup(
|
||||
source: str,
|
||||
*,
|
||||
column: str,
|
||||
distance_threshold: float,
|
||||
dataset_version: Optional[int] = None,
|
||||
) -> VectorDedupSource:
|
||||
"""Declare an indexed dedup source for a materialized view.
|
||||
|
||||
The default policy keeps the smallest unassigned snapshot row ID and removes
|
||||
only its direct qualifying neighbors: an A-B-C chain without an A-C edge
|
||||
retains A and C. Rows with no pair, including null vectors, survive.
|
||||
Only same-segment/same-IVF-partition pairs are considered; quantized distances
|
||||
use reconstructed index vectors. The index must cover the selected snapshot.
|
||||
|
||||
``source`` is an exact table name in the connection's root namespace, just
|
||||
as for ``create_materialized_view``. ``dataset_version=None`` captures the
|
||||
source's version when creation is submitted. Pass an explicit positive
|
||||
version to reproduce a previous snapshot. The result contains original
|
||||
source columns and rows. Do not combine this source with ``select``,
|
||||
``where`` or ``limit``; query the resulting table for further transforms.
|
||||
|
||||
Creation follows the usual MV contract: by default it waits for the result.
|
||||
Use ``with_no_data=True`` to declare only, or
|
||||
``create_materialized_view_async`` to get a job handle. On Cloud/Enterprise,
|
||||
execution uses server-side MV registry jobs; local connections run in process.
|
||||
|
||||
Examples
|
||||
--------
|
||||
Given a connected ``db`` and an indexed ``images.phash`` column::
|
||||
|
||||
from lancedb import vector_dedup
|
||||
|
||||
view = db.create_materialized_view(
|
||||
"images_clean",
|
||||
vector_dedup("images", column="phash", distance_threshold=4),
|
||||
)
|
||||
cleaned = view.table
|
||||
"""
|
||||
return VectorDedupSource(source, column, distance_threshold, dataset_version)
|
||||
|
||||
|
||||
MaterializedViewSource = Union[str, VectorDedupSource]
|
||||
|
||||
|
||||
DEFINITION_FORMAT = 4
|
||||
DEFINITION_FORMAT = 2
|
||||
"""The newest stored layout this version reads: ``{"format": N, "query": "<SQL>"}``,
|
||||
format 2 being a query with ``GROUP BY``, format 3 native pairs and format 4
|
||||
native dedup results. A ``kind`` key beside it is for
|
||||
format 2 being a query with ``GROUP BY``. A ``kind`` key beside it is for
|
||||
readers older than the format number."""
|
||||
|
||||
|
||||
@@ -272,38 +179,31 @@ class AsyncMaterializedView:
|
||||
return _definition_from_json(raw)
|
||||
|
||||
async def refresh(
|
||||
self, *, full: bool = False, source_version: Optional[int] = None
|
||||
self, *, source_version: Optional[int] = None
|
||||
) -> "RefreshMaterializedViewResult":
|
||||
"""Recompute the view from its source.
|
||||
|
||||
The refresh is incremental when the source's changes can be
|
||||
reconciled into the view -- rows added, changed or removed since the
|
||||
last one -- and otherwise rebuilds. ``full=True`` forces a rebuild;
|
||||
Every refresh rebuilds the complete view. For local views,
|
||||
``source_version`` refreshes to that source version instead of the
|
||||
latest.
|
||||
|
||||
Concurrent refreshes of one view do not duplicate its rows. Two that
|
||||
plan the same source rows conflict on commit, and the loser raises
|
||||
rather than writing them a second time.
|
||||
latest; remote SQL refreshes do not support source-version pinning.
|
||||
"""
|
||||
return await self._table._inner.refresh_materialized_view(
|
||||
full=full, source_version=source_version
|
||||
source_version=source_version
|
||||
)
|
||||
|
||||
async def refresh_async(
|
||||
self, *, full: bool = False, source_version: Optional[int] = None
|
||||
self, *, source_version: Optional[int] = None
|
||||
) -> "AsyncJob[RefreshMaterializedViewResult]":
|
||||
"""Submit a refresh and return its job without waiting.
|
||||
|
||||
The job may already be complete for a local view. On LanceDB Cloud
|
||||
and Enterprise, its ``id`` is the server job identifier returned by
|
||||
the refresh endpoint.
|
||||
The job may already be complete for a local view. Remote refreshes are
|
||||
submitted through the SQL service.
|
||||
"""
|
||||
from ._lancedb import RefreshMaterializedViewResult
|
||||
|
||||
return _typed_job(
|
||||
await self._table._inner.refresh_materialized_view_async(
|
||||
full=full, source_version=source_version
|
||||
source_version=source_version
|
||||
),
|
||||
RefreshMaterializedViewResult.from_json,
|
||||
)
|
||||
@@ -335,22 +235,18 @@ class MaterializedView:
|
||||
return LOOP.run(self._async.definition())
|
||||
|
||||
def refresh(
|
||||
self, *, full: bool = False, source_version: Optional[int] = None
|
||||
self, *, source_version: Optional[int] = None
|
||||
) -> "RefreshMaterializedViewResult":
|
||||
"""Recompute the view from its source. See
|
||||
[AsyncMaterializedView.refresh][lancedb.materialized_view.AsyncMaterializedView.refresh]."""
|
||||
return LOOP.run(self._async.refresh(full=full, source_version=source_version))
|
||||
return LOOP.run(self._async.refresh(source_version=source_version))
|
||||
|
||||
def refresh_async(
|
||||
self, *, full: bool = False, source_version: Optional[int] = None
|
||||
self, *, source_version: Optional[int] = None
|
||||
) -> "Job[RefreshMaterializedViewResult]":
|
||||
"""Submit a refresh and return its job without waiting.
|
||||
|
||||
See
|
||||
[AsyncMaterializedView.refresh_async][lancedb.materialized_view.AsyncMaterializedView.refresh_async].
|
||||
"""
|
||||
return Job(
|
||||
LOOP.run(
|
||||
self._async.refresh_async(full=full, source_version=source_version)
|
||||
)
|
||||
)
|
||||
return Job(LOOP.run(self._async.refresh_async(source_version=source_version)))
|
||||
|
||||
@@ -37,7 +37,7 @@ class LanceMergeInsertBuilder(object):
|
||||
self._when_not_matched_by_source_condition_expr = None
|
||||
self._timeout = None
|
||||
self._use_index = True
|
||||
self._use_lsm = None
|
||||
self._use_memwal = None
|
||||
self._validate_single_shard = None
|
||||
|
||||
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_memwal(self, enable: bool) -> LanceMergeInsertBuilder:
|
||||
"""
|
||||
Control MemWAL routing for this merge.
|
||||
|
||||
By default (unset), a `merge_insert` on a table with an LSM write spec is
|
||||
By default (unset), a `merge_insert` on a table with a MemWAL write spec is
|
||||
routed through Lance's MemWAL shard writer, and a table without one uses
|
||||
the standard path.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
enable: bool
|
||||
``True`` forces MemWAL routing and errors if the table has no LSM
|
||||
``True`` forces MemWAL routing and errors if the table has no MemWAL
|
||||
write spec. ``False`` forces the standard write path even when a spec
|
||||
is set.
|
||||
"""
|
||||
self._use_lsm = enable
|
||||
self._use_memwal = enable
|
||||
return self
|
||||
|
||||
def validate_single_shard(
|
||||
@@ -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 a MemWAL write spec, every row in a `merge_insert` call
|
||||
must route to the same shard. When `True` (the default), every row is
|
||||
inspected to verify this. When `False`, only the first row is inspected
|
||||
and the shard it routes to is used for the whole input — a faster path
|
||||
for callers that have already pre-sharded their input.
|
||||
|
||||
Has no effect on tables without an LSM write spec.
|
||||
Has no effect on tables without a MemWAL write spec.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
@@ -11,7 +11,7 @@ through a namespace abstraction.
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
|
||||
from typing import Any, Dict, Iterable, List, Optional, Union
|
||||
from uuid import UUID
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
@@ -19,8 +19,7 @@ if sys.version_info >= (3, 12):
|
||||
else:
|
||||
from overrides import override
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lancedb.query import Query
|
||||
from lancedb.query import Query
|
||||
|
||||
from datetime import timedelta
|
||||
import pyarrow as pa
|
||||
@@ -52,8 +51,9 @@ from lancedb.background_loop import LOOP
|
||||
from lancedb.arrow import AsyncRecordBatchReader
|
||||
from lancedb.db import AsyncConnection, DBConnection
|
||||
from lancedb.job import AsyncJob, Job
|
||||
from lancedb.listing import AsyncListing, Listing
|
||||
from lancedb.sql import AsyncQuery as AsyncSqlQuery
|
||||
from lancedb.sql import QueryDescription
|
||||
from lancedb.sql import QueryDescription, QueryParameters
|
||||
from lance_namespace import (
|
||||
LanceNamespace,
|
||||
connect as namespace_connect,
|
||||
@@ -105,6 +105,11 @@ def _query_to_namespace_request(
|
||||
|
||||
if not isinstance(query, Query):
|
||||
raise TypeError(f"Expected Query, got {type(query)}")
|
||||
if query.search_effort is not None:
|
||||
raise ValueError(
|
||||
"namespace QueryTable does not support search_effort; "
|
||||
"execute the query locally"
|
||||
)
|
||||
|
||||
vector_request = None
|
||||
if query.vector is not None:
|
||||
@@ -190,8 +195,12 @@ def _query_to_namespace_request(
|
||||
kwargs["with_row_id"] = query.with_row_id
|
||||
if query.ef is not None:
|
||||
kwargs["ef"] = query.ef
|
||||
if query.nprobes is not None:
|
||||
kwargs["nprobes"] = query.nprobes
|
||||
if query.minimum_nprobes is not None:
|
||||
kwargs["nprobes"] = query.minimum_nprobes
|
||||
kwargs["minimum_nprobes"] = query.minimum_nprobes
|
||||
if query.maximum_nprobes is not None:
|
||||
kwargs["maximum_nprobes"] = query.maximum_nprobes
|
||||
if query.refine_factor is not None:
|
||||
kwargs["refine_factor"] = query.refine_factor
|
||||
if query.lower_bound is not None:
|
||||
@@ -205,6 +214,13 @@ def _query_to_namespace_request(
|
||||
if query.distance_type is not None:
|
||||
kwargs["distance_type"] = query.distance_type
|
||||
|
||||
if kwargs.get("maximum_nprobes") == 0:
|
||||
# The generated 0.13 model rejects LanceDB's unbounded sentinel;
|
||||
# model_copy preserves direct wire pass-through to the Lance receiver.
|
||||
maximum_nprobes = kwargs.pop("maximum_nprobes")
|
||||
return QueryTableRequest(**kwargs).model_copy(
|
||||
update={"maximum_nprobes": maximum_nprobes}
|
||||
)
|
||||
return QueryTableRequest(**kwargs)
|
||||
|
||||
|
||||
@@ -489,6 +505,14 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
)
|
||||
self._uri = self._inner.uri
|
||||
|
||||
@override
|
||||
def is_open(self) -> bool:
|
||||
return self._inner.is_open()
|
||||
|
||||
@override
|
||||
def close(self) -> None:
|
||||
self._inner.close()
|
||||
|
||||
@override
|
||||
def serialize(self) -> str:
|
||||
import json
|
||||
@@ -514,13 +538,15 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
def table_names(
|
||||
self,
|
||||
page_token: Optional[str] = None,
|
||||
limit: int = 10,
|
||||
limit: Optional[int] = None,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> Iterable[str]:
|
||||
"""
|
||||
List table names in the database.
|
||||
|
||||
If limit is None, return all table names; zero returns an empty list.
|
||||
|
||||
.. deprecated::
|
||||
Use :meth:`list_tables` instead, which provides proper pagination support.
|
||||
"""
|
||||
@@ -690,9 +716,20 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
return view
|
||||
|
||||
@override
|
||||
def list_materialized_views(self) -> List[str]:
|
||||
"""The names of the materialized views in the root namespace."""
|
||||
return LOOP.run(self._inner.list_materialized_views())
|
||||
def list_materialized_views(
|
||||
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||
) -> Listing[str]:
|
||||
"""The names of the materialized views in the root namespace.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
return Listing(
|
||||
self._inner.list_materialized_views(
|
||||
page_token=page_token, page_limit=page_limit
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def drop_materialized_view(
|
||||
@@ -1138,13 +1175,15 @@ class AsyncLanceNamespaceDBConnection:
|
||||
async def table_names(
|
||||
self,
|
||||
page_token: Optional[str] = None,
|
||||
limit: int = 10,
|
||||
limit: Optional[int] = None,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> Iterable[str]:
|
||||
"""
|
||||
List table names in the namespace.
|
||||
|
||||
If limit is None, return all table names; zero returns an empty list.
|
||||
|
||||
.. deprecated::
|
||||
Use :meth:`list_tables` instead, which provides proper pagination support.
|
||||
"""
|
||||
@@ -1289,9 +1328,20 @@ class AsyncLanceNamespaceDBConnection:
|
||||
await view.definition()
|
||||
return view
|
||||
|
||||
async def list_materialized_views(self) -> List[str]:
|
||||
"""The names of the materialized views in the root namespace."""
|
||||
return await self._inner.list_materialized_views()
|
||||
def list_materialized_views(
|
||||
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||
) -> AsyncListing[str]:
|
||||
"""The names of the materialized views in the root namespace.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
return AsyncListing(
|
||||
self._inner.list_materialized_views(
|
||||
page_token=page_token, page_limit=page_limit
|
||||
)
|
||||
)
|
||||
|
||||
async def drop_materialized_view(
|
||||
self, name: str, namespace_path: Optional[List[str]] = None
|
||||
@@ -1554,11 +1604,13 @@ class AsyncLanceNamespaceDBConnection:
|
||||
query: str,
|
||||
*,
|
||||
default_namespace_path: Optional[List[str]] = None,
|
||||
parameters: Optional[QueryParameters] = None,
|
||||
) -> AsyncRecordBatchReader:
|
||||
"""Execute SQL when supported by the underlying connection."""
|
||||
return await self._inner.execute_query(
|
||||
query,
|
||||
default_namespace_path=default_namespace_path,
|
||||
parameters=parameters,
|
||||
)
|
||||
|
||||
async def execute_query_async(
|
||||
@@ -1566,6 +1618,7 @@ class AsyncLanceNamespaceDBConnection:
|
||||
query: str,
|
||||
*,
|
||||
default_namespace_path: Optional[List[str]] = None,
|
||||
parameters: Optional[QueryParameters] = None,
|
||||
) -> AsyncSqlQuery:
|
||||
"""Start executing SQL when supported by the underlying connection.
|
||||
|
||||
@@ -1574,6 +1627,7 @@ class AsyncLanceNamespaceDBConnection:
|
||||
return await self._inner.execute_query_async(
|
||||
query,
|
||||
default_namespace_path=default_namespace_path,
|
||||
parameters=parameters,
|
||||
)
|
||||
|
||||
async def describe_query(self, query_id: UUID) -> QueryDescription:
|
||||
|
||||
@@ -48,7 +48,7 @@ class PermutationBuilder:
|
||||
By default, the permutation builder will create a single split that contains all
|
||||
rows in the same order as the base table.
|
||||
|
||||
Tables with an LSM write spec are rejected: unflushed rows have no row id.
|
||||
Tables with a MemWAL write spec are rejected: unflushed rows have no row id.
|
||||
"""
|
||||
self._async = async_permutation_builder(table)
|
||||
|
||||
|
||||
@@ -88,6 +88,17 @@ def Vector(
|
||||
... pa.field("url", pa.utf8(), False),
|
||||
... pa.field("embeddings", pa.list_(pa.float32(), 768))
|
||||
... ])
|
||||
|
||||
Notes
|
||||
-----
|
||||
``Vector`` creates a type dynamically, so calls such as ``Vector(768)`` are
|
||||
not valid static type annotations. For an embedding field, use the standard
|
||||
``list[float]`` annotation when running mypy; ``VectorField`` supplies the
|
||||
fixed dimension to LanceDB::
|
||||
|
||||
class MyModel(LanceModel):
|
||||
text: str = embeddings.SourceField()
|
||||
vector: list[float] = embeddings.VectorField()
|
||||
"""
|
||||
|
||||
# TODO: make a public parameterized type.
|
||||
@@ -313,6 +324,10 @@ def _unwrap_optional_annotation(annotation: Any) -> Any | None:
|
||||
|
||||
def _pydantic_to_arrow_type(field: FieldInfo) -> pa.DataType:
|
||||
"""Convert a Pydantic FieldInfo to Arrow DataType"""
|
||||
embedding_vector_type = _embedding_vector_to_arrow_type(field)
|
||||
if embedding_vector_type is not None:
|
||||
return embedding_vector_type
|
||||
|
||||
unwrapped = _unwrap_optional_annotation(field.annotation)
|
||||
if unwrapped is not None:
|
||||
return _pydantic_type_to_arrow_type(unwrapped, field)
|
||||
@@ -326,8 +341,32 @@ def _pydantic_to_arrow_type(field: FieldInfo) -> pa.DataType:
|
||||
return _pydantic_type_to_arrow_type(field.annotation, field)
|
||||
|
||||
|
||||
def _embedding_vector_to_arrow_type(field: FieldInfo) -> pa.DataType | None:
|
||||
"""Infer a fixed-size vector type from ``VectorField`` metadata."""
|
||||
if not _is_embedding_vector_annotation(field):
|
||||
return None
|
||||
|
||||
function = get_extras(field, "vector_column_for")
|
||||
return pa.list_(pa.float32(), function.ndims())
|
||||
|
||||
|
||||
def _is_embedding_vector_annotation(field: FieldInfo) -> bool:
|
||||
if get_extras(field, "vector_column_for") is None:
|
||||
return False
|
||||
|
||||
annotation = _unwrap_optional_annotation(field.annotation)
|
||||
if annotation is None:
|
||||
annotation = field.annotation
|
||||
|
||||
origin = getattr(annotation, "__origin__", None)
|
||||
args = getattr(annotation, "__args__", ())
|
||||
return origin is list and args == (float,)
|
||||
|
||||
|
||||
def is_nullable(field: FieldInfo) -> bool:
|
||||
"""Check if a Pydantic FieldInfo is nullable."""
|
||||
if _is_embedding_vector_annotation(field):
|
||||
return True
|
||||
if _unwrap_optional_annotation(field.annotation) is not None:
|
||||
return True
|
||||
if isinstance(field.annotation, (_GenericAlias, GenericAlias)):
|
||||
|
||||
+587
-160
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
|
||||
@@ -2,6 +2,8 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import replace
|
||||
from datetime import timedelta
|
||||
import json
|
||||
@@ -34,11 +36,12 @@ from lancedb.remote import ClientConfig, RetryConfig, TimeoutConfig, TlsConfig
|
||||
import pyarrow as pa
|
||||
|
||||
from ..common import DATA
|
||||
from lancedb.listing import Listing
|
||||
from ..db import DBConnection, LOOP
|
||||
from ..functions import FunctionVersion, UdfDefinition
|
||||
from ..job import AsyncJob, Job
|
||||
from ..sql import Query as SqlQuery
|
||||
from ..sql import QueryDescription
|
||||
from ..sql import QueryDescription, QueryParameters
|
||||
from ..materialized_view import MaterializedView, MaterializedViewSource, SelectArg
|
||||
from ..secrets import EnvVarSecret, SecretInfo
|
||||
from ..view import ViewDescription
|
||||
@@ -109,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),
|
||||
@@ -376,7 +380,7 @@ class RemoteDBConnection(DBConnection):
|
||||
page_token: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> ListTablesResponse:
|
||||
"""List all tables in this database with pagination support.
|
||||
"""List a page of tables in this database.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
@@ -387,7 +391,8 @@ class RemoteDBConnection(DBConnection):
|
||||
Token for pagination. Use the token from a previous response
|
||||
to get the next page of results.
|
||||
limit: int, optional
|
||||
The maximum number of results to return.
|
||||
The maximum number of results to return, default 100.
|
||||
Zero returns an empty page without a continuation token.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -406,7 +411,7 @@ class RemoteDBConnection(DBConnection):
|
||||
def table_names(
|
||||
self,
|
||||
page_token: Optional[str] = None,
|
||||
limit: int = 10,
|
||||
limit: Optional[int] = None,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> Iterable[str]:
|
||||
@@ -422,8 +427,9 @@ class RemoteDBConnection(DBConnection):
|
||||
Empty list represents root namespace.
|
||||
page_token: str
|
||||
The last token to start the new page.
|
||||
limit: int, default 10
|
||||
The maximum number of tables to return for each page.
|
||||
limit: int, optional
|
||||
The maximum number of tables to return. None returns all tables;
|
||||
zero returns an empty list.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -483,12 +489,12 @@ class RemoteDBConnection(DBConnection):
|
||||
namespace_path = []
|
||||
if storage_options is not None:
|
||||
logging.info(
|
||||
"storage_options is ignored in LanceDb Cloud"
|
||||
"storage_options is ignored for remote tables"
|
||||
" (storage is managed; set storage_options on connect() instead)"
|
||||
)
|
||||
if index_cache_size is not None:
|
||||
logging.info(
|
||||
"index_cache_size is ignored in LanceDb Cloud"
|
||||
"index_cache_size is ignored for remote tables"
|
||||
" (there is no local cache to configure)"
|
||||
)
|
||||
|
||||
@@ -608,6 +614,9 @@ class RemoteDBConnection(DBConnection):
|
||||
One of "error", "drop", "fill".
|
||||
fill_value: float
|
||||
The value to use when filling vectors. Only used if on_bad_vectors="fill".
|
||||
embedding_functions: list of EmbeddingFunctionConfig, optional
|
||||
The embedding functions to store in the table schema and use to
|
||||
generate vectors when creating the table or adding data.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -682,12 +691,6 @@ class RemoteDBConnection(DBConnection):
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
validate_table_name(name)
|
||||
if embedding_functions is not None:
|
||||
logging.warning(
|
||||
"embedding_functions is not yet supported on LanceDB Cloud."
|
||||
"Please vote https://github.com/lancedb/lancedb/issues/626 "
|
||||
"for this feature."
|
||||
)
|
||||
|
||||
from .table import RemoteTable
|
||||
|
||||
@@ -700,6 +703,7 @@ class RemoteDBConnection(DBConnection):
|
||||
schema=schema,
|
||||
on_bad_vectors=on_bad_vectors,
|
||||
fill_value=fill_value,
|
||||
embedding_functions=embedding_functions,
|
||||
)
|
||||
)
|
||||
return RemoteTable(
|
||||
@@ -771,8 +775,20 @@ class RemoteDBConnection(DBConnection):
|
||||
return view
|
||||
|
||||
@override
|
||||
def list_materialized_views(self) -> List[str]:
|
||||
return LOOP.run(self._conn.list_materialized_views())
|
||||
def list_materialized_views(
|
||||
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||
) -> Listing[str]:
|
||||
"""List materialized views.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
return Listing(
|
||||
self._conn.list_materialized_views(
|
||||
page_token=page_token, page_limit=page_limit
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def drop_materialized_view(
|
||||
@@ -862,25 +878,78 @@ class RemoteDBConnection(DBConnection):
|
||||
definition: UdfDefinition,
|
||||
*,
|
||||
secrets: Optional[Sequence[EnvVarSecret]] = None,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> Job[FunctionVersion]:
|
||||
job = LOOP.run(self._conn.create_function_async(definition, secrets=secrets))
|
||||
job = LOOP.run(
|
||||
self._conn.create_function_async(
|
||||
definition, secrets=secrets, namespace_path=namespace_path
|
||||
)
|
||||
)
|
||||
return Job(job)
|
||||
|
||||
@override
|
||||
def get_function(self, name: str, *, version: str) -> FunctionVersion:
|
||||
return LOOP.run(self._conn.get_function(name, version=version))
|
||||
def get_function(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
version: str,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> FunctionVersion:
|
||||
return LOOP.run(
|
||||
self._conn.get_function(
|
||||
name, version=version, namespace_path=namespace_path
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def list_functions(self) -> List[FunctionVersion]:
|
||||
return LOOP.run(self._conn.list_functions())
|
||||
def list_functions(
|
||||
self,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
page_limit: Optional[int] = None,
|
||||
) -> Listing[FunctionVersion]:
|
||||
"""List functions.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
return Listing(
|
||||
self._conn.list_functions(
|
||||
namespace_path=namespace_path,
|
||||
page_token=page_token,
|
||||
page_limit=page_limit,
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def drop_function(self, name: str, *, version: str) -> bool:
|
||||
return LOOP.run(self._conn.drop_function(name, version=version))
|
||||
def drop_function(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
version: str,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> bool:
|
||||
return LOOP.run(
|
||||
self._conn.drop_function(
|
||||
name, version=version, namespace_path=namespace_path
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def drop_function_async(self, name: str, *, version: str) -> Tuple[bool, Job]:
|
||||
dropped, job = LOOP.run(self._conn.drop_function_async(name, version=version))
|
||||
def drop_function_async(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
version: str,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
) -> Tuple[bool, Job]:
|
||||
dropped, job = LOOP.run(
|
||||
self._conn.drop_function_async(
|
||||
name, version=version, namespace_path=namespace_path
|
||||
)
|
||||
)
|
||||
return dropped, Job(job)
|
||||
|
||||
@override
|
||||
@@ -902,8 +971,26 @@ class RemoteDBConnection(DBConnection):
|
||||
return LOOP.run(self._conn.describe_secret(name, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def list_secrets(self, *, namespace_path: Optional[List[str]] = None) -> List[str]:
|
||||
return LOOP.run(self._conn.list_secrets(namespace_path=namespace_path))
|
||||
def list_secrets(
|
||||
self,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
page_limit: Optional[int] = None,
|
||||
) -> Listing[str]:
|
||||
"""List secrets.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
return Listing(
|
||||
self._conn.list_secrets(
|
||||
namespace_path=namespace_path,
|
||||
page_token=page_token,
|
||||
page_limit=page_limit,
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def drop_secret(
|
||||
@@ -939,13 +1026,40 @@ class RemoteDBConnection(DBConnection):
|
||||
return Job(job)
|
||||
|
||||
@override
|
||||
def list_views(self, *, namespace_path: Optional[List[str]] = None) -> List[str]:
|
||||
return LOOP.run(self._conn.list_views(namespace_path=namespace_path))
|
||||
def list_views(
|
||||
self,
|
||||
*,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
page_token: Optional[str] = None,
|
||||
page_limit: Optional[int] = None,
|
||||
) -> Listing[str]:
|
||||
"""List views.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
return Listing(
|
||||
self._conn.list_views(
|
||||
namespace_path=namespace_path,
|
||||
page_token=page_token,
|
||||
page_limit=page_limit,
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def list_jobs(self) -> List["JobInfo"]:
|
||||
"""List server-side jobs across the database's tables."""
|
||||
return LOOP.run(self._conn.list_jobs())
|
||||
def list_jobs(
|
||||
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
|
||||
) -> Listing[JobInfo]:
|
||||
"""List server-side jobs across the database's tables.
|
||||
|
||||
Returns a lazy iterator. ``page_limit`` limits each request, not the total;
|
||||
``page_token`` resumes from a saved token. Requests and errors occur during
|
||||
iteration. See [pagination state][lancedb.listing.AsyncListing].
|
||||
"""
|
||||
return Listing(
|
||||
self._conn.list_jobs(page_token=page_token, page_limit=page_limit)
|
||||
)
|
||||
|
||||
@override
|
||||
def cancel_job(self, job_id: str) -> bool:
|
||||
@@ -979,18 +1093,22 @@ class RemoteDBConnection(DBConnection):
|
||||
query: str,
|
||||
*,
|
||||
default_namespace_path: Optional[List[str]] = None,
|
||||
parameters: Optional[QueryParameters] = None,
|
||||
) -> SqlQuery:
|
||||
"""Start executing SQL through this remote connection.
|
||||
|
||||
Unqualified tables use this connection's database and the
|
||||
``["public"]`` namespace by default. Fully qualified table names may
|
||||
reference other databases available to the same deployment.
|
||||
``parameters`` binds values to ``$1`` / ``$name`` placeholders; see
|
||||
[DBConnection.execute_query_async][lancedb.db.DBConnection.execute_query_async].
|
||||
"""
|
||||
return SqlQuery(
|
||||
LOOP.run(
|
||||
self._conn.execute_query_async(
|
||||
query,
|
||||
default_namespace_path=default_namespace_path,
|
||||
parameters=parameters,
|
||||
)
|
||||
)
|
||||
)
|
||||
@@ -1017,6 +1135,10 @@ class RemoteDBConnection(DBConnection):
|
||||
"""
|
||||
return LOOP.run(self._conn.namespace_client())
|
||||
|
||||
async def close(self):
|
||||
"""Close the connection to the database."""
|
||||
@override
|
||||
def is_open(self) -> bool:
|
||||
return self._conn.is_open()
|
||||
|
||||
@override
|
||||
def close(self) -> None:
|
||||
self._conn.close()
|
||||
|
||||
@@ -31,7 +31,7 @@ from lancedb._lancedb import (
|
||||
DropColumnsResult,
|
||||
FtsToken,
|
||||
IndexConfig,
|
||||
LsmWriteSpec,
|
||||
MemWalWriteSpec,
|
||||
MergeResult,
|
||||
UpdateResult,
|
||||
)
|
||||
@@ -42,6 +42,7 @@ from lancedb.index import (
|
||||
BTree,
|
||||
Bitmap,
|
||||
HnswFlat,
|
||||
HnswPq,
|
||||
HnswSq,
|
||||
IvfFlat,
|
||||
IvfPq,
|
||||
@@ -56,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
|
||||
@@ -63,6 +65,7 @@ from lancedb.table import _normalize_progress
|
||||
from ..query import (
|
||||
AnalyzePlanDistributedMetrics,
|
||||
DocumentGranularity,
|
||||
FullTextQuery,
|
||||
LanceQueryBuilder,
|
||||
LanceTakeQueryBuilder,
|
||||
LanceVectorQueryBuilder,
|
||||
@@ -78,6 +81,7 @@ from ..table import (
|
||||
Tags,
|
||||
)
|
||||
from ..types import BaseTokenizerType
|
||||
from ..util import infer_vector_column_name
|
||||
|
||||
|
||||
class RemoteTable(Table):
|
||||
@@ -237,13 +241,15 @@ class RemoteTable(Table):
|
||||
"""List all versions of the table"""
|
||||
return LOOP.run(self._table.list_versions())
|
||||
|
||||
def to_arrow(self) -> pa.Table:
|
||||
"""to_arrow() is not yet supported on LanceDB cloud."""
|
||||
raise NotImplementedError("to_arrow() is not yet supported on LanceDB cloud.")
|
||||
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 yet supported on LanceDB cloud."""
|
||||
raise NotImplementedError("to_pandas() is not yet supported on LanceDB cloud.")
|
||||
"""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]):
|
||||
result = LOOP.run(self._table.checkout(version))
|
||||
@@ -352,8 +358,8 @@ class RemoteTable(Table):
|
||||
language: str = "English",
|
||||
max_token_length: Optional[int] = 40,
|
||||
lower_case: bool = True,
|
||||
stem: bool = True,
|
||||
remove_stop_words: bool = True,
|
||||
stem: Optional[bool] = None,
|
||||
remove_stop_words: Optional[bool] = None,
|
||||
custom_stop_words: Optional[List[str]] = None,
|
||||
ascii_folding: bool = True,
|
||||
ngram_min_length: int = 3,
|
||||
@@ -500,14 +506,8 @@ class RemoteTable(Table):
|
||||
|
||||
if accelerator is not None:
|
||||
logging.warning(
|
||||
"GPU accelerator is not yet supported on LanceDB cloud."
|
||||
"If you have 100M+ vectors to index,"
|
||||
"please contact us at contact@lancedb.com"
|
||||
)
|
||||
if replace is not None:
|
||||
logging.warning(
|
||||
"replace is not supported on LanceDB cloud."
|
||||
"Existing indexes will always be replaced."
|
||||
"GPU accelerator is not supported for remote tables "
|
||||
"and will be ignored."
|
||||
)
|
||||
|
||||
idx_type = index_type.upper()
|
||||
@@ -527,9 +527,11 @@ class RemoteTable(Table):
|
||||
elif idx_type == "IVF_SQ":
|
||||
config = IvfSq(distance_type=metric, num_partitions=num_partitions)
|
||||
elif idx_type == "IVF_HNSW_PQ":
|
||||
raise ValueError(
|
||||
"IVF_HNSW_PQ is not supported on LanceDB cloud."
|
||||
"Please use IVF_HNSW_SQ instead."
|
||||
config = HnswPq(
|
||||
distance_type=metric,
|
||||
num_partitions=num_partitions,
|
||||
num_sub_vectors=num_sub_vectors,
|
||||
num_bits=num_bits,
|
||||
)
|
||||
elif idx_type == "IVF_HNSW_SQ":
|
||||
config = HnswSq(distance_type=metric, num_partitions=num_partitions)
|
||||
@@ -658,7 +660,7 @@ class RemoteTable(Table):
|
||||
capped at the number of CPU cores. Lower this if bulk ingestion is
|
||||
using too much memory.
|
||||
allow_external_blob_outside_bases: bool, default False
|
||||
Not supported on LanceDB Cloud. Setting this raises.
|
||||
Not supported for remote tables. Setting this raises.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -688,7 +690,7 @@ class RemoteTable(Table):
|
||||
vector_column_name: Optional[str] = None,
|
||||
query_type="auto",
|
||||
fts_columns: Optional[Union[str, List[str]]] = None,
|
||||
fast_search: bool = False,
|
||||
fast_search: Optional[bool] = None,
|
||||
) -> LanceVectorQueryBuilder:
|
||||
"""Create a search query to find the nearest neighbors
|
||||
of the given query vector. We currently support
|
||||
@@ -735,11 +737,17 @@ class RemoteTable(Table):
|
||||
- If the table has multiple vector columns then the *vector_column_name*
|
||||
needs to be specified. Otherwise, an error is raised.
|
||||
|
||||
query_type: str, default "auto"
|
||||
The type of search: "auto", "vector", "fts", or "hybrid".
|
||||
With "auto", a string query uses the embedding function for the
|
||||
inferred or specified vector column if available, otherwise it uses
|
||||
full-text search.
|
||||
|
||||
fast_search: bool, optional
|
||||
Skip a flat search of unindexed data. This may improve
|
||||
search performance but search results will not include unindexed data.
|
||||
|
||||
- *default False*.
|
||||
- *default None*, which leaves fast search disabled.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -754,12 +762,27 @@ class RemoteTable(Table):
|
||||
- and also the "_distance" column which is the distance between the query
|
||||
vector and the returned vector.
|
||||
"""
|
||||
if isinstance(query, FullTextQuery):
|
||||
query_type = "fts"
|
||||
elif (
|
||||
query_type == "auto"
|
||||
and isinstance(query, str)
|
||||
and not self.embedding_functions
|
||||
):
|
||||
query_type = "fts"
|
||||
vector_column_name = infer_vector_column_name(
|
||||
schema=self.schema,
|
||||
query_type=query_type,
|
||||
query=query,
|
||||
vector_column_name=vector_column_name,
|
||||
)
|
||||
|
||||
return LanceQueryBuilder.create(
|
||||
self,
|
||||
query,
|
||||
query_type,
|
||||
vector_column_name=vector_column_name,
|
||||
fts_columns=fts_columns,
|
||||
fts_columns=fts_columns or [],
|
||||
fast_search=fast_search,
|
||||
)
|
||||
|
||||
@@ -934,26 +957,20 @@ class RemoteTable(Table):
|
||||
|
||||
def cleanup_old_versions(self, *_):
|
||||
"""
|
||||
cleanup_old_versions() is a no-op on LanceDB Cloud.
|
||||
|
||||
Tables are automatically cleaned up and optimized.
|
||||
cleanup_old_versions() is a no-op for remote tables.
|
||||
"""
|
||||
warnings.warn(
|
||||
"cleanup_old_versions() is a no-op on LanceDB Cloud. "
|
||||
"Tables are automatically cleaned up and optimized.",
|
||||
"cleanup_old_versions() is a no-op for remote tables.",
|
||||
stacklevel=2,
|
||||
)
|
||||
pass
|
||||
|
||||
def compact_files(self, *_):
|
||||
"""
|
||||
compact_files() is a no-op on LanceDB Cloud.
|
||||
|
||||
Tables are automatically compacted and optimized.
|
||||
compact_files() is a no-op for remote tables.
|
||||
"""
|
||||
warnings.warn(
|
||||
"compact_files() is a no-op on LanceDB Cloud. "
|
||||
"Tables are automatically compacted and optimized.",
|
||||
"compact_files() is a no-op for remote tables.",
|
||||
stacklevel=2,
|
||||
)
|
||||
pass
|
||||
@@ -965,17 +982,13 @@ class RemoteTable(Table):
|
||||
delete_unverified: bool = False,
|
||||
compaction_options: Optional[CompactionOptions] = None,
|
||||
):
|
||||
"""
|
||||
optimize() is a no-op on LanceDB Cloud.
|
||||
|
||||
Indices are optimized automatically.
|
||||
"""
|
||||
warnings.warn(
|
||||
"optimize() is a no-op on LanceDB Cloud. "
|
||||
"Indices are optimized automatically.",
|
||||
stacklevel=2,
|
||||
"""optimize() is not supported for remote tables."""
|
||||
return LOOP.run(
|
||||
self._table.optimize(
|
||||
cleanup_older_than=cleanup_older_than,
|
||||
delete_unverified=delete_unverified,
|
||||
)
|
||||
)
|
||||
pass
|
||||
|
||||
def count_rows(self, filter: Optional[str] = None) -> int:
|
||||
return LOOP.run(self._table.count_rows(filter))
|
||||
@@ -1023,20 +1036,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 on LanceDB Cloud."""
|
||||
"""Set the unenforced primary key for this table to a single column."""
|
||||
return LOOP.run(self._table.set_unenforced_primary_key(columns))
|
||||
|
||||
def set_lsm_write_spec(self, spec: "LsmWriteSpec") -> None:
|
||||
"""Install an LsmWriteSpec."""
|
||||
return LOOP.run(self._table.set_lsm_write_spec(spec))
|
||||
def enable_memwal(self, spec: "MemWalWriteSpec") -> None:
|
||||
"""Install an MemWalWriteSpec."""
|
||||
return LOOP.run(self._table.enable_memwal(spec))
|
||||
|
||||
def unset_lsm_write_spec(self) -> None:
|
||||
"""Remove the LsmWriteSpec."""
|
||||
return LOOP.run(self._table.unset_lsm_write_spec())
|
||||
def disable_memwal(self) -> None:
|
||||
"""Remove the MemWalWriteSpec."""
|
||||
return LOOP.run(self._table.disable_memwal())
|
||||
|
||||
def get_lsm_write_spec(self) -> Optional["LsmWriteSpec"]:
|
||||
"""Read the installed LsmWriteSpec, or ``None``."""
|
||||
return LOOP.run(self._table.get_lsm_write_spec())
|
||||
def get_memwal_spec(self) -> Optional["MemWalWriteSpec"]:
|
||||
"""Read the installed MemWalWriteSpec, or ``None``."""
|
||||
return LOOP.run(self._table.get_memwal_spec())
|
||||
|
||||
# WAL-PK-FUSION: delete both hooks.
|
||||
def _hybrid_pk_fusion_learned(self) -> bool:
|
||||
@@ -1046,31 +1059,31 @@ class RemoteTable(Table):
|
||||
def _note_hybrid_pk_fusion(self) -> None:
|
||||
self._table._note_hybrid_pk_fusion()
|
||||
|
||||
def checkpoint_lsm(self) -> None:
|
||||
def checkpoint_memwal(self) -> None:
|
||||
"""Synchronous version of
|
||||
[`AsyncTable.checkpoint_lsm`][lancedb.AsyncTable.checkpoint_lsm]."""
|
||||
return LOOP.run(self._table.checkpoint_lsm())
|
||||
[`AsyncTable.checkpoint_memwal`][lancedb.AsyncTable.checkpoint_memwal]."""
|
||||
return LOOP.run(self._table.checkpoint_memwal())
|
||||
|
||||
def flush_lsm(self) -> None:
|
||||
def flush_memwal(self) -> None:
|
||||
"""Synchronous version of
|
||||
[`AsyncTable.flush_lsm`][lancedb.AsyncTable.flush_lsm]."""
|
||||
return LOOP.run(self._table.flush_lsm())
|
||||
[`AsyncTable.flush_memwal`][lancedb.AsyncTable.flush_memwal]."""
|
||||
return LOOP.run(self._table.flush_memwal())
|
||||
|
||||
def compact_lsm(self) -> None:
|
||||
def compact_memwal(self) -> None:
|
||||
"""Synchronous version of
|
||||
[`AsyncTable.compact_lsm`][lancedb.AsyncTable.compact_lsm]."""
|
||||
return LOOP.run(self._table.compact_lsm())
|
||||
[`AsyncTable.compact_memwal`][lancedb.AsyncTable.compact_memwal]."""
|
||||
return LOOP.run(self._table.compact_memwal())
|
||||
|
||||
def get_lsm_stats(self, *, include_generation_rows: bool = False) -> Optional[dict]:
|
||||
def get_memwal_stats(self, *, include_sstable_rows: bool = False) -> Optional[dict]:
|
||||
"""Synchronous version of
|
||||
[`AsyncTable.get_lsm_stats`][lancedb.AsyncTable.get_lsm_stats]."""
|
||||
[`AsyncTable.get_memwal_stats`][lancedb.AsyncTable.get_memwal_stats]."""
|
||||
return LOOP.run(
|
||||
self._table.get_lsm_stats(include_generation_rows=include_generation_rows)
|
||||
self._table.get_memwal_stats(include_sstable_rows=include_sstable_rows)
|
||||
)
|
||||
|
||||
def close_lsm_writers(self) -> None:
|
||||
"""No-op on LanceDB Cloud (no local shard writers)."""
|
||||
return LOOP.run(self._table.close_lsm_writers())
|
||||
def close_memwal_writers(self) -> None:
|
||||
"""No-op for remote tables (no local shard writers)."""
|
||||
return LOOP.run(self._table.close_memwal_writers())
|
||||
|
||||
def drop_index(self, index_name: str):
|
||||
return LOOP.run(self._table.drop_index(index_name))
|
||||
@@ -1137,30 +1150,28 @@ class RemoteTable(Table):
|
||||
return LanceTakeQueryBuilder(self._table.take_row_ids(row_ids))
|
||||
|
||||
def uses_v2_manifest_paths(self) -> bool:
|
||||
raise NotImplementedError(
|
||||
"uses_v2_manifest_paths() is not supported on the LanceDB Cloud"
|
||||
)
|
||||
return LOOP.run(self._table.uses_v2_manifest_paths())
|
||||
|
||||
def migrate_v2_manifest_paths(self):
|
||||
raise NotImplementedError(
|
||||
"migrate_v2_manifest_paths() is not supported on the LanceDB Cloud"
|
||||
)
|
||||
return LOOP.run(self._table.migrate_manifest_paths_v2())
|
||||
|
||||
def blob_columns(self) -> list[str]:
|
||||
return LOOP.run(self._table.blob_columns())
|
||||
|
||||
def fetch_blobs(
|
||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
||||
self,
|
||||
column: str,
|
||||
row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table],
|
||||
) -> pa.LargeBinaryArray:
|
||||
return LOOP.run(self._table.fetch_blobs(column, row_ids))
|
||||
|
||||
def fetch_blob_ranges(self, column: str, requests) -> pa.LargeBinaryArray:
|
||||
raise NotImplementedError(
|
||||
"fetch_blob_ranges() is not supported on LanceDB Cloud"
|
||||
)
|
||||
return LOOP.run(self._table.fetch_blob_ranges(column, list(requests)))
|
||||
|
||||
def fetch_blob_files(
|
||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
||||
self,
|
||||
column: str,
|
||||
row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table],
|
||||
) -> "list[Optional[BlobFile]]":
|
||||
return LOOP.run(self._table.fetch_blob_files(column, row_ids))
|
||||
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
"""Schema helpers for Lance blob columns."""
|
||||
|
||||
import importlib
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import pyarrow as pa
|
||||
import pyarrow.ipc
|
||||
@@ -181,13 +181,60 @@ def _resolve_blob_type():
|
||||
return _resolved_blob_type
|
||||
|
||||
|
||||
def blob(name: str, nullable: bool = True) -> pa.Field:
|
||||
def blob(
|
||||
name: str,
|
||||
nullable: bool = True,
|
||||
*,
|
||||
inline_size_threshold: Optional[int] = None,
|
||||
dedicated_size_threshold: Optional[int] = None,
|
||||
pack_file_size_threshold: Optional[int] = None,
|
||||
) -> pa.Field:
|
||||
"""Create a Lance blob v2 column field.
|
||||
|
||||
When pylance is installed this is ``lance.blob.BlobType``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name:
|
||||
Field name.
|
||||
nullable:
|
||||
Whether the column accepts null values.
|
||||
inline_size_threshold:
|
||||
Maximum payload size in bytes to keep inline in the data file
|
||||
before using packed blob storage. Zero is allowed.
|
||||
dedicated_size_threshold:
|
||||
Maximum payload size in bytes to store in packed blob storage
|
||||
before using dedicated blob storage. Must be positive.
|
||||
pack_file_size_threshold:
|
||||
Maximum size in bytes of a single packed blob sidecar file before
|
||||
a new one is started. Must be positive.
|
||||
|
||||
The thresholds are forwarded to ``lance.blob_field``, which validates
|
||||
them and stores them as field metadata, so pylance is required whenever
|
||||
any threshold is set. When all thresholds are omitted the behavior is
|
||||
unchanged.
|
||||
"""
|
||||
blob_type = _resolve_blob_type()
|
||||
return pa.field(name, blob_type(), nullable=nullable)
|
||||
if (
|
||||
inline_size_threshold is None
|
||||
and dedicated_size_threshold is None
|
||||
and pack_file_size_threshold is None
|
||||
):
|
||||
blob_type = _resolve_blob_type()
|
||||
return pa.field(name, blob_type(), nullable=nullable)
|
||||
try:
|
||||
lance = importlib.import_module("lance")
|
||||
except ModuleNotFoundError as err:
|
||||
raise ValueError(
|
||||
"Setting blob tier thresholds requires pylance to be installed"
|
||||
) from err
|
||||
_resolve_blob_type()
|
||||
return lance.blob_field(
|
||||
name,
|
||||
nullable=nullable,
|
||||
inline_size_threshold=inline_size_threshold,
|
||||
dedicated_size_threshold=dedicated_size_threshold,
|
||||
pack_file_size_threshold=pack_file_size_threshold,
|
||||
)
|
||||
|
||||
|
||||
def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataType:
|
||||
|
||||
@@ -3,8 +3,10 @@
|
||||
|
||||
"""Handles to SQL queries running on a remote database."""
|
||||
|
||||
from typing import Any, Mapping, Optional, Sequence, Union
|
||||
from uuid import UUID
|
||||
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
|
||||
from lancedb.background_loop import LOOP
|
||||
@@ -14,6 +16,82 @@ from .arrow import AsyncRecordBatchReader
|
||||
|
||||
QueryDescription = _lancedb.QueryDescription
|
||||
|
||||
QueryParameters = Union[pa.RecordBatch, pa.Table, Sequence[Any], Mapping[str, Any]]
|
||||
"""Values for a statement's ``$1`` / ``$name`` placeholders.
|
||||
|
||||
A sequence binds by position (``$1``, ``$2``, ...), a mapping binds by name
|
||||
(``$name``), and a one-row ``pyarrow.RecordBatch`` or ``pyarrow.Table`` binds
|
||||
both ways: column ``i`` is ``$<i + 1>`` and, when its name is not a number,
|
||||
also ``$<name>``.
|
||||
"""
|
||||
|
||||
|
||||
def _parameter_array(value: Any) -> pa.Array:
|
||||
"""One parameter as a one-element array of its own type."""
|
||||
if isinstance(value, pa.Scalar):
|
||||
return pa.array([value.as_py()], type=value.type)
|
||||
# A 1-D array is a vector: keep its element type and make it the
|
||||
# fixed-size list a vector column stores, rather than a list of Python
|
||||
# floats that would come back as float64.
|
||||
if isinstance(value, pa.ChunkedArray):
|
||||
value = value.combine_chunks()
|
||||
if isinstance(value, np.ndarray):
|
||||
if value.ndim == 0:
|
||||
return pa.array(value.reshape(1))
|
||||
if value.ndim != 1:
|
||||
raise ValueError(
|
||||
"a numpy query parameter must be a scalar or a 1-D vector, "
|
||||
f"got {value.ndim} dimensions"
|
||||
)
|
||||
value = pa.array(value)
|
||||
if isinstance(value, pa.Array):
|
||||
return pa.FixedSizeListArray.from_arrays(value, len(value))
|
||||
return pa.array([value])
|
||||
|
||||
|
||||
def to_parameter_batch(
|
||||
parameters: Optional[QueryParameters],
|
||||
) -> Optional[pa.RecordBatch]:
|
||||
"""Turn ``parameters`` into the one-row batch a parameterized query sends.
|
||||
|
||||
Scalars keep the type pyarrow infers for them (numpy scalars keep their
|
||||
dtype), a ``pyarrow.Scalar`` keeps its own type, and a 1-D numpy or
|
||||
pyarrow array becomes a ``FixedSizeList`` of its element type, which is
|
||||
what a vector column stores.
|
||||
"""
|
||||
if parameters is None:
|
||||
return None
|
||||
batch: pa.RecordBatch
|
||||
if isinstance(parameters, pa.RecordBatch):
|
||||
batch = parameters
|
||||
elif isinstance(parameters, pa.Table):
|
||||
batches = parameters.combine_chunks().to_batches()
|
||||
batch = (
|
||||
batches[0]
|
||||
if len(batches) == 1
|
||||
else pa.RecordBatch.from_pylist([], schema=parameters.schema)
|
||||
)
|
||||
elif isinstance(parameters, Mapping):
|
||||
batch = pa.RecordBatch.from_arrays(
|
||||
[_parameter_array(value) for value in parameters.values()],
|
||||
names=[str(name) for name in parameters.keys()],
|
||||
)
|
||||
elif isinstance(parameters, Sequence) and not isinstance(parameters, (str, bytes)):
|
||||
batch = pa.RecordBatch.from_arrays(
|
||||
[_parameter_array(value) for value in parameters],
|
||||
names=[str(position + 1) for position in range(len(parameters))],
|
||||
)
|
||||
else:
|
||||
raise TypeError(
|
||||
"query parameters must be a sequence, a mapping, or a one-row "
|
||||
f"pyarrow RecordBatch or Table, not {type(parameters).__name__}"
|
||||
)
|
||||
if batch.num_rows != 1:
|
||||
raise ValueError(
|
||||
f"query parameters must be exactly one row, got {batch.num_rows}"
|
||||
)
|
||||
return batch
|
||||
|
||||
|
||||
class AsyncQuery:
|
||||
"""A handle to a submitted SQL query on an asynchronous connection."""
|
||||
@@ -85,4 +163,4 @@ class Query:
|
||||
LOOP.run(self._inner.cancel())
|
||||
|
||||
|
||||
__all__ = ["AsyncQuery", "Query", "QueryDescription"]
|
||||
__all__ = ["AsyncQuery", "Query", "QueryDescription", "QueryParameters"]
|
||||
|
||||
+459
-150
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
@@ -235,6 +236,19 @@ def inf_vector_column_query(schema: pa.Schema, dim: Optional[int] = None) -> str
|
||||
vector_col_names = []
|
||||
for field in schema:
|
||||
vector_col_names.extend(_iter_vector_columns(field, [], dim))
|
||||
|
||||
if len(vector_col_names) == 0 and dim is not None:
|
||||
# The query's dimension matched no column. If the schema is
|
||||
# unambiguous without that filter, the single vector column is still
|
||||
# the right one to search: a dimension mismatch is a query-execution
|
||||
# error, not a missing-column error.
|
||||
unfiltered = []
|
||||
for field in schema:
|
||||
unfiltered.extend(_iter_vector_columns(field, []))
|
||||
if len(unfiltered) == 1:
|
||||
return unfiltered[0]
|
||||
vector_col_names = unfiltered
|
||||
|
||||
if len(vector_col_names) > 1:
|
||||
raise ValueError(
|
||||
"Schema has more than one vector column. "
|
||||
@@ -280,6 +294,14 @@ def infer_vector_column_dim(data_type: pa.DataType) -> Optional[int]:
|
||||
return None
|
||||
|
||||
|
||||
def _validate_query_vector(query: Any) -> None:
|
||||
"""Reject empty vector inputs before vector-column inference or execution."""
|
||||
if (isinstance(query, list) and not query) or (
|
||||
isinstance(query, np.ndarray) and query.size == 0
|
||||
):
|
||||
raise ValueError("Query vector must not be empty")
|
||||
|
||||
|
||||
def _query_vector_dim(query: Optional[Any]) -> Optional[int]:
|
||||
if query is None:
|
||||
return None
|
||||
@@ -301,6 +323,9 @@ def infer_vector_column_name(
|
||||
query: Optional[Any], # inferred later in query builder
|
||||
vector_column_name: Optional[str],
|
||||
):
|
||||
if query_type != "fts":
|
||||
_validate_query_vector(query)
|
||||
|
||||
if vector_column_name is not None:
|
||||
return vector_column_name
|
||||
|
||||
@@ -362,6 +387,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()
|
||||
|
||||
@@ -0,0 +1,821 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import asyncio
|
||||
from dataclasses import FrozenInstanceError
|
||||
import json
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
import lancedb
|
||||
from lancedb.authz import (
|
||||
AccessControlEntry,
|
||||
AsyncAuthorization,
|
||||
Authorization,
|
||||
GroupInfo,
|
||||
IdentityInfo,
|
||||
Object,
|
||||
PrincipalInfo,
|
||||
Privilege,
|
||||
RoleBinding,
|
||||
Subject,
|
||||
)
|
||||
from lancedb.remote import HeaderProvider
|
||||
from lancedb.remote.errors import HttpError
|
||||
from test_catalog import catalog_server as _catalog_server
|
||||
|
||||
catalog_server = _catalog_server
|
||||
|
||||
|
||||
class DynamicHeaders(HeaderProvider):
|
||||
def get_headers(self):
|
||||
return {
|
||||
"Authorization": "Bearer rotating-token",
|
||||
"X-LanceDB-Database": "wrong-dynamic-db",
|
||||
"X-LanceDB-Database-Prefix": "wrong-prefix",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
def test_authz_operations_and_scope(catalog_server, asynchronous):
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.extend(
|
||||
[
|
||||
(200, {}),
|
||||
(200, {}),
|
||||
(200, {}),
|
||||
(200, {}),
|
||||
(
|
||||
200,
|
||||
{
|
||||
"entries": [
|
||||
{
|
||||
"subject_id": "g:team-id",
|
||||
"object": "table:tenant/db:a$b$events",
|
||||
"privilege": "PRIVILEGE20",
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(
|
||||
200,
|
||||
{
|
||||
"bindings": [
|
||||
{
|
||||
"subject_id": "p:alice-id",
|
||||
"subject_name": "Alice",
|
||||
"role": "reader",
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(
|
||||
200,
|
||||
{
|
||||
"principal_id": "alice-id",
|
||||
"groups": [{"id": "team-id"}],
|
||||
"roles": ["reader"],
|
||||
},
|
||||
),
|
||||
(
|
||||
200,
|
||||
{
|
||||
"group_id": "team-id",
|
||||
"group_name": "Engineering",
|
||||
"principals": [{"id": "alice-id", "name": "Alice"}],
|
||||
"roles": ["reader"],
|
||||
},
|
||||
),
|
||||
]
|
||||
)
|
||||
config = {
|
||||
"header_provider": DynamicHeaders(),
|
||||
"extra_headers": {"x-lancedb-database": "wrong-static-db"},
|
||||
}
|
||||
obj = Object.table(database="tenant/db", namespace_path=["a", "b"], name="events")
|
||||
|
||||
async def exercise():
|
||||
if asynchronous:
|
||||
catalog = await lancedb.connect_catalog_async(
|
||||
endpoint, client_config=config
|
||||
)
|
||||
else:
|
||||
catalog = lancedb.connect_catalog(endpoint, client_config=config)
|
||||
authz = catalog.authz
|
||||
assert authz is catalog.authz
|
||||
assert requests == [] # Constructing handles does not probe a database.
|
||||
|
||||
async def call(method, *args, **kwargs):
|
||||
value = method(*args, **kwargs)
|
||||
return await value if asynchronous else value
|
||||
|
||||
await call(
|
||||
authz.add_acl,
|
||||
object=obj,
|
||||
subject=Subject.group_name("Engineering"),
|
||||
privilege=Privilege.SELECT,
|
||||
)
|
||||
await call(
|
||||
authz.delete_acl,
|
||||
object=obj,
|
||||
subject=Subject.group_id("team-id"),
|
||||
privilege="SELECT",
|
||||
)
|
||||
await call(
|
||||
authz.add_role_binding,
|
||||
subject=Subject.principal_name("Alice"),
|
||||
role="reader",
|
||||
)
|
||||
await call(
|
||||
authz.delete_role_binding,
|
||||
subject=Subject.principal_api_key("subject-secret"),
|
||||
role="reader",
|
||||
)
|
||||
iterator = authz.list_acls(object=obj, subject=Subject.role("reader"))
|
||||
entries = (
|
||||
[entry async for entry in iterator] if asynchronous else list(iterator)
|
||||
)
|
||||
assert entries == [
|
||||
AccessControlEntry(
|
||||
str(obj), Subject.group_id("team-id"), None, "PRIVILEGE20"
|
||||
)
|
||||
]
|
||||
iterator = authz.list_role_bindings(
|
||||
subject=Subject.group_id("team-id"), role="reader"
|
||||
)
|
||||
bindings = (
|
||||
[binding async for binding in iterator] if asynchronous else list(iterator)
|
||||
)
|
||||
assert bindings == [
|
||||
RoleBinding(Subject.principal_id("alice-id"), "Alice", "reader")
|
||||
]
|
||||
assert await call(
|
||||
authz.get_principal, Subject.principal_id("alice-id")
|
||||
) == PrincipalInfo("alice-id", None, [IdentityInfo("team-id")], ["reader"])
|
||||
assert await call(
|
||||
authz.get_group, Subject.group_name("Engineering")
|
||||
) == GroupInfo(
|
||||
"team-id", "Engineering", [IdentityInfo("alice-id", "Alice")], ["reader"]
|
||||
)
|
||||
|
||||
asyncio.run(exercise())
|
||||
assert [path for path, _, _ in requests] == [
|
||||
"/admin/authz/acl/add",
|
||||
"/admin/authz/acl/delete",
|
||||
"/admin/authz/role_binding/add",
|
||||
"/admin/authz/role_binding/delete",
|
||||
"/admin/authz/acl/list",
|
||||
"/admin/authz/role_binding/list",
|
||||
"/admin/authz/principal",
|
||||
"/admin/authz/group",
|
||||
]
|
||||
assert [body for _, _, body in requests] == [
|
||||
{"object": str(obj), "subject": "G:Engineering", "privilege": "SELECT"},
|
||||
{"object": str(obj), "subject": "g:team-id", "privilege": "SELECT"},
|
||||
{"subject": "P:Alice", "role": "reader"},
|
||||
{"subject": "a:subject-secret", "role": "reader"},
|
||||
{"object": str(obj), "subject": "r:reader", "max_results": 100},
|
||||
{"subject": "g:team-id", "role": "reader", "max_results": 100},
|
||||
{"principal": "p:alice-id"},
|
||||
{"group": "G:Engineering"},
|
||||
]
|
||||
for _, headers, _ in requests:
|
||||
headers = {key.lower(): value for key, value in headers.items()}
|
||||
assert "x-lancedb-database" not in headers
|
||||
assert "x-lancedb-database-prefix" not in headers
|
||||
assert headers["authorization"] == "Bearer rotating-token"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
@pytest.mark.parametrize("limit", [None, 0, 1, 3])
|
||||
def test_lazy_expanded_acl_pagination(catalog_server, asynchronous, limit):
|
||||
endpoint, requests, responses = catalog_server
|
||||
|
||||
def entry(privilege):
|
||||
return {
|
||||
"object": "database:db",
|
||||
"subject_id": "r:reader",
|
||||
"privilege": privilege,
|
||||
}
|
||||
|
||||
responses.extend(
|
||||
[
|
||||
(
|
||||
200,
|
||||
{
|
||||
"entries": [entry("SELECT"), entry("INSERT")],
|
||||
"continuation_token": "opaque/first",
|
||||
},
|
||||
),
|
||||
(200, {"entries": [], "continuation_token": "opaque/second"}),
|
||||
(200, {"entries": [entry("UPDATE"), entry("DELETE")]}),
|
||||
]
|
||||
)
|
||||
|
||||
async def exercise():
|
||||
catalog = (
|
||||
await lancedb.connect_catalog_async(endpoint, api_key="caller-secret")
|
||||
if asynchronous
|
||||
else lancedb.connect_catalog(endpoint, api_key="caller-secret")
|
||||
)
|
||||
iterator = catalog.authz.list_acls(subject=Subject.role("reader"), limit=limit)
|
||||
assert requests == []
|
||||
return [entry async for entry in iterator] if asynchronous else list(iterator)
|
||||
|
||||
values = asyncio.run(exercise())
|
||||
expected = ["SELECT", "INSERT", "UPDATE", "DELETE"]
|
||||
assert [entry.privilege for entry in values] == (
|
||||
expected if limit is None else expected[:limit]
|
||||
)
|
||||
assert len(requests) == (0 if limit == 0 else 1 if limit == 1 else 3)
|
||||
for _, headers, body in requests:
|
||||
assert headers["x-api-key"] == "caller-secret"
|
||||
assert body["subject"] == "r:reader"
|
||||
if len(requests) == 3:
|
||||
assert requests[1][2]["continuation_token"] == "opaque/first"
|
||||
assert requests[2][2]["continuation_token"] == "opaque/second"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", [400, 401, 403, 404, 409, 501, 503])
|
||||
def test_http_errors_preserve_status_without_replaying_writes(catalog_server, status):
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.append((status, {"error": "authorization failure"}))
|
||||
authz = lancedb.connect_catalog(endpoint).authz
|
||||
with pytest.raises(HttpError) as error:
|
||||
authz.add_acl(
|
||||
object="system", subject=Subject.role("operator"), privilege="OPERATE"
|
||||
)
|
||||
assert error.value.status_code == status
|
||||
assert error.value.request_id
|
||||
assert len(requests) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
def test_role_pagination_and_late_failure(catalog_server, asynchronous):
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.extend(
|
||||
[
|
||||
(
|
||||
200,
|
||||
{
|
||||
"bindings": [{"subject_id": "g:team", "role": "reader"}],
|
||||
"continuation_token": "next",
|
||||
},
|
||||
),
|
||||
(403, {"error": "denied"}),
|
||||
]
|
||||
)
|
||||
|
||||
async def exercise():
|
||||
catalog = (
|
||||
await lancedb.connect_catalog_async(endpoint)
|
||||
if asynchronous
|
||||
else lancedb.connect_catalog(endpoint)
|
||||
)
|
||||
iterator = catalog.authz.list_role_bindings(role="reader")
|
||||
first = await anext(iterator) if asynchronous else next(iterator)
|
||||
assert first.subject == Subject.group_id("team")
|
||||
assert len(requests) == 1
|
||||
with pytest.raises(HttpError) as error:
|
||||
if asynchronous:
|
||||
await anext(iterator)
|
||||
else:
|
||||
next(iterator)
|
||||
assert error.value.status_code == 403
|
||||
|
||||
asyncio.run(exercise())
|
||||
assert requests[1][2] == {
|
||||
"role": "reader",
|
||||
"max_results": 100,
|
||||
"continuation_token": "next",
|
||||
}
|
||||
|
||||
|
||||
def test_early_close_does_not_fetch_next_page(catalog_server):
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.append(
|
||||
(
|
||||
200,
|
||||
{
|
||||
"bindings": [{"subject_id": "g:team", "role": "reader"}],
|
||||
"continuation_token": "next",
|
||||
},
|
||||
)
|
||||
)
|
||||
iterator = lancedb.connect_catalog(endpoint).authz.list_role_bindings()
|
||||
next(iterator)
|
||||
iterator.close()
|
||||
assert len(requests) == 1
|
||||
|
||||
|
||||
def test_selectors_objects_and_redaction():
|
||||
secret = Subject.principal_api_key("do-not-print-this")
|
||||
assert "do-not-print-this" not in repr(secret)
|
||||
assert "do-not-print-this" not in str(secret)
|
||||
with pytest.raises(FrozenInstanceError):
|
||||
secret.kind = "role"
|
||||
assert str(Object.system()) == "system"
|
||||
assert str(Object.database("tenant/db")) == "database:tenant/db"
|
||||
assert (
|
||||
str(Object.namespace(database="db", namespace_path=["a", "b"]))
|
||||
== "namespace:db:a$b"
|
||||
)
|
||||
assert str(Object.table(database="db", name="events")) == "table:db:public$events"
|
||||
assert str(Object.view(database="db", name="recent")) == "view:db:public$recent"
|
||||
assert Subject.principal_id("p:id")._wire() == "p:p:id"
|
||||
assert Subject.group_id("team-id").value == "team-id"
|
||||
for namespace in (["a$b"], [""], "public", [".."]):
|
||||
with pytest.raises(ValueError):
|
||||
Object.table(database="db", namespace_path=namespace, name="events")
|
||||
for database in ("", "a//b", "a/../b", "a:b", "a%2Fb", "a$b"):
|
||||
with pytest.raises(ValueError):
|
||||
Object.database(database)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validation_before_io_and_empty_lists():
|
||||
inner = AsyncMock()
|
||||
inner.list_acls.return_value = json.dumps({"entries": []})
|
||||
inner.list_role_bindings.return_value = json.dumps({"bindings": []})
|
||||
authz = AsyncAuthorization(inner)
|
||||
for limit in (-1, 1.5, True):
|
||||
with pytest.raises(ValueError):
|
||||
authz.list_acls(limit=limit)
|
||||
with pytest.raises(ValueError):
|
||||
Authorization(authz).list_role_bindings(limit=limit)
|
||||
with pytest.raises(ValueError):
|
||||
await authz.get_group(Subject.principal_id("alice"))
|
||||
with pytest.raises(ValueError):
|
||||
await authz.get_principal(Subject.group_id("team"))
|
||||
for method in (authz.add_role_binding, authz.delete_role_binding):
|
||||
with pytest.raises(ValueError):
|
||||
await method(subject=Subject.role("reader"), role="writer")
|
||||
with pytest.raises(ValueError):
|
||||
authz.list_role_bindings(subject=Subject.role("reader"))
|
||||
assert [entry async for entry in authz.list_acls(limit=0)] == []
|
||||
assert inner.mock_calls == []
|
||||
assert [entry async for entry in authz.list_acls()] == []
|
||||
assert [entry async for entry in authz.list_role_bindings()] == []
|
||||
|
||||
|
||||
def test_api_key_subject_is_not_logged(catalog_server):
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.append((200, {"principal_id": "public-id", "groups": [], "roles": []}))
|
||||
result = subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-c",
|
||||
"import lancedb, sys; from lancedb.authz import Subject; "
|
||||
"lancedb.connect_catalog(sys.argv[1]).authz.get_principal("
|
||||
"Subject.principal_api_key('test-subject-secret-never-log'))",
|
||||
endpoint,
|
||||
],
|
||||
env={**os.environ, "LANCEDB_LOG": "debug"},
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
timeout=30,
|
||||
)
|
||||
assert requests[0][2] == {"principal": "a:test-subject-secret-never-log"}
|
||||
assert "test-subject-secret-never-log" not in result.stdout + result.stderr
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"builder, kwargs, wire",
|
||||
[
|
||||
("system", {}, "system"),
|
||||
("secret", {"database": "db", "name": "api-key"}, "secret:db:public$api-key"),
|
||||
(
|
||||
"function",
|
||||
{"database": "db", "name": "caption"},
|
||||
"function:db:public$caption",
|
||||
),
|
||||
(
|
||||
"secret",
|
||||
{"database": "tenant/db", "namespace_path": ["a", "b"], "name": "api-key"},
|
||||
"secret:tenant/db:a$b$api-key",
|
||||
),
|
||||
(
|
||||
"function",
|
||||
{"database": "tenant/db", "namespace_path": ["a", "b"], "name": "caption"},
|
||||
"function:tenant/db:a$b$caption",
|
||||
),
|
||||
("database", {"name": "tenant/db"}, "database:tenant/db"),
|
||||
(
|
||||
"namespace",
|
||||
{"database": "tenant/db", "namespace_path": ["a", "b"]},
|
||||
"namespace:tenant/db:a$b",
|
||||
),
|
||||
(
|
||||
"table",
|
||||
{"database": "tenant/db", "namespace_path": ["a", "b"], "name": "events"},
|
||||
"table:tenant/db:a$b$events",
|
||||
),
|
||||
(
|
||||
"view",
|
||||
{"database": "db", "namespace_path": ["public"], "name": "recent"},
|
||||
"view:db:public$recent",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_rust_backed_object_builders(builder, kwargs, wire):
|
||||
import pickle
|
||||
from lancedb import _lancedb
|
||||
|
||||
value = getattr(Object, builder)(**kwargs)
|
||||
assert str(value) == wire
|
||||
assert str(value._inner) == wire
|
||||
assert str(_lancedb.AuthzObject(wire)) == wire
|
||||
assert value == Object(wire)
|
||||
assert hash(value) == hash(Object(wire))
|
||||
assert pickle.loads(pickle.dumps(value)) == value
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kind, prefix",
|
||||
[
|
||||
("principal_id", "p:"),
|
||||
("principal_name", "P:"),
|
||||
("group_id", "g:"),
|
||||
("group_name", "G:"),
|
||||
("role", "r:"),
|
||||
("principal_api_key", "a:"),
|
||||
],
|
||||
)
|
||||
def test_rust_backed_subject_selectors(kind, prefix):
|
||||
import pickle
|
||||
from lancedb import _lancedb
|
||||
|
||||
value = "p:literal@example.com / 名前"
|
||||
subject = getattr(Subject, kind)(value)
|
||||
assert subject.kind == kind
|
||||
assert subject.value == value
|
||||
assert subject._wire() == prefix + value
|
||||
assert _lancedb.AuthzSubject(kind, value).to_wire() == prefix + value
|
||||
assert pickle.loads(pickle.dumps(subject)) == subject
|
||||
assert hash(subject) == hash(getattr(Subject, kind)(value))
|
||||
with pytest.raises(ValueError):
|
||||
getattr(Subject, kind)("")
|
||||
if kind == "principal_api_key":
|
||||
assert value not in repr(subject._inner)
|
||||
assert value not in str(subject._inner)
|
||||
if kind in ("principal_id", "group_id", "role"):
|
||||
assert Subject._canonical(prefix + value) == subject
|
||||
else:
|
||||
with pytest.raises(ValueError) as error:
|
||||
Subject._canonical(prefix + value)
|
||||
assert value not in str(error.value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"wire",
|
||||
[
|
||||
"",
|
||||
"database:",
|
||||
"database:a//b",
|
||||
"database:a/../b",
|
||||
"database:a%2Fb",
|
||||
"namespace:db:a$$b",
|
||||
"table:db:events",
|
||||
"table:db:public$",
|
||||
"table:db:public$..",
|
||||
"table:db:a/b$events",
|
||||
"view:db:public$café",
|
||||
"secret:db:public$",
|
||||
"function:db:public$..",
|
||||
"secret:db:a$$key",
|
||||
"function:db:caption",
|
||||
"secret:db:public$bad/name",
|
||||
"function:db:public$bad:name",
|
||||
"system:db",
|
||||
],
|
||||
)
|
||||
def test_invalid_raw_objects_fail_before_io(catalog_server, wire):
|
||||
endpoint, requests, _ = catalog_server
|
||||
authz = lancedb.connect_catalog(endpoint).authz
|
||||
with pytest.raises(ValueError):
|
||||
Object(wire)
|
||||
with pytest.raises(ValueError):
|
||||
authz.list_acls(object=wire)
|
||||
with pytest.raises(ValueError):
|
||||
authz.add_acl(object=wire, subject=Subject.role("reader"), privilege="SELECT")
|
||||
assert requests == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"wire",
|
||||
[
|
||||
"future:db:resource",
|
||||
"table_extension:db:resource",
|
||||
"secret:tenant/db:a$b$key",
|
||||
"function:tenant/db:a$b$caption",
|
||||
],
|
||||
)
|
||||
def test_objects_round_trip_through_acl_requests(catalog_server, wire):
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.extend(
|
||||
[
|
||||
(200, {}),
|
||||
(
|
||||
200,
|
||||
{
|
||||
"entries": [
|
||||
{
|
||||
"object": wire,
|
||||
"subject_id": "r:reader",
|
||||
"privilege": "SELECT",
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(200, {}),
|
||||
]
|
||||
)
|
||||
authz = lancedb.connect_catalog(endpoint).authz
|
||||
obj = Object(wire)
|
||||
assert str(obj) == wire
|
||||
authz.add_acl(object=obj, subject=Subject.role("reader"), privilege="SELECT")
|
||||
entries = list(authz.list_acls(object=wire))
|
||||
assert entries[0].object == wire
|
||||
authz.delete_acl(
|
||||
object=entries[0].object, subject=Subject.role("reader"), privilege="SELECT"
|
||||
)
|
||||
assert len(requests) == 3
|
||||
for _, _, body in requests:
|
||||
assert body["object"] == wire
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_requests_validate_selector_context_before_io(catalog_server):
|
||||
endpoint, requests, _ = catalog_server
|
||||
catalog = await lancedb.connect_catalog_async(endpoint)
|
||||
native = catalog.authz._inner
|
||||
for method, request in [
|
||||
(native.list_role_bindings, {"subject": "r:reader"}),
|
||||
(native.add_role_binding, {"subject": "r:reader", "role": "writer"}),
|
||||
(native.delete_role_binding, {"subject": "r:reader", "role": "writer"}),
|
||||
(native.get_principal, {"principal": "g:group"}),
|
||||
(native.get_group, {"group": "a:never-log-this-key"}),
|
||||
]:
|
||||
with pytest.raises(ValueError) as error:
|
||||
await method(json.dumps(request))
|
||||
assert "never-log-this-key" not in str(error.value)
|
||||
assert requests == []
|
||||
|
||||
|
||||
def test_privilege_values_use_rust_serialization():
|
||||
import pickle
|
||||
from lancedb import _lancedb
|
||||
|
||||
names = [value.value for value in Privilege]
|
||||
assert names == _lancedb.AuthzPrivilege.known()
|
||||
for privilege in Privilege:
|
||||
assert json.loads(json.dumps(privilege)) == privilege.value
|
||||
assert str(_lancedb.AuthzPrivilege(privilege)) == privilege.value
|
||||
assert pickle.loads(pickle.dumps(privilege)) is privilege
|
||||
for index in range(64):
|
||||
alias = f"PRIVILEGE{index}"
|
||||
expected = names[index] if index < len(names) else alias
|
||||
assert str(_lancedb.AuthzPrivilege(alias)) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"creation_privilege", [None, "CREATE_SECRET", "CREATE_FUNCTION"]
|
||||
)
|
||||
def test_privilege_names_and_future_values_round_trip(
|
||||
catalog_server, asynchronous, creation_privilege
|
||||
):
|
||||
endpoint, requests, responses = catalog_server
|
||||
inputs = [
|
||||
Privilege.SELECT,
|
||||
"SELECT",
|
||||
"PRIVILEGE2",
|
||||
"PRIVILEGE20",
|
||||
"PRIVILEGE63",
|
||||
"PRIVILEGE64",
|
||||
"PRIVILEGE256",
|
||||
"PRIVILEGE18446744073709551616",
|
||||
]
|
||||
expected = ["SELECT", "SELECT", "SELECT", *inputs[3:]]
|
||||
object_id = "table:db:public$t"
|
||||
if creation_privilege is not None:
|
||||
inputs = [Privilege[creation_privilege], creation_privilege]
|
||||
expected = [creation_privilege, creation_privilege]
|
||||
object_id = "namespace:db:public"
|
||||
for privilege in inputs:
|
||||
responses.extend(
|
||||
[
|
||||
(200, {}),
|
||||
(
|
||||
200,
|
||||
{
|
||||
"entries": [
|
||||
{
|
||||
"object": object_id,
|
||||
"subject_id": "r:reader",
|
||||
"privilege": privilege,
|
||||
}
|
||||
]
|
||||
},
|
||||
),
|
||||
(200, {}),
|
||||
]
|
||||
)
|
||||
|
||||
async def exercise():
|
||||
catalog = (
|
||||
await lancedb.connect_catalog_async(endpoint)
|
||||
if asynchronous
|
||||
else lancedb.connect_catalog(endpoint)
|
||||
)
|
||||
authz = catalog.authz
|
||||
for privilege, canonical in zip(inputs, expected):
|
||||
added = authz.add_acl(
|
||||
object=object_id,
|
||||
subject=Subject.role("reader"),
|
||||
privilege=privilege,
|
||||
)
|
||||
if asynchronous:
|
||||
await added
|
||||
iterator = authz.list_acls()
|
||||
entries = (
|
||||
[entry async for entry in iterator] if asynchronous else list(iterator)
|
||||
)
|
||||
assert entries[0].privilege == canonical
|
||||
deleted = authz.delete_acl(
|
||||
object=entries[0].object,
|
||||
subject=entries[0].subject,
|
||||
privilege=entries[0].privilege,
|
||||
)
|
||||
if asynchronous:
|
||||
await deleted
|
||||
|
||||
asyncio.run(exercise())
|
||||
assert len(requests) == 3 * len(inputs)
|
||||
for index, canonical in enumerate(expected):
|
||||
assert requests[index * 3][2]["object"] == object_id
|
||||
assert requests[index * 3][2]["privilege"] == canonical
|
||||
assert requests[index * 3 + 2][2]["object"] == object_id
|
||||
assert requests[index * 3 + 2][2]["privilege"] == canonical
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
def test_invalid_privileges_fail_before_io(catalog_server, asynchronous):
|
||||
endpoint, requests, _ = catalog_server
|
||||
invalid = [
|
||||
"",
|
||||
"select",
|
||||
"SELECT ",
|
||||
"UNKNOWN",
|
||||
"PRIVILEGE",
|
||||
"PRIVILEGE-1",
|
||||
"PRIVILEGE+1",
|
||||
"PRIVILEGE1x",
|
||||
None,
|
||||
2,
|
||||
True,
|
||||
["SELECT"],
|
||||
]
|
||||
|
||||
async def exercise():
|
||||
catalog = (
|
||||
await lancedb.connect_catalog_async(endpoint)
|
||||
if asynchronous
|
||||
else lancedb.connect_catalog(endpoint)
|
||||
)
|
||||
async_authz = catalog.authz if asynchronous else catalog.authz._inner
|
||||
native = async_authz._inner
|
||||
for privilege in invalid:
|
||||
for method in [catalog.authz.add_acl, catalog.authz.delete_acl]:
|
||||
with pytest.raises(ValueError, match="privilege"):
|
||||
result = method(
|
||||
object="system",
|
||||
subject=Subject.role("reader"),
|
||||
privilege=privilege,
|
||||
)
|
||||
if asynchronous:
|
||||
await result
|
||||
for method in [
|
||||
native.add_acl,
|
||||
native.delete_acl,
|
||||
]:
|
||||
with pytest.raises(ValueError):
|
||||
await method(
|
||||
json.dumps(
|
||||
{
|
||||
"object": "system",
|
||||
"subject": "r:reader",
|
||||
"privilege": privilege,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
asyncio.run(exercise())
|
||||
assert requests == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
def test_acl_page_preserves_unknown_subjects(catalog_server, asynchronous):
|
||||
import pickle
|
||||
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.extend(
|
||||
[
|
||||
(
|
||||
200,
|
||||
{
|
||||
"entries": [
|
||||
{
|
||||
"subject_id": "p:known",
|
||||
"object": "system",
|
||||
"privilege": "USAGE",
|
||||
},
|
||||
{
|
||||
"subject_id": "x:future-subject",
|
||||
"object": "system",
|
||||
"privilege": "USAGE",
|
||||
},
|
||||
]
|
||||
},
|
||||
),
|
||||
(200, {}),
|
||||
]
|
||||
)
|
||||
|
||||
async def exercise():
|
||||
catalog = (
|
||||
await lancedb.connect_catalog_async(endpoint)
|
||||
if asynchronous
|
||||
else lancedb.connect_catalog(endpoint)
|
||||
)
|
||||
iterator = catalog.authz.list_acls()
|
||||
entries = (
|
||||
[entry async for entry in iterator] if asynchronous else list(iterator)
|
||||
)
|
||||
assert len(entries) == 2
|
||||
assert entries[0].subject == Subject.principal_id("known")
|
||||
subject = entries[1].subject
|
||||
assert subject.kind == "unknown"
|
||||
assert subject.value == "x:future-subject"
|
||||
assert subject._wire() == "x:future-subject"
|
||||
assert pickle.loads(pickle.dumps(subject)) == subject
|
||||
result = catalog.authz.delete_acl(
|
||||
object=entries[1].object, subject=subject, privilege=entries[1].privilege
|
||||
)
|
||||
if asynchronous:
|
||||
await result
|
||||
|
||||
asyncio.run(exercise())
|
||||
assert len(requests) == 2
|
||||
assert requests[1][2]["subject"] == "x:future-subject"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("builder", ["secret", "function"])
|
||||
def test_secret_and_function_builders_validate_components(builder):
|
||||
from lancedb import _lancedb
|
||||
|
||||
for database, namespace, name in [
|
||||
("a//b", ["public"], "valid"),
|
||||
("db", ["a$b"], "valid"),
|
||||
("db", "public", "valid"),
|
||||
("db", ["public"], "bad$name"),
|
||||
("db", ["public"], ".."),
|
||||
]:
|
||||
with pytest.raises(ValueError):
|
||||
getattr(Object, builder)(
|
||||
database=database, namespace_path=namespace, name=name
|
||||
)
|
||||
with pytest.raises(ValueError):
|
||||
getattr(_lancedb.AuthzObject, builder)(database, namespace, name)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"builder", ["namespace", "table", "view", "secret", "function"]
|
||||
)
|
||||
def test_empty_and_omitted_namespace_paths_select_public(builder):
|
||||
from lancedb import _lancedb
|
||||
|
||||
kwargs = {"database": "db"}
|
||||
native_args = ()
|
||||
wire = "namespace:db:public"
|
||||
if builder != "namespace":
|
||||
kwargs["name"] = "item"
|
||||
native_args = ("item",)
|
||||
wire = f"{builder}:db:public$item"
|
||||
|
||||
build = getattr(Object, builder)
|
||||
expected = build(**kwargs, namespace_path=["public"])
|
||||
assert str(expected) == wire
|
||||
assert build(**kwargs) == expected
|
||||
for namespace in [[], ()]:
|
||||
assert build(**kwargs, namespace_path=namespace) == expected
|
||||
assert (
|
||||
str(getattr(_lancedb.AuthzObject, builder)("db", namespace, *native_args))
|
||||
== wire
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -55,7 +55,7 @@ def test_catalog_sync_scope_and_serialization(catalog_server):
|
||||
(200, {}),
|
||||
(200, {"tables": []}),
|
||||
(200, {"tables": []}),
|
||||
(200, {"namespaces": ["team/search"], "page_token": "next"}),
|
||||
(200, {"namespaces": ["team/search"]}),
|
||||
(204, None),
|
||||
]
|
||||
)
|
||||
@@ -85,13 +85,12 @@ def test_catalog_sync_scope_and_serialization(catalog_server):
|
||||
with pytest.raises(ValueError, match="sql_host_override must use"):
|
||||
connection.execute_query_async("SELECT 1")
|
||||
assert restored.table_names() == []
|
||||
page = catalog.list_databases(limit=1, page_token="a/b")
|
||||
assert page == lancedb.ListDatabasesResponse(["team/search"], "next")
|
||||
assert list(catalog.list_databases()) == ["team/search"]
|
||||
catalog.drop_database("team/search", ignore_missing=True)
|
||||
assert requests[0][0] == "/v1/namespace/team%2Fsearch/create"
|
||||
assert requests[0][2] == {"mode": "ExistOk"}
|
||||
assert requests[1][0] == "/v1/namespace/team%2Fsearch/describe"
|
||||
assert requests[4][0] == "/v1/namespace/$/list?limit=1&page_token=a%2Fb"
|
||||
assert requests[4][0] == "/v1/namespace/$/list"
|
||||
assert requests[5][2] == {"mode": "Skip", "behavior": "Restrict"}
|
||||
for i, (_, headers, _) in enumerate(requests):
|
||||
headers = {key.lower(): value for key, value in headers.items()}
|
||||
@@ -138,3 +137,171 @@ async def test_catalog_async_and_errors(catalog_server):
|
||||
def test_catalog_requires_remote_endpoint(endpoint):
|
||||
with pytest.raises(ValueError, match="endpoint"):
|
||||
lancedb.connect_catalog(endpoint)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("page_token,page_limit", [(None, None), ("start/token", 2)])
|
||||
async def test_database_names_iteration(
|
||||
catalog_server, asynchronous, page_token, page_limit
|
||||
):
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.extend(
|
||||
[
|
||||
(200, {"namespaces": ["first", "second"], "page_token": "a/b"}),
|
||||
(200, {"namespaces": [], "page_token": "next"}),
|
||||
(200, {"namespaces": ["last", "final"], "page_token": ""}),
|
||||
]
|
||||
)
|
||||
catalog = (
|
||||
await lancedb.connect_catalog_async(endpoint)
|
||||
if asynchronous
|
||||
else lancedb.connect_catalog(endpoint)
|
||||
)
|
||||
names = catalog.list_databases(page_token=page_token, page_limit=page_limit)
|
||||
assert requests == []
|
||||
assert names.num_page_results() == 0
|
||||
assert names.page_token() == page_token
|
||||
for cached, expected in [(1, "first"), (0, "second")]:
|
||||
assert (await names.__anext__() if asynchronous else next(names)) == expected
|
||||
assert names.num_page_results() == cached
|
||||
assert names.page_token() == "a/b"
|
||||
assert len(requests) == 1
|
||||
assert (await names.__anext__() if asynchronous else next(names)) == "last"
|
||||
assert names.num_page_results() == 1
|
||||
assert names.page_token() is None
|
||||
assert (await names.__anext__() if asynchronous else next(names)) == "final"
|
||||
assert names.num_page_results() == 0
|
||||
assert names.page_token() is None
|
||||
for _ in range(2):
|
||||
with pytest.raises(StopAsyncIteration if asynchronous else StopIteration):
|
||||
if asynchronous:
|
||||
await names.__anext__()
|
||||
else:
|
||||
next(names)
|
||||
from urllib.parse import parse_qs, urlsplit
|
||||
|
||||
for request, token in zip(requests, [page_token, "a/b", "next"]):
|
||||
url = urlsplit(request[0])
|
||||
assert url.path == "/v1/namespace/$/list"
|
||||
expected = {}
|
||||
if token is not None:
|
||||
expected["page_token"] = [token]
|
||||
if page_limit is not None:
|
||||
expected["limit"] = [str(page_limit)]
|
||||
assert parse_qs(url.query) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
@pytest.mark.parametrize("fail", [False, True])
|
||||
@pytest.mark.asyncio
|
||||
async def test_database_names_empty_or_error(catalog_server, asynchronous, fail):
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.append(
|
||||
(401, {"error": "unauthorized"}) if fail else (200, {"namespaces": []})
|
||||
)
|
||||
catalog = (
|
||||
await lancedb.connect_catalog_async(endpoint)
|
||||
if asynchronous
|
||||
else lancedb.connect_catalog(endpoint)
|
||||
)
|
||||
names = catalog.list_databases()
|
||||
if fail:
|
||||
with pytest.raises(HttpError):
|
||||
if asynchronous:
|
||||
await names.__anext__()
|
||||
else:
|
||||
next(names)
|
||||
assert ([name async for name in names] if asynchronous else list(names)) == []
|
||||
assert len(requests) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
@pytest.mark.asyncio
|
||||
async def test_database_names_resume(catalog_server, asynchronous):
|
||||
endpoint, requests, responses = catalog_server
|
||||
responses.extend(
|
||||
[
|
||||
(200, {"namespaces": ["first"], "page_token": "resume/token"}),
|
||||
(200, {"namespaces": ["second"]}),
|
||||
]
|
||||
)
|
||||
catalog = (
|
||||
await lancedb.connect_catalog_async(endpoint)
|
||||
if asynchronous
|
||||
else lancedb.connect_catalog(endpoint)
|
||||
)
|
||||
names = catalog.list_databases(page_limit=1)
|
||||
assert (await names.__anext__() if asynchronous else next(names)) == "first"
|
||||
assert names.num_page_results() == 0
|
||||
resumed = catalog.list_databases(page_token=names.page_token(), page_limit=1)
|
||||
assert resumed.page_token() == "resume/token"
|
||||
assert len(requests) == 1
|
||||
assert ([name async for name in resumed] if asynchronous else list(resumed)) == [
|
||||
"second"
|
||||
]
|
||||
assert resumed.page_token() is None
|
||||
assert resumed.num_page_results() == 0
|
||||
assert requests[1][0] == "/v1/namespace/$/list?limit=1&page_token=resume%2Ftoken"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
@pytest.mark.parametrize("page_limit", [0, 2**32 - 1])
|
||||
@pytest.mark.asyncio
|
||||
async def test_database_names_invalid_limit(catalog_server, asynchronous, page_limit):
|
||||
endpoint, requests, _ = catalog_server
|
||||
catalog = (
|
||||
await lancedb.connect_catalog_async(endpoint)
|
||||
if asynchronous
|
||||
else lancedb.connect_catalog(endpoint)
|
||||
)
|
||||
names = catalog.list_databases(page_limit=page_limit)
|
||||
with pytest.raises(ValueError, match="limit"):
|
||||
if asynchronous:
|
||||
await names.__anext__()
|
||||
else:
|
||||
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 == []
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
|
||||
import contextlib
|
||||
import inspect
|
||||
import re
|
||||
import sys
|
||||
@@ -244,6 +245,40 @@ def test_db_contains_and_len_include_all_table_name_pages(tmp_db: lancedb.DBConn
|
||||
assert "does_not_exist" not in tmp_db
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_table_listing_defaults_and_zero_limit(tmp_path):
|
||||
names = [f"t{i:03}" for i in range(130)]
|
||||
for name in names:
|
||||
(tmp_path / f"{name}.lance").mkdir()
|
||||
|
||||
db = lancedb.connect(tmp_path)
|
||||
assert db.table_names() == names
|
||||
assert db.table_names(limit=None) == names
|
||||
assert db.table_names(limit=0) == []
|
||||
first = db.list_tables()
|
||||
assert first.tables == names[:100]
|
||||
assert first.page_token
|
||||
second = db.list_tables(page_token=first.page_token)
|
||||
assert second.tables == names[100:]
|
||||
assert second.page_token is None
|
||||
assert db.list_tables(limit=200).tables == names
|
||||
empty = db.list_tables(limit=0)
|
||||
assert empty.tables == []
|
||||
assert empty.page_token is None
|
||||
|
||||
adb = await lancedb.connect_async(tmp_path)
|
||||
assert await adb.table_names() == names
|
||||
assert await adb.table_names(limit=0) == []
|
||||
first = await adb.list_tables()
|
||||
assert first.tables == names[:100]
|
||||
second = await adb.list_tables(page_token=first.page_token)
|
||||
assert second.tables == names[100:]
|
||||
assert second.page_token is None
|
||||
empty = await adb.list_tables(limit=0)
|
||||
assert empty.tables == []
|
||||
assert empty.page_token is None
|
||||
|
||||
|
||||
def test_db_contains_stops_after_matching_table_page(
|
||||
tmp_db: lancedb.DBConnection, monkeypatch
|
||||
):
|
||||
@@ -350,6 +385,40 @@ async def test_create_table_from_iterator_async(mem_db_async: lancedb.AsyncConne
|
||||
assert await table.count_rows() == 10
|
||||
|
||||
|
||||
def test_create_exist_ok_mode(tmp_db: lancedb.DBConnection):
|
||||
db = tmp_db
|
||||
data = [{"vector": [1.1, 1.2]}, {"vector": [0.2, 1.8]}]
|
||||
table = db.create_table("t1", data, mode="exist_ok")
|
||||
assert table.count_rows() == 2
|
||||
original = table.to_arrow()
|
||||
|
||||
reopened = db.create_table("t1", [{"vector": [9.0, 8.0]}], mode="exist_ok")
|
||||
assert reopened.to_arrow().equals(original)
|
||||
|
||||
|
||||
def test_create_empty_exist_ok_mode(tmp_db: lancedb.DBConnection):
|
||||
schema = pa.schema([pa.field("id", pa.int64())])
|
||||
table = tmp_db.create_table("empty", schema=schema, mode="exist_ok")
|
||||
assert table.schema == schema
|
||||
assert table.count_rows() == 0
|
||||
table.add([{"id": 7}])
|
||||
|
||||
reopened = tmp_db.create_table("empty", schema=schema, mode="exist_ok")
|
||||
assert reopened.to_arrow().to_pylist() == [{"id": 7}]
|
||||
|
||||
bad_schema = pa.schema([pa.field("other", pa.int64())])
|
||||
with pytest.raises(ValueError):
|
||||
tmp_db.create_table("empty", schema=bad_schema, mode="exist_ok")
|
||||
|
||||
|
||||
def test_create_table_invalid_mode(tmp_db: lancedb.DBConnection):
|
||||
db = tmp_db
|
||||
data = [{"vector": [1.1, 1.2]}, {"vector": [0.2, 1.8]}]
|
||||
with pytest.raises(ValueError, match="Invalid mode bogus"):
|
||||
db.create_table("tb", data, mode="bogus")
|
||||
assert "tb" not in db
|
||||
|
||||
|
||||
def test_create_exist_ok(tmp_db: lancedb.DBConnection):
|
||||
from conftest import pandas_string_type
|
||||
|
||||
@@ -405,6 +474,49 @@ async def test_connect(tmp_path):
|
||||
assert str(db) == f"ListingDatabase(uri={tmp_path}, read_consistency_interval=5s)"
|
||||
|
||||
|
||||
@pytest.fixture(params=["memory", "local", "namespace", "from_inner"])
|
||||
def sync_db(request, tmp_path):
|
||||
if request.param == "memory":
|
||||
return lancedb.connect("memory://")
|
||||
if request.param == "namespace":
|
||||
return lancedb.connect_namespace("dir", {"root": str(tmp_path)})
|
||||
db = lancedb.connect(tmp_path)
|
||||
if request.param == "from_inner":
|
||||
return lancedb.db.LanceDBConnection.from_inner(db._inner, None)
|
||||
return db
|
||||
|
||||
|
||||
def test_sync_close(sync_db):
|
||||
assert sync_db.is_open()
|
||||
assert sync_db.close() is None
|
||||
assert not sync_db.is_open()
|
||||
assert sync_db.close() is None
|
||||
|
||||
with pytest.warns(DeprecationWarning, match="table_names"):
|
||||
with pytest.raises(RuntimeError, match="Connection is closed"):
|
||||
sync_db.table_names()
|
||||
with pytest.raises(RuntimeError, match="Connection is closed"):
|
||||
sync_db.list_tables()
|
||||
with pytest.raises(RuntimeError, match="Connection is closed"):
|
||||
sync_db.open_table("test")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raise_error", [False, True])
|
||||
def test_sync_context_manager(sync_db, raise_error):
|
||||
with contextlib.ExitStack() as stack:
|
||||
if raise_error:
|
||||
stack.enter_context(pytest.raises(ValueError, match="test error"))
|
||||
with sync_db as db:
|
||||
assert db is sync_db
|
||||
assert db.is_open()
|
||||
if raise_error:
|
||||
raise ValueError("test error")
|
||||
|
||||
assert not sync_db.is_open()
|
||||
with pytest.raises(RuntimeError, match="Connection is closed"):
|
||||
sync_db.list_tables()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_close(mem_db_async: lancedb.AsyncConnection):
|
||||
assert mem_db_async.is_open()
|
||||
@@ -495,6 +607,27 @@ async def test_create_exist_ok_async(tmp_db_async: lancedb.AsyncConnection):
|
||||
# await db.create_table("test", schema=bad_schema, exist_ok=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_v2_manifest_paths", [False, True])
|
||||
@pytest.mark.parametrize("empty", [False, True])
|
||||
def test_create_table_deprecated_v2_manifest_paths(
|
||||
tmp_path, enable_v2_manifest_paths, empty
|
||||
):
|
||||
db = lancedb.connect(tmp_path)
|
||||
data = None if empty else [{"id": 1}]
|
||||
|
||||
with pytest.warns(DeprecationWarning, match="enable_v2_manifest_paths"):
|
||||
table = db.create_table(
|
||||
"test",
|
||||
data=data,
|
||||
schema=pa.schema([("id", pa.int64())]),
|
||||
enable_v2_manifest_paths=enable_v2_manifest_paths,
|
||||
)
|
||||
|
||||
assert table.uses_v2_manifest_paths() == enable_v2_manifest_paths
|
||||
assert table.to_arrow().to_pylist() == ([] if empty else data)
|
||||
assert db.open_table("test").uses_v2_manifest_paths() == enable_v2_manifest_paths
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_table_v2_manifest_paths_async(tmp_path):
|
||||
db_with_v2_paths = await lancedb.connect_async(
|
||||
|
||||
@@ -2,13 +2,78 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
|
||||
import os
|
||||
from datetime import timedelta
|
||||
from uuid import uuid4
|
||||
|
||||
import lancedb
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from lancedb import LanceDBConnection
|
||||
from lancedb.index import FTS
|
||||
|
||||
# TODO: setup integ test mark and script
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("LANCEDB_URI"), reason="Set LANCEDB_URI to a JSON FTS server"
|
||||
)
|
||||
def test_json_fts_against_server():
|
||||
"""Requires a writable database and capacity to run the distributed indexer.
|
||||
|
||||
Configure LANCEDB_URI, LANCEDB_API_KEY, and optionally LANCEDB_REGION and
|
||||
LANCEDB_HOST_OVERRIDE before running this test.
|
||||
"""
|
||||
uri = os.environ["LANCEDB_URI"]
|
||||
assert uri.startswith("db://"), "LANCEDB_URI must identify a remote database"
|
||||
conn = lancedb.connect(
|
||||
uri,
|
||||
api_key=os.environ.get("LANCEDB_API_KEY"),
|
||||
region=os.environ.get("LANCEDB_REGION", "us-east-1"),
|
||||
host_override=os.environ.get("LANCEDB_HOST_OVERRIDE"),
|
||||
)
|
||||
name = f"json_fts_{uuid4().hex}"
|
||||
table = conn.create_table(
|
||||
name,
|
||||
pa.table(
|
||||
{
|
||||
"id": [0, 1, 2],
|
||||
"doc": pa.array(
|
||||
[
|
||||
'{"meta":{"tag":"nature"},"author":"alice"}',
|
||||
'{"meta":{"tag":"city"},"author":"nature"}',
|
||||
'{"meta":{"tag":"nature"},"author":"bob"}',
|
||||
],
|
||||
type=pa.json_(),
|
||||
),
|
||||
"raw": pa.array([b"nature"] * 3, type=pa.large_binary()),
|
||||
}
|
||||
),
|
||||
)
|
||||
try:
|
||||
table.create_index("doc", config=FTS(), name="doc_fts")
|
||||
table.wait_for_index(["doc_fts"], timeout=timedelta(minutes=5))
|
||||
stats = table.index_stats("doc_fts")
|
||||
assert stats.num_indexed_rows == 3
|
||||
assert stats.num_unindexed_rows == 0
|
||||
assert stats.index_type == "FTS"
|
||||
for query, expected in [
|
||||
("meta.tag,str,nature", [0, 2]),
|
||||
("meta.tag,str,alice", []),
|
||||
("meta.tag,str,tag", []),
|
||||
]:
|
||||
search = table.search(query, query_type="fts", fts_columns="doc")
|
||||
assert sorted(row["id"] for row in search.to_list()) == expected
|
||||
plan = search.explain_plan(True)
|
||||
assert "MatchQuery" in plan, plan
|
||||
assert "FlatMatchQuery" not in plan, plan
|
||||
with pytest.raises(ValueError, match="FTS.*lance.json.*raw binary"):
|
||||
table.create_index("raw", config=FTS())
|
||||
finally:
|
||||
conn.drop_table(name)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Need to set up a local server")
|
||||
def test_against_local_server():
|
||||
conn = LanceDBConnection("lancedb+http://localhost:10024")
|
||||
|
||||
@@ -162,6 +162,29 @@ def test_function_version_binds_named_columns_as_one_immutable_application():
|
||||
] == [("text", "column", "documents.body")]
|
||||
|
||||
|
||||
def test_function_version_names_its_namespace_as_a_separate_field():
|
||||
"""A Function is named by its parts, as a Secret is; a root one states no
|
||||
namespace."""
|
||||
value = job_result("remote_function_job.json")
|
||||
root = FunctionVersion.from_json(json.dumps(value))
|
||||
assert root.namespace_path == ()
|
||||
assert "namespace_path" not in json.loads(root.to_canonical_json())
|
||||
|
||||
nested = FunctionVersion(**{**value, "namespace_path": ["analytics", "features"]})
|
||||
assert nested.name == "embed"
|
||||
assert nested.namespace_path == ("analytics", "features")
|
||||
assert json.loads(nested.to_canonical_json())["namespace_path"] == [
|
||||
"analytics",
|
||||
"features",
|
||||
]
|
||||
application = nested(text=col("documents.body"))
|
||||
assert application.function.name == "embed"
|
||||
assert application.function.namespace_path == ("analytics", "features")
|
||||
assert json.loads(application.to_canonical_json())["function"][
|
||||
"namespace_path"
|
||||
] == ["analytics", "features"]
|
||||
|
||||
|
||||
def test_function_version_binding_validates_names_and_direct_columns():
|
||||
version = FunctionVersion.from_json(
|
||||
json.dumps(job_result("remote_function_job.json"))
|
||||
|
||||
@@ -1499,7 +1499,7 @@ def test_local_function_catalog_operations_are_not_supported(tmp_path):
|
||||
with pytest.raises(NotImplementedError, match=message):
|
||||
db.get_function("normalize_score", version=FUNCTION_VERSION)
|
||||
with pytest.raises(NotImplementedError, match=message):
|
||||
db.list_functions()
|
||||
list(db.list_functions())
|
||||
with pytest.raises(NotImplementedError, match=message):
|
||||
db.drop_function("normalize_score", version=FUNCTION_VERSION)
|
||||
|
||||
@@ -1539,8 +1539,10 @@ def _mock_remote_function_catalog():
|
||||
else (None, None)
|
||||
)
|
||||
if function_action[1] == "create":
|
||||
*namespace_path, name = function_action[0].split("$")
|
||||
state["version"] = {
|
||||
"name": "normalize_score",
|
||||
"name": name,
|
||||
**({"namespace_path": namespace_path} if namespace_path else {}),
|
||||
"version": FUNCTION_VERSION,
|
||||
"object_id": "fixture",
|
||||
"location": "memory:///fixture",
|
||||
@@ -1553,7 +1555,7 @@ def _mock_remote_function_catalog():
|
||||
)["image"],
|
||||
"signature": body["signature"],
|
||||
"secret_bindings": body.get("secret_bindings", []),
|
||||
"created_at": "2026-08-21T00:00:00Z",
|
||||
"created_at_millis": 1787270400000, # 2026-08-21T00:00:00Z
|
||||
}
|
||||
response = {"job_id": "job-register"}
|
||||
status = 202
|
||||
@@ -1565,10 +1567,18 @@ def _mock_remote_function_catalog():
|
||||
"job_state": "DONE",
|
||||
"result": state["version"],
|
||||
}
|
||||
elif self.path == "/v1/function/normalize_score/describe":
|
||||
# A namespaced Function's id is its namespace path and name joined
|
||||
# by `$`, so the name is the last segment of it.
|
||||
elif (
|
||||
function_action[1] == "describe"
|
||||
and function_action[0].split("$")[-1] == "normalize_score"
|
||||
):
|
||||
assert body == {"version": FUNCTION_VERSION}
|
||||
response = state["version"]
|
||||
elif self.path == "/v1/function/normalize_score/drop":
|
||||
elif (
|
||||
function_action[1] == "drop"
|
||||
and function_action[0].split("$")[-1] == "normalize_score"
|
||||
):
|
||||
assert body == {"version": FUNCTION_VERSION}
|
||||
response = {"dropped": True}
|
||||
elif secret_action[1] in ("create", "alter"):
|
||||
@@ -1602,7 +1612,8 @@ def _mock_remote_function_catalog():
|
||||
assert query["page_token"] == "next"
|
||||
self._write_response(200, {"secrets": [{"name": "hf-prod"}]})
|
||||
return
|
||||
if url.path != "/v1/namespace/$/function/list":
|
||||
parts = url.path.strip("/").split("/")
|
||||
if parts[:2] != ["v1", "namespace"] or parts[3:] != ["function", "list"]:
|
||||
self._write_response(404, {"error": "not found"})
|
||||
return
|
||||
assert query["include_definition"] == "true"
|
||||
@@ -1709,7 +1720,7 @@ def test_remote_secret_verbs_round_trip():
|
||||
)
|
||||
assert db.create_secret("openai-prod", "sk-live-0001") is None
|
||||
assert db.alter_secret("openai-prod", "sk-live-0002") is None
|
||||
assert db.list_secrets() == ["openai-prod", "hf-prod"]
|
||||
assert list(db.list_secrets()) == ["openai-prod", "hf-prod"]
|
||||
assert db.drop_secret("openai-prod") is None
|
||||
|
||||
routes = [path for path, _ in state["requests"]]
|
||||
@@ -1770,7 +1781,7 @@ def test_remote_list_functions_paginates_and_returns_typed_versions():
|
||||
)
|
||||
created = db.create_function(normalize_score)
|
||||
state["requests"].clear()
|
||||
functions = db.list_functions()
|
||||
functions = list(db.list_functions())
|
||||
|
||||
assert functions == [created]
|
||||
assert state["requests"] == [
|
||||
@@ -1794,7 +1805,7 @@ async def test_async_remote_list_functions_returns_typed_versions():
|
||||
registration = await db.create_function_async(normalize_score)
|
||||
created = await registration.wait()
|
||||
state["requests"].clear()
|
||||
functions = await db.list_functions()
|
||||
functions = [item async for item in db.list_functions()]
|
||||
|
||||
assert functions == [created]
|
||||
assert [path for path, _ in state["requests"]] == [
|
||||
@@ -1840,3 +1851,102 @@ async def test_async_remote_drop_function_sends_exact_version():
|
||||
{"version": FUNCTION_VERSION},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def test_remote_function_verbs_address_a_namespace_in_the_path():
|
||||
"""A Function's namespace is part of the identifier each route addresses;
|
||||
no body carries it, so a namespaced request differs only in its path."""
|
||||
namespace = ["analytics", "features"]
|
||||
with _mock_remote_function_catalog() as (host, state):
|
||||
db = lancedb.connect(
|
||||
"db://dev",
|
||||
api_key="fake",
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
created = db.create_function(normalize_score, namespace_path=namespace)
|
||||
# The returned parts are what the other verbs take, so a created
|
||||
# Function reopens from its own fields.
|
||||
assert created.name == "normalize_score"
|
||||
assert created.namespace_path == tuple(namespace)
|
||||
reopened = db.get_function(
|
||||
created.name,
|
||||
version=created.version,
|
||||
namespace_path=list(created.namespace_path),
|
||||
)
|
||||
assert reopened == created
|
||||
assert reopened(value=lancedb.col("score")).function.namespace_path == tuple(
|
||||
namespace
|
||||
)
|
||||
list(db.list_functions(namespace_path=namespace))
|
||||
assert (
|
||||
db.drop_function(
|
||||
"normalize_score", version=FUNCTION_VERSION, namespace_path=namespace
|
||||
)
|
||||
is True
|
||||
)
|
||||
dropped, _ = db.drop_function_async(
|
||||
"normalize_score", version=FUNCTION_VERSION, namespace_path=namespace
|
||||
)
|
||||
assert dropped is True
|
||||
|
||||
assert [path for path, _ in state["requests"]] == [
|
||||
"/v1/function/analytics$features$normalize_score/create",
|
||||
"/v1/jobs/describe",
|
||||
"/v1/function/analytics$features$normalize_score/describe",
|
||||
"/v1/namespace/analytics$features/function/list",
|
||||
"/v1/namespace/analytics$features/function/list",
|
||||
"/v1/function/analytics$features$normalize_score/drop",
|
||||
"/v1/function/analytics$features$normalize_score/drop",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_remote_function_verbs_address_a_namespace_in_the_path():
|
||||
namespace = ["analytics", "features"]
|
||||
with _mock_remote_function_catalog() as (host, state):
|
||||
db = await lancedb.connect_async(
|
||||
"db://dev",
|
||||
api_key="fake",
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
registration = await db.create_function_async(
|
||||
normalize_score, namespace_path=namespace
|
||||
)
|
||||
await registration.wait()
|
||||
await db.get_function(
|
||||
"normalize_score", version=FUNCTION_VERSION, namespace_path=namespace
|
||||
)
|
||||
[item async for item in db.list_functions(namespace_path=namespace)]
|
||||
assert (
|
||||
await db.drop_function(
|
||||
"normalize_score", version=FUNCTION_VERSION, namespace_path=namespace
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
assert [path for path, _ in state["requests"]] == [
|
||||
"/v1/function/analytics$features$normalize_score/create",
|
||||
"/v1/jobs/describe",
|
||||
"/v1/function/analytics$features$normalize_score/describe",
|
||||
"/v1/namespace/analytics$features/function/list",
|
||||
"/v1/namespace/analytics$features/function/list",
|
||||
"/v1/function/analytics$features$normalize_score/drop",
|
||||
]
|
||||
|
||||
|
||||
def test_a_function_namespace_path_must_be_a_list_of_segments():
|
||||
"""A bare string is refused rather than read as a sequence of one-character
|
||||
segments, and nothing reaches the service."""
|
||||
with _mock_remote_function_catalog() as (host, state):
|
||||
db = lancedb.connect(
|
||||
"db://dev",
|
||||
api_key="fake",
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
with pytest.raises(TypeError):
|
||||
list(db.list_functions(namespace_path="analytics"))
|
||||
|
||||
assert state["requests"] == []
|
||||
|
||||
@@ -13,6 +13,7 @@
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import shutil
|
||||
@@ -44,6 +45,75 @@ from utils import exception_output
|
||||
TEST_LANGUAGE_MODEL_HOME = Path(__file__).parent / "models"
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
@pytest.mark.parametrize("block_size", [128, 256])
|
||||
def test_json_fts(tmp_path, block_size):
|
||||
db = ldb.connect(tmp_path)
|
||||
data = pa.table(
|
||||
{
|
||||
"id": [0, 1, 2, 3],
|
||||
"doc": pa.array(
|
||||
[
|
||||
'{"title":"brown fox","meta":{"tag":"nature"},"author":"alice"}',
|
||||
'{"title":"blue car","meta":{"tag":"city"}}',
|
||||
'{"title":"brown bear","meta":{"tag":"nature"}}',
|
||||
'{"title":"red fox","author":"alice"}',
|
||||
],
|
||||
type=pa.json_(),
|
||||
),
|
||||
}
|
||||
)
|
||||
table = db.create_table("json_fts", data)
|
||||
table.create_index("doc", config=FTS(block_size=block_size), name="doc_fts")
|
||||
table = db.open_table("json_fts")
|
||||
index = table.list_indices()[0]
|
||||
assert index.name == "doc_fts"
|
||||
assert index.index_type == "FTS"
|
||||
assert index.columns == ["doc"]
|
||||
stats = table.index_stats("doc_fts")
|
||||
assert stats.num_indexed_rows == 4
|
||||
assert stats.num_unindexed_rows == 0
|
||||
|
||||
for query, expected in [
|
||||
("title,str,brown", [0, 2]),
|
||||
("meta.tag,str,nature", [0, 2]),
|
||||
("author,str,alice", [0, 3]),
|
||||
("title,str,alice", []),
|
||||
("title,str,title", []),
|
||||
]:
|
||||
search = table.search(query, query_type="fts", fts_columns="doc")
|
||||
assert sorted(row["id"] for row in search.to_list()) == expected
|
||||
plan = search.explain_plan(True)
|
||||
assert "MatchQuery" in plan, plan
|
||||
assert "FlatMatchQuery" not in plan, plan
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
@pytest.mark.asyncio
|
||||
async def test_json_fts_async(tmp_path):
|
||||
db = await ldb.connect_async(tmp_path)
|
||||
table = await db.create_table(
|
||||
"json_fts",
|
||||
pa.table({"doc": pa.array(['{"meta":{"tag":"nature"}}'], type=pa.json_())}),
|
||||
)
|
||||
await table.create_index("doc", config=FTS())
|
||||
search = table.query().nearest_to_text(MatchQuery("meta.tag,str,nature", "doc"))
|
||||
assert len(await search.to_list()) == 1
|
||||
plan = await search.explain_plan(True)
|
||||
assert "MatchQuery" in plan, plan
|
||||
assert "FlatMatchQuery" not in plan, plan
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [pa.binary(), pa.large_binary()])
|
||||
def test_fts_rejects_raw_binary(tmp_path, dtype):
|
||||
table = ldb.connect(tmp_path).create_table(
|
||||
"binary_fts", pa.table({"doc": pa.array([b'{"title":"brown"}'], type=dtype)})
|
||||
)
|
||||
with pytest.raises(ValueError, match="FTS.*lance.json.*raw binary"):
|
||||
table.create_index("doc", config=FTS())
|
||||
assert table.list_indices() == []
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def table(tmp_path) -> ldb.table.LanceTable:
|
||||
# Use local random state to avoid affecting other tests
|
||||
@@ -894,6 +964,78 @@ def test_tokenize_uses_simple_index_tokenizer(mem_db: DBConnection):
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("legacy", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
("base_tokenizer", "options", "expected_filters", "expected_tokens"),
|
||||
[
|
||||
(
|
||||
"code",
|
||||
{},
|
||||
(False, False),
|
||||
[("running", 0), ("in", 1), ("cafe", 2)],
|
||||
),
|
||||
(
|
||||
"code",
|
||||
{"stem": None, "remove_stop_words": None},
|
||||
(False, False),
|
||||
[("running", 0), ("in", 1), ("cafe", 2)],
|
||||
),
|
||||
(
|
||||
"code",
|
||||
{"stem": True},
|
||||
(True, False),
|
||||
[("run", 0), ("in", 1), ("cafe", 2)],
|
||||
),
|
||||
(
|
||||
"code",
|
||||
{"remove_stop_words": True},
|
||||
(False, True),
|
||||
[("running", 0), ("cafe", 2)],
|
||||
),
|
||||
(
|
||||
"code",
|
||||
{"stem": True, "remove_stop_words": True},
|
||||
(True, True),
|
||||
[("run", 0), ("cafe", 2)],
|
||||
),
|
||||
(
|
||||
"simple",
|
||||
{},
|
||||
(True, True),
|
||||
[("run", 0), ("cafe", 2)],
|
||||
),
|
||||
(
|
||||
"simple",
|
||||
{"stem": False, "remove_stop_words": False},
|
||||
(False, False),
|
||||
[("running", 0), ("in", 1), ("cafe", 2)],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_fts_tokenizer_filter_defaults(
|
||||
mem_db, legacy, base_tokenizer, options, expected_filters, expected_tokens
|
||||
):
|
||||
text = "Running in café"
|
||||
table = mem_db.create_table("test_filter_defaults", data=[{"text": text}])
|
||||
if legacy:
|
||||
with pytest.warns(DeprecationWarning, match="create_fts_index"):
|
||||
table.create_fts_index("text", base_tokenizer=base_tokenizer, **options)
|
||||
else:
|
||||
table.create_index("text", config=FTS(base_tokenizer=base_tokenizer, **options))
|
||||
|
||||
details = table.list_indices()[0].index_details
|
||||
assert (details["stem"], details["remove_stop_words"]) == expected_filters
|
||||
assert [
|
||||
(token.text, token.position) for token in table.tokenize(text, column="text")
|
||||
] == expected_tokens
|
||||
assert [
|
||||
(token.text, token.position)
|
||||
for token in ldb.tokenize(text, base_tokenizer=base_tokenizer, **options)
|
||||
] == expected_tokens
|
||||
if base_tokenizer == "code" and not expected_filters[1]:
|
||||
assert table.search("in", query_type="fts").to_list()[0]["text"] == text
|
||||
|
||||
|
||||
def test_tokenize_uses_icu_index_tokenizer_by_name(mem_db: DBConnection):
|
||||
data = pa.table({"text": ["Hello, こんにちは世界!"]})
|
||||
table = mem_db.create_table("test_tokenize_icu", data=data)
|
||||
@@ -1137,6 +1279,34 @@ def test_fts_lindera_tokenizer(
|
||||
assert [row["text"] for row in results] == ["成田国際空港"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fuzziness", [None, 0, 1, 2])
|
||||
def test_match_query_fuzziness_to_json(fuzziness):
|
||||
query = MatchQuery("puppy", "text", fuzziness=fuzziness)
|
||||
|
||||
assert query.fuzziness == fuzziness
|
||||
assert json.loads(query.to_json())["match"]["fuzziness"] == fuzziness
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"options, expected",
|
||||
[
|
||||
pytest.param({}, ["ab"], id="default-exact"),
|
||||
pytest.param({"fuzziness": 0}, ["ab"], id="explicit-exact"),
|
||||
pytest.param({"fuzziness": None}, ["ab"], id="automatic"),
|
||||
pytest.param({"fuzziness": 1}, ["ab", "ac"], id="one-edit"),
|
||||
pytest.param({"fuzziness": 2}, ["ab", "ac", "zz"], id="two-edits"),
|
||||
],
|
||||
)
|
||||
def test_match_query_fuzziness_search(tmp_db, options, expected):
|
||||
# Automatic fuzziness is zero for two-character terms.
|
||||
table = tmp_db.create_table("fuzzy", pa.table({"text": ["ab", "ac", "zz"]}))
|
||||
table.create_index("text", config=FTS(stem=False, remove_stop_words=False))
|
||||
|
||||
results = table.search(MatchQuery("ab", "text", **options)).to_list()
|
||||
|
||||
assert sorted(row["text"] for row in results) == sorted(expected)
|
||||
|
||||
|
||||
def test_fts_query_to_json():
|
||||
"""Test that FTS query to_json() produces valid JSON strings with exact format."""
|
||||
|
||||
@@ -1267,7 +1437,10 @@ def test_fts_query_to_json():
|
||||
json_str = multi_match.to_json()
|
||||
expected = (
|
||||
'{"multi_match":{"query":"python","columns":["tags","title"],'
|
||||
'"boost":[1.0,1.0]}}'
|
||||
'"boost":[1.0,1.0],"match_queries":[{"column":"tags","terms":"python",'
|
||||
'"boost":1.0,"fuzziness":0,"max_expansions":50,"operator":"Or",'
|
||||
'"prefix_length":0},{"column":"title","terms":"python","boost":1.0,'
|
||||
'"fuzziness":0,"max_expansions":50,"operator":"Or","prefix_length":0}]}}'
|
||||
)
|
||||
assert json_str == expected
|
||||
|
||||
|
||||
@@ -355,6 +355,20 @@ def test_hybrid_query_minimum_nprobes_zero_raises(sync_table: Table):
|
||||
)
|
||||
|
||||
|
||||
def test_hybrid_query_nprobes_remains_independent(sync_table: Table):
|
||||
query = sync_table.search(query_type="hybrid").nprobes(20)
|
||||
|
||||
assert query._nprobes == 20
|
||||
assert query._minimum_nprobes is None
|
||||
assert query._maximum_nprobes is None
|
||||
|
||||
query = query.minimum_nprobes(5).nprobes(20)
|
||||
|
||||
assert query._nprobes == 20
|
||||
assert query._minimum_nprobes == 5
|
||||
assert query._maximum_nprobes is None
|
||||
|
||||
|
||||
def test_hybrid_query_distance_range(sync_table: Table):
|
||||
reranker = RRFReranker(return_score="all")
|
||||
result = (
|
||||
|
||||
@@ -0,0 +1,194 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import asyncio
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
import json
|
||||
from pathlib import Path
|
||||
from threading import Thread
|
||||
from urllib.parse import parse_qs, urlsplit
|
||||
|
||||
import lancedb
|
||||
from lancedb.remote.errors import HttpError
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def listing_server():
|
||||
requests, responses = [], []
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def handle_request(self):
|
||||
body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
|
||||
requests.append((self.command, self.path, json.loads(body) if body else {}))
|
||||
status, response = responses.pop(0)
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.end_headers()
|
||||
self.wfile.write(json.dumps(response).encode())
|
||||
|
||||
do_GET = handle_request
|
||||
do_POST = handle_request
|
||||
|
||||
with ThreadingHTTPServer(("127.0.0.1", 0), Handler) as server:
|
||||
thread = Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
yield f"http://127.0.0.1:{server.server_port}", requests, responses
|
||||
finally:
|
||||
server.shutdown()
|
||||
thread.join()
|
||||
|
||||
|
||||
def response_page(kind, names, token):
|
||||
if kind == "secrets":
|
||||
items = [{"name": name} for name in names]
|
||||
elif kind == "jobs":
|
||||
items = [
|
||||
{
|
||||
"job_id": name,
|
||||
"table": "t",
|
||||
"job_type": "create_index",
|
||||
"state": "in_progress",
|
||||
"created_at_millis": 1,
|
||||
}
|
||||
for name in names
|
||||
]
|
||||
elif kind == "functions":
|
||||
fixture = (
|
||||
Path(__file__).resolve().parents[3]
|
||||
/ "rust/lancedb/tests/fixtures/first_class_functions/v1"
|
||||
/ "remote_function_version.canonical.json"
|
||||
)
|
||||
definition = json.loads(fixture.read_text())
|
||||
items = [
|
||||
{"name": "embed", "version": "1", "definition": definition} for _ in names
|
||||
]
|
||||
else:
|
||||
items = names
|
||||
key = "views" if kind == "materialized_views" else kind
|
||||
return {key: items, "page_token": token}
|
||||
|
||||
|
||||
async def connection(endpoint, asynchronous):
|
||||
kwargs = dict(
|
||||
api_key="fake",
|
||||
host_override=endpoint,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
if asynchronous:
|
||||
return await lancedb.connect_async("db://dev", **kwargs)
|
||||
return lancedb.connect("db://dev", **kwargs)
|
||||
|
||||
|
||||
async def advance(items, asynchronous):
|
||||
return await items.__anext__() if asynchronous else next(items)
|
||||
|
||||
|
||||
KINDS = ["secrets", "views", "jobs", "functions", "materialized_views"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("kind", KINDS)
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
async def test_listing_lazy_pages_and_resume(listing_server, kind, asynchronous):
|
||||
endpoint, requests, responses = listing_server
|
||||
responses.extend(
|
||||
(200, response_page(kind, names, token))
|
||||
for names, token in [(["a", "b"], "empty"), ([], "last"), (["c", "d"], "")]
|
||||
)
|
||||
db = await connection(endpoint, asynchronous)
|
||||
method = getattr(db, "list_" + kind)
|
||||
kwargs = (
|
||||
{} if kind in ("jobs", "materialized_views") else {"namespace_path": ["team"]}
|
||||
)
|
||||
items = method(page_token="start/token", page_limit=2, **kwargs)
|
||||
assert requests == []
|
||||
assert items.page_token() == "start/token"
|
||||
assert items.num_page_results() == 0
|
||||
first = await advance(items, asynchronous)
|
||||
if kind == "functions":
|
||||
assert first.name == "embed"
|
||||
assert first.version == "1"
|
||||
elif kind == "jobs":
|
||||
assert first.job_id == "a"
|
||||
assert first.state == "running"
|
||||
else:
|
||||
assert first == "a"
|
||||
assert items.num_page_results() == 1
|
||||
assert items.page_token() == "empty"
|
||||
await advance(items, asynchronous)
|
||||
assert items.num_page_results() == 0
|
||||
assert len(requests) == 1
|
||||
resumed = method(page_token=items.page_token(), page_limit=2, **kwargs)
|
||||
await advance(resumed, asynchronous)
|
||||
assert len(requests) == 3
|
||||
assert resumed.page_token() is None
|
||||
assert resumed.num_page_results() == 1
|
||||
remaining = [item async for item in resumed] if asynchronous else list(resumed)
|
||||
assert len(remaining) == 1
|
||||
assert resumed.num_page_results() == 0
|
||||
assert ([item async for item in resumed] if asynchronous else list(resumed)) == []
|
||||
for (verb, url, body), token in zip(requests, ["start/token", "empty", "last"]):
|
||||
path = urlsplit(url)
|
||||
if kind == "jobs":
|
||||
assert verb == "POST"
|
||||
assert path.path == "/v1/jobs/list"
|
||||
assert body == {"page_token": token, "limit": 2}
|
||||
else:
|
||||
namespace = "$" if kind == "materialized_views" else "team"
|
||||
resource = kind.removesuffix("s")
|
||||
assert path.path == f"/v1/namespace/{namespace}/{resource}/list"
|
||||
expected = {"page_token": [token], "limit": ["2"]}
|
||||
if kind == "functions":
|
||||
expected["include_definition"] = ["true"]
|
||||
assert parse_qs(path.query) == expected
|
||||
assert responses == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("kind", KINDS)
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
async def test_listing_error_retains_token(listing_server, kind, asynchronous):
|
||||
endpoint, requests, responses = listing_server
|
||||
responses.extend([(401, {}), (200, response_page(kind, ["a"], None))])
|
||||
db = await connection(endpoint, asynchronous)
|
||||
method = getattr(db, "list_" + kind)
|
||||
items = method(page_token="retry")
|
||||
with pytest.raises(HttpError):
|
||||
await advance(items, asynchronous)
|
||||
assert items.page_token() == "retry"
|
||||
assert items.num_page_results() == 0
|
||||
assert ([item async for item in items] if asynchronous else list(items)) == []
|
||||
assert len(requests) == 1
|
||||
resumed = method(page_token=items.page_token())
|
||||
assert len([item async for item in resumed] if asynchronous else list(resumed)) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("kind", KINDS)
|
||||
async def test_listing_serializes_concurrent_advances(listing_server, kind):
|
||||
endpoint, requests, responses = listing_server
|
||||
responses.append((200, response_page(kind, ["a", "b"], None)))
|
||||
db = await connection(endpoint, True)
|
||||
items = getattr(db, "list_" + kind)()
|
||||
results = await asyncio.gather(
|
||||
items.__anext__(), items.__anext__(), items.__anext__(), return_exceptions=True
|
||||
)
|
||||
assert (
|
||||
len([result for result in results if isinstance(result, StopAsyncIteration)])
|
||||
== 1
|
||||
)
|
||||
assert len(requests) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("kind", KINDS)
|
||||
@pytest.mark.parametrize("page_limit", [0, 2**32 - 1])
|
||||
async def test_listing_invalid_limit_does_not_fetch(listing_server, kind, page_limit):
|
||||
endpoint, requests, _ = listing_server
|
||||
db = await connection(endpoint, True)
|
||||
items = getattr(db, "list_" + kind)(page_limit=page_limit)
|
||||
with pytest.raises(ValueError, match="limit"):
|
||||
await items.__anext__()
|
||||
assert requests == []
|
||||
@@ -93,28 +93,10 @@ def test_remote_list_uses_namespace_route():
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
assert db.list_materialized_views() == ["daily_sales"]
|
||||
assert list(db.list_materialized_views()) == ["daily_sales"]
|
||||
assert requests == ["/v1/namespace/$/materialized_view/list"]
|
||||
|
||||
|
||||
def test_remote_create_async_returns_server_job():
|
||||
with mock_remote_materialized_view_create() as (host, requests):
|
||||
db = lancedb.connect(
|
||||
"db://dev",
|
||||
api_key="fake",
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
job = db.create_materialized_view_async("adults", "people", where="age >= 18")
|
||||
assert job.id == "mv-create-123"
|
||||
assert requests == [
|
||||
(
|
||||
"/v1/materialized_view/adults/create",
|
||||
{"query": 'SELECT * FROM "people" WHERE age >= 18', "with_no_data": False},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def test_remote_drop_async_returns_server_job():
|
||||
with mock_remote_materialized_view_create() as (host, requests):
|
||||
db = lancedb.connect(
|
||||
@@ -208,7 +190,7 @@ def test_create_and_refresh_jobs(tmp_path):
|
||||
drop_job = db.drop_materialized_view_async("adults")
|
||||
assert drop_job.id is None
|
||||
assert drop_job.wait() is None
|
||||
assert "adults" not in db.list_materialized_views()
|
||||
assert "adults" not in list(db.list_materialized_views())
|
||||
|
||||
|
||||
def test_definition_round_trips(tmp_path):
|
||||
@@ -221,29 +203,29 @@ def test_definition_round_trips(tmp_path):
|
||||
)
|
||||
|
||||
|
||||
def test_incremental_refresh_after_append(tmp_path):
|
||||
def test_full_refresh_after_append(tmp_path):
|
||||
db = make_db(tmp_path)
|
||||
view = db.create_materialized_view("copy", "people", with_no_data=True)
|
||||
view.refresh()
|
||||
|
||||
db.open_table("people").add([{"name": "alan", "age": 41}])
|
||||
result = view.refresh()
|
||||
assert result.mode == "incremental"
|
||||
assert result.rows_written == 1
|
||||
assert result.mode == "rebuild"
|
||||
assert result.rows_written == 4
|
||||
assert view.table.count_rows() == 4
|
||||
|
||||
assert view.refresh().mode == "no_op"
|
||||
assert view.refresh().mode == "rebuild"
|
||||
|
||||
|
||||
def test_incremental_refresh_after_update(tmp_path):
|
||||
def test_full_refresh_after_update(tmp_path):
|
||||
db = make_db(tmp_path)
|
||||
view = db.create_materialized_view("copy", "people", with_no_data=True)
|
||||
view.refresh()
|
||||
|
||||
db.open_table("people").update(where="name = 'kid'", values={"age": 8})
|
||||
result = view.refresh()
|
||||
assert result.mode == "incremental"
|
||||
assert result.rows_written == 1
|
||||
assert result.mode == "rebuild"
|
||||
assert result.rows_written == 3
|
||||
rows = view.table.search().to_list()
|
||||
assert sorted(row["age"] for row in rows) == [8, 36, 85]
|
||||
|
||||
@@ -264,18 +246,43 @@ def test_legacy_storage_source_update_rebuilds(tmp_path):
|
||||
assert sorted(row["age"] for row in rows) == [8, 36]
|
||||
|
||||
|
||||
def test_direct_writes_to_a_view_are_refused(tmp_path):
|
||||
db = make_db(tmp_path)
|
||||
view = db.create_materialized_view("copy", "people")
|
||||
table = view.table
|
||||
row = [{"name": "eve", "age": 1}]
|
||||
refused = pytest.raises(NotImplementedError, match="is a materialized view")
|
||||
|
||||
with refused:
|
||||
table.add(row)
|
||||
with refused:
|
||||
table.update(where="name = 'kid'", values={"age": 8})
|
||||
with refused:
|
||||
table.delete("age < 18")
|
||||
with refused:
|
||||
table.merge_insert("name").when_not_matched_insert_all().execute(row)
|
||||
with refused:
|
||||
table.add_columns({"older": "age + 1"})
|
||||
with refused:
|
||||
table.drop_columns(["age"])
|
||||
assert table.count_rows() == 3
|
||||
|
||||
db.open_table("people").add(row)
|
||||
assert view.refresh().rows_written == 4
|
||||
|
||||
|
||||
def test_list_and_not_a_view(tmp_path):
|
||||
db = make_db(tmp_path)
|
||||
db.create_materialized_view("adults", "people", where="age >= 18")
|
||||
|
||||
assert db.list_materialized_views() == ["adults"]
|
||||
assert list(db.list_materialized_views()) == ["adults"]
|
||||
with pytest.raises(ValueError, match="not a materialized view"):
|
||||
db.open_materialized_view("people")
|
||||
with pytest.raises(ValueError, match="not a materialized view"):
|
||||
db.drop_materialized_view("people")
|
||||
|
||||
db.drop_materialized_view("adults")
|
||||
assert db.list_materialized_views() == []
|
||||
assert list(db.list_materialized_views()) == []
|
||||
|
||||
|
||||
def test_invalid_expression_fails_at_create(tmp_path):
|
||||
@@ -303,7 +310,7 @@ async def test_async_create_refresh_and_open(tmp_path):
|
||||
reopened = await db.open_materialized_view("shouts")
|
||||
definition = await reopened.definition()
|
||||
assert definition.query == "SELECT upper(name) AS shout FROM people"
|
||||
assert await db.list_materialized_views() == ["shouts"]
|
||||
assert [item async for item in db.list_materialized_views()] == ["shouts"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -327,12 +334,12 @@ async def test_async_create_and_refresh_jobs(tmp_path):
|
||||
drop_job = await db.drop_materialized_view_async("adults")
|
||||
assert drop_job.id is None
|
||||
assert await drop_job.wait() is None
|
||||
assert "adults" not in await db.list_materialized_views()
|
||||
assert "adults" not in [item async for item in db.list_materialized_views()]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_incremental(tmp_path):
|
||||
db = await lancedb.connect_async(tmp_path, storage_options=STABLE_ROW_IDS)
|
||||
async def test_async_refresh_rebuilds(tmp_path):
|
||||
db = await lancedb.connect_async(tmp_path)
|
||||
await db.create_table("people", [{"name": "ada", "age": 36}])
|
||||
view = await db.create_materialized_view("copy", "people", with_no_data=True)
|
||||
await view.refresh()
|
||||
@@ -340,15 +347,15 @@ async def test_async_incremental(tmp_path):
|
||||
table = await db.open_table("people")
|
||||
await table.add([{"name": "alan", "age": 41}])
|
||||
result = await view.refresh()
|
||||
assert result.mode == "incremental"
|
||||
assert result.rows_written == 1
|
||||
assert result.mode == "rebuild"
|
||||
assert result.rows_written == 2
|
||||
|
||||
|
||||
def test_source_requires_stable_row_ids(tmp_path):
|
||||
def test_source_does_not_require_stable_row_ids(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
db.create_table("plain", [{"x": 1}])
|
||||
with pytest.raises(Exception, match="stable row ids"):
|
||||
db.create_materialized_view("v", "plain")
|
||||
view = db.create_materialized_view("v", "plain")
|
||||
assert view.table.count_rows() == 1
|
||||
|
||||
|
||||
def test_bare_select_names_are_quoted(tmp_path):
|
||||
@@ -373,7 +380,7 @@ def test_scalar_select_is_one_column(tmp_path):
|
||||
view = db.create_materialized_view("just_name", "people", select="name")
|
||||
view.refresh()
|
||||
rows = view.table.search().to_list()
|
||||
assert set(rows[0]) - {"__source_row_id"} == {"name"}
|
||||
assert set(rows[0]) == {"name"}
|
||||
assert sorted(row["name"] for row in rows) == ["ada", "grace", "kid"]
|
||||
|
||||
|
||||
@@ -384,7 +391,7 @@ async def test_async_scalar_select_is_one_column(tmp_path):
|
||||
view = await db.create_materialized_view("just_name", "people", select="name")
|
||||
await view.refresh()
|
||||
rows = await view.table.query().to_list()
|
||||
assert set(rows[0]) - {"__source_row_id"} == {"name"}
|
||||
assert set(rows[0]) == {"name"}
|
||||
|
||||
|
||||
def test_limit_above_i64_max_is_refused(tmp_path):
|
||||
@@ -417,7 +424,7 @@ def test_namespace_connection_materialized_views(tmp_path):
|
||||
view = db.create_materialized_view("adults", "people", where="age >= 18")
|
||||
view.refresh()
|
||||
assert view.table.count_rows() == 1
|
||||
assert db.list_materialized_views() == ["adults"]
|
||||
assert list(db.list_materialized_views()) == ["adults"]
|
||||
|
||||
reopened = db.open_materialized_view("adults")
|
||||
assert reopened.definition.query.startswith("SELECT name, age FROM ")
|
||||
@@ -433,7 +440,7 @@ def test_namespace_connection_materialized_views(tmp_path):
|
||||
|
||||
assert db.drop_materialized_view_async("job_view").wait() is None
|
||||
db.drop_materialized_view("adults")
|
||||
assert db.list_materialized_views() == []
|
||||
assert list(db.list_materialized_views()) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -452,7 +459,7 @@ async def test_async_namespace_connection_materialized_views(tmp_path):
|
||||
view = await db.create_materialized_view("adults", "people", where="age >= 18")
|
||||
await view.refresh()
|
||||
assert await view.table.count_rows() == 1
|
||||
assert await db.list_materialized_views() == ["adults"]
|
||||
assert [item async for item in db.list_materialized_views()] == ["adults"]
|
||||
|
||||
reopened = await db.open_materialized_view("adults")
|
||||
assert (await reopened.definition()).query.startswith("SELECT name, age FROM ")
|
||||
@@ -478,7 +485,7 @@ async def test_async_namespace_connection_materialized_views(tmp_path):
|
||||
drop_job = await db.drop_materialized_view_async("job_view")
|
||||
assert await drop_job.wait() is None
|
||||
await db.drop_materialized_view("adults")
|
||||
assert await db.list_materialized_views() == []
|
||||
assert [item async for item in db.list_materialized_views()] == []
|
||||
|
||||
|
||||
def test_stored_queries_and_legacy_layouts_are_read():
|
||||
@@ -520,166 +527,26 @@ def test_stored_queries_and_legacy_layouts_are_read():
|
||||
"SELECT * FROM people"
|
||||
)
|
||||
|
||||
for fmt, function in [(3, "vector_duplicate_pairs"), (4, "vector_dedup")]:
|
||||
native = f"SELECT * FROM {function}('images', 3, 'phash', 4)"
|
||||
assert read({"format": fmt, "query": native}).query == native
|
||||
|
||||
# A newer writer's layout is reported, never guessed at.
|
||||
for newer in (
|
||||
{"format": 5, "query": query},
|
||||
{"format": 3, "query": query},
|
||||
{"kind": "select_v3", "source_table": "people"},
|
||||
):
|
||||
with pytest.raises(NotImplementedError, match="cannot refresh"):
|
||||
read(newer)
|
||||
|
||||
|
||||
def indexed_dedup_source(db, name="images"):
|
||||
import pyarrow as pa
|
||||
from lancedb.index import IvfFlat
|
||||
|
||||
vectors = pa.FixedSizeListArray.from_arrays(
|
||||
pa.array(
|
||||
[v for i in range(256) for v in (float(i if i < 3 else i * 10), 0.0)],
|
||||
type=pa.float32(),
|
||||
),
|
||||
2,
|
||||
)
|
||||
table = db.create_table(name, pa.table({"id": range(256), "vector": vectors}))
|
||||
table.create_index("vector", config=IvfFlat(num_partitions=1))
|
||||
return table
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stable", [False, True])
|
||||
def test_vector_dedup_source_captures_snapshot_and_preserves_source(tmp_path, stable):
|
||||
db = lancedb.connect(
|
||||
tmp_path,
|
||||
storage_options={"new_table_enable_stable_row_ids": str(stable).lower()},
|
||||
)
|
||||
source = indexed_dedup_source(db, "images.with.dot")
|
||||
version = source.version
|
||||
expression = lancedb.vector_dedup(
|
||||
source.name, column="vector", distance_threshold=1
|
||||
)
|
||||
view = db.create_materialized_view("clean", expression)
|
||||
assert "vector_dedup(" in view.definition.query
|
||||
assert f", {version}, 'vector', 1.0)" in view.definition.query
|
||||
assert view.table.count_rows() == 255
|
||||
assert [r["id"] for r in view.table.search().where("id < 3").to_list()] == [0, 2]
|
||||
assert source.version == version
|
||||
assert source.count_rows() == 256
|
||||
|
||||
source.delete("id = 0")
|
||||
view.refresh_async().wait()
|
||||
assert view.table.search().where("id = 0").to_list() == [
|
||||
{"id": 0, "vector": [0.0, 0.0]}
|
||||
]
|
||||
assert expression.dataset_version is None # reusable immutable declaration
|
||||
|
||||
|
||||
def test_vector_dedup_create_job_and_explicit_snapshot(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
source = indexed_dedup_source(db)
|
||||
version = source.version
|
||||
source.delete("id = 0")
|
||||
job = db.create_materialized_view_async(
|
||||
"clean",
|
||||
lancedb.vector_dedup(
|
||||
"images", column="vector", distance_threshold=1, dataset_version=version
|
||||
),
|
||||
)
|
||||
assert job.wait() is None
|
||||
view = db.open_materialized_view("clean")
|
||||
assert view.table.count_rows() == 255
|
||||
assert len(view.table.search().where("id = 0").to_list()) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_vector_dedup_declaration_and_creation(tmp_path):
|
||||
sync_db = lancedb.connect(tmp_path)
|
||||
source = indexed_dedup_source(sync_db)
|
||||
db = await lancedb.connect_async(tmp_path)
|
||||
expression = lancedb.vector_dedup("images", column="vector", distance_threshold=1)
|
||||
view = await db.create_materialized_view("clean", expression, with_no_data=True)
|
||||
assert await view.table.count_rows() == 0
|
||||
# The snapshot is captured at declaration even if refresh comes later.
|
||||
source.delete("id = 0")
|
||||
job = await view.refresh_async()
|
||||
await job.wait()
|
||||
assert await view.table.count_rows() == 255
|
||||
another = await db.create_materialized_view_async(
|
||||
"later", expression, with_no_data=True
|
||||
)
|
||||
await another.wait()
|
||||
later = await db.open_materialized_view("later")
|
||||
assert (await later.definition()).query != (await view.definition()).query
|
||||
assert await later.table.count_rows() == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"options", [{"select": ["id"]}, {"select": []}, {"where": "id > 0"}, {"limit": 1}]
|
||||
)
|
||||
def test_vector_dedup_rejects_extra_clauses_before_opening_source(tmp_path, options):
|
||||
db = lancedb.connect(tmp_path)
|
||||
expression = lancedb.vector_dedup("missing", column="vector", distance_threshold=1)
|
||||
with pytest.raises(ValueError, match="cannot be combined"):
|
||||
db.create_materialized_view("bad", expression, **options)
|
||||
assert db.list_tables().tables == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("version", [0, -1, True, 1.5, "2", 2**64])
|
||||
def test_vector_dedup_rejects_invalid_snapshot(version):
|
||||
with pytest.raises(ValueError, match="dataset_version"):
|
||||
lancedb.vector_dedup(
|
||||
"images", column="vector", distance_threshold=1, dataset_version=version
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"threshold", [float("nan"), float("inf"), -float("inf"), 1e100, 10**1000, True, "4"]
|
||||
)
|
||||
def test_vector_dedup_rejects_invalid_threshold(threshold):
|
||||
with pytest.raises((ValueError, TypeError), match="distance_threshold"):
|
||||
lancedb.vector_dedup("images", column="vector", distance_threshold=threshold)
|
||||
|
||||
|
||||
def test_vector_dedup_descriptor_is_frozen():
|
||||
from dataclasses import FrozenInstanceError
|
||||
|
||||
expression = lancedb.vector_dedup("images", column="vector", distance_threshold=1)
|
||||
with pytest.raises(FrozenInstanceError):
|
||||
expression.column = "other"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("with_no_data", [False, True])
|
||||
def test_remote_vector_dedup_uses_mv_endpoint_without_sql_client(with_no_data):
|
||||
with mock_remote_materialized_view_create() as (host, requests):
|
||||
db = lancedb.connect(
|
||||
"db://dev",
|
||||
api_key="fake",
|
||||
host_override=host,
|
||||
sql_host_override="grpc://127.0.0.1:1",
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
job = db.create_materialized_view_async(
|
||||
"clean",
|
||||
lancedb.vector_dedup(
|
||||
"images.with'dot",
|
||||
column="vector'field",
|
||||
distance_threshold=4,
|
||||
dataset_version=7,
|
||||
),
|
||||
with_no_data=with_no_data,
|
||||
)
|
||||
assert job.id == "mv-create-123"
|
||||
assert requests == [
|
||||
(
|
||||
"/v1/materialized_view/clean/create",
|
||||
{
|
||||
"query": (
|
||||
"SELECT * FROM vector_dedup('`images.with''dot`', "
|
||||
"7, 'vector''field', 4.0)"
|
||||
),
|
||||
"with_no_data": with_no_data,
|
||||
},
|
||||
)
|
||||
]
|
||||
def test_local_materialized_view_listing_pages_and_resume(tmp_path):
|
||||
db = make_db(tmp_path)
|
||||
db.create_materialized_view("a_view", "people", with_no_data=True)
|
||||
db.create_materialized_view("z_view", "people", with_no_data=True)
|
||||
db.create_table("ordinary", [{"id": 1}])
|
||||
names = db.list_materialized_views(page_limit=1)
|
||||
assert names.num_page_results() == 0
|
||||
first = next(names)
|
||||
assert names.num_page_results() == 0
|
||||
assert names.page_token() is not None
|
||||
resumed = db.list_materialized_views(page_token=names.page_token(), page_limit=1)
|
||||
assert sorted([first, *resumed]) == ["a_view", "z_view"]
|
||||
assert resumed.page_token() is None
|
||||
assert list(resumed) == []
|
||||
|
||||
+61
-61
@@ -1,8 +1,8 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Tests for installing and clearing an LsmWriteSpec via
|
||||
`Table.set_lsm_write_spec` / `Table.unset_lsm_write_spec`.
|
||||
"""Tests for installing and clearing an MemWalWriteSpec via
|
||||
`Table.enable_memwal` / `Table.disable_memwal`.
|
||||
"""
|
||||
|
||||
from datetime import timedelta
|
||||
@@ -10,7 +10,7 @@ from datetime import timedelta
|
||||
import lancedb
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from lancedb._lancedb import LsmWriteSpec
|
||||
from lancedb._lancedb import MemWalWriteSpec
|
||||
from lancedb.index import BTree
|
||||
|
||||
SCHEMA = pa.schema(
|
||||
@@ -21,9 +21,9 @@ SCHEMA = pa.schema(
|
||||
)
|
||||
|
||||
|
||||
def test_lsm_write_spec_module_metadata():
|
||||
assert lancedb.LsmWriteSpec is LsmWriteSpec
|
||||
assert LsmWriteSpec.__module__ == "lancedb._lancedb"
|
||||
def test_memwal_spec_module_metadata():
|
||||
assert lancedb.MemWalWriteSpec is MemWalWriteSpec
|
||||
assert MemWalWriteSpec.__module__ == "lancedb._lancedb"
|
||||
|
||||
|
||||
def _batch(ids, vs):
|
||||
@@ -43,35 +43,35 @@ def _make_table(tmp_path):
|
||||
return db, table
|
||||
|
||||
|
||||
def test_set_lsm_write_spec_validates(tmp_path):
|
||||
def test_enable_memwal_validates(tmp_path):
|
||||
_db, table = _make_table(tmp_path)
|
||||
|
||||
# Out-of-range num_buckets.
|
||||
with pytest.raises(Exception, match="num_buckets"):
|
||||
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 0))
|
||||
table.enable_memwal(MemWalWriteSpec.bucket("id", 0))
|
||||
with pytest.raises(Exception, match="num_buckets"):
|
||||
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 1025))
|
||||
table.enable_memwal(MemWalWriteSpec.bucket("id", 1025))
|
||||
|
||||
# Happy path then mutation rejected.
|
||||
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 4))
|
||||
with pytest.raises(Exception, match="mutation"):
|
||||
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 8))
|
||||
# Happy path, then a set over the installed spec is refused.
|
||||
table.enable_memwal(MemWalWriteSpec.bucket("id", 4))
|
||||
with pytest.raises(Exception, match="already set"):
|
||||
table.enable_memwal(MemWalWriteSpec.bucket("id", 8))
|
||||
|
||||
|
||||
def test_unset_lsm_write_spec(tmp_path):
|
||||
def test_disable_memwal(tmp_path):
|
||||
_db, table = _make_table(tmp_path)
|
||||
|
||||
# unset errors when no spec is set.
|
||||
with pytest.raises(Exception, match="no LSM write spec"):
|
||||
table.unset_lsm_write_spec()
|
||||
with pytest.raises(Exception, match="no MemWAL write spec"):
|
||||
table.disable_memwal()
|
||||
|
||||
# Install a spec, then remove it; afterwards a fresh spec can be set.
|
||||
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 4))
|
||||
table.unset_lsm_write_spec()
|
||||
table.enable_memwal(MemWalWriteSpec.bucket("id", 4))
|
||||
table.disable_memwal()
|
||||
# A second unset errors — there is no spec left to remove.
|
||||
with pytest.raises(Exception, match="no LSM write spec"):
|
||||
table.unset_lsm_write_spec()
|
||||
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 8))
|
||||
with pytest.raises(Exception, match="no MemWAL write spec"):
|
||||
table.disable_memwal()
|
||||
table.enable_memwal(MemWalWriteSpec.bucket("id", 8))
|
||||
|
||||
|
||||
def test_set_unsharded_spec(tmp_path):
|
||||
@@ -79,12 +79,12 @@ def test_set_unsharded_spec(tmp_path):
|
||||
# Lance MemWAL still requires a primary key on the dataset; Unsharded
|
||||
# just skips per-row hashing.
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
|
||||
table.unset_lsm_write_spec()
|
||||
table.enable_memwal(MemWalWriteSpec.unsharded())
|
||||
table.disable_memwal()
|
||||
|
||||
|
||||
def test_lsm_write_spec_repr():
|
||||
s = LsmWriteSpec.bucket("id", 4)
|
||||
def test_memwal_spec_repr():
|
||||
s = MemWalWriteSpec.bucket("id", 4)
|
||||
assert s.spec_type == "bucket"
|
||||
assert s.column == "id"
|
||||
assert s.num_buckets == 4
|
||||
@@ -95,20 +95,20 @@ def test_lsm_write_spec_repr():
|
||||
assert "id" in repr(s)
|
||||
assert "4" in repr(s)
|
||||
|
||||
u = LsmWriteSpec.unsharded()
|
||||
u = MemWalWriteSpec.unsharded()
|
||||
assert u.spec_type == "unsharded"
|
||||
assert u.column is None
|
||||
assert u.num_buckets is None
|
||||
assert "unsharded" in repr(u)
|
||||
|
||||
|
||||
def test_lsm_write_spec_with_maintained_indexes():
|
||||
s = LsmWriteSpec.bucket("id", 4).with_maintained_indexes(["idx_a", "idx_b"])
|
||||
def test_memwal_spec_with_maintained_indexes():
|
||||
s = MemWalWriteSpec.bucket("id", 4).with_maintained_indexes(["idx_a", "idx_b"])
|
||||
assert s.maintained_indexes == ["idx_a", "idx_b"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_set_unset_lsm_write_spec(tmp_path):
|
||||
async def test_async_set_disable_memwal(tmp_path):
|
||||
db = await lancedb.connect_async(
|
||||
tmp_path, read_consistency_interval=timedelta(seconds=0)
|
||||
)
|
||||
@@ -118,11 +118,11 @@ async def test_async_set_unset_lsm_write_spec(tmp_path):
|
||||
)
|
||||
|
||||
await table.set_unenforced_primary_key("id")
|
||||
await table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 4))
|
||||
await table.unset_lsm_write_spec()
|
||||
await table.enable_memwal(MemWalWriteSpec.bucket("id", 4))
|
||||
await table.disable_memwal()
|
||||
# A second unset errors.
|
||||
with pytest.raises(Exception, match="no LSM write spec"):
|
||||
await table.unset_lsm_write_spec()
|
||||
with pytest.raises(Exception, match="no MemWAL write spec"):
|
||||
await table.disable_memwal()
|
||||
|
||||
|
||||
def test_set_identity_spec(tmp_path):
|
||||
@@ -130,12 +130,12 @@ def test_set_identity_spec(tmp_path):
|
||||
# Identity sharding still requires an unenforced primary key on the
|
||||
# table; it shards by the raw value of the given column.
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.set_lsm_write_spec(LsmWriteSpec.identity("v"))
|
||||
table.unset_lsm_write_spec()
|
||||
table.enable_memwal(MemWalWriteSpec.identity("v"))
|
||||
table.disable_memwal()
|
||||
|
||||
|
||||
def test_lsm_write_spec_identity_and_writer_config_defaults():
|
||||
s = LsmWriteSpec.identity("v")
|
||||
def test_memwal_spec_identity_and_writer_config_defaults():
|
||||
s = MemWalWriteSpec.identity("v")
|
||||
assert s.spec_type == "identity"
|
||||
assert s.column == "v"
|
||||
assert s.num_buckets is None
|
||||
@@ -146,12 +146,12 @@ def test_lsm_write_spec_identity_and_writer_config_defaults():
|
||||
assert "durable_write" in repr(s)
|
||||
|
||||
|
||||
def test_get_lsm_write_spec(tmp_path):
|
||||
def test_get_memwal_spec(tmp_path):
|
||||
_db, table = _make_table(tmp_path)
|
||||
table.set_unenforced_primary_key("id")
|
||||
|
||||
# None when nothing is installed.
|
||||
assert table.get_lsm_write_spec() is None
|
||||
assert table.get_memwal_spec() is None
|
||||
|
||||
# A real scalar index is needed to name it as a maintained index.
|
||||
table.create_index("id", config=BTree())
|
||||
@@ -159,12 +159,12 @@ def test_get_lsm_write_spec(tmp_path):
|
||||
|
||||
# Bucket spec round-trips, including maintained indexes and writer config
|
||||
# defaults.
|
||||
table.set_lsm_write_spec(
|
||||
LsmWriteSpec.bucket("id", 4)
|
||||
table.enable_memwal(
|
||||
MemWalWriteSpec.bucket("id", 4)
|
||||
.with_maintained_indexes([idx_name])
|
||||
.with_writer_config_defaults({"durable_write": "false"})
|
||||
)
|
||||
spec = table.get_lsm_write_spec()
|
||||
spec = table.get_memwal_spec()
|
||||
assert spec is not None
|
||||
assert spec.spec_type == "bucket"
|
||||
assert spec.column == "id"
|
||||
@@ -173,30 +173,30 @@ def test_get_lsm_write_spec(tmp_path):
|
||||
assert spec.writer_config_defaults == {"durable_write": "false"}
|
||||
|
||||
# After unset, None again.
|
||||
table.unset_lsm_write_spec()
|
||||
assert table.get_lsm_write_spec() is None
|
||||
table.disable_memwal()
|
||||
assert table.get_memwal_spec() is None
|
||||
|
||||
# Identity round-trips (column recovered from the schema). Leaving the
|
||||
# maintained set to be inferred picks up the index on the table, so the
|
||||
# spec reads back naming it rather than as "infer".
|
||||
table.set_lsm_write_spec(LsmWriteSpec.identity("id"))
|
||||
spec = table.get_lsm_write_spec()
|
||||
# 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.enable_memwal(MemWalWriteSpec.identity("id"))
|
||||
spec = table.get_memwal_spec()
|
||||
assert spec.spec_type == "identity"
|
||||
assert spec.column == "id"
|
||||
assert spec.maintained_indexes == [idx_name]
|
||||
table.unset_lsm_write_spec()
|
||||
assert spec.maintained_indexes is None
|
||||
table.disable_memwal()
|
||||
|
||||
# Unsharded round-trips (no routing column). Opting out is distinct from
|
||||
# the inferred default.
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([]))
|
||||
spec = table.get_lsm_write_spec()
|
||||
# the unnamed default.
|
||||
table.enable_memwal(MemWalWriteSpec.unsharded().with_maintained_indexes([]))
|
||||
spec = table.get_memwal_spec()
|
||||
assert spec.spec_type == "unsharded"
|
||||
assert spec.column is None
|
||||
assert spec.maintained_indexes == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_get_lsm_write_spec(tmp_path):
|
||||
async def test_async_get_memwal_spec(tmp_path):
|
||||
db = await lancedb.connect_async(
|
||||
tmp_path, read_consistency_interval=timedelta(seconds=0)
|
||||
)
|
||||
@@ -205,20 +205,20 @@ async def test_async_get_lsm_write_spec(tmp_path):
|
||||
pa.RecordBatchReader.from_batches(SCHEMA, [_batch(["seed"], [0])]),
|
||||
)
|
||||
|
||||
assert await table.get_lsm_write_spec() is None
|
||||
assert await table.get_memwal_spec() is None
|
||||
|
||||
# A real scalar index is needed to name it as a maintained index.
|
||||
await table.create_index("id", config=BTree())
|
||||
idx_name = (await table.list_indices())[0].name
|
||||
|
||||
await table.set_lsm_write_spec(
|
||||
LsmWriteSpec.bucket("id", 8).with_maintained_indexes([idx_name])
|
||||
await table.enable_memwal(
|
||||
MemWalWriteSpec.bucket("id", 8).with_maintained_indexes([idx_name])
|
||||
)
|
||||
spec = await table.get_lsm_write_spec()
|
||||
spec = await table.get_memwal_spec()
|
||||
assert spec is not None
|
||||
assert spec.spec_type == "bucket"
|
||||
assert spec.column == "id"
|
||||
assert spec.num_buckets == 8
|
||||
assert spec.maintained_indexes == [idx_name]
|
||||
await table.unset_lsm_write_spec()
|
||||
assert await table.get_lsm_write_spec() is None
|
||||
await table.disable_memwal()
|
||||
assert await table.get_memwal_spec() is None
|
||||
@@ -8,7 +8,7 @@ from datetime import timedelta
|
||||
import lancedb
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from lancedb._lancedb import LsmWriteSpec
|
||||
from lancedb._lancedb import MemWalWriteSpec
|
||||
from lancedb.index import FTS, IvfPq
|
||||
|
||||
SCHEMA = pa.schema(
|
||||
@@ -49,12 +49,12 @@ def _region_reader(rows):
|
||||
|
||||
|
||||
def _bucket_table(tmp_path):
|
||||
"""A table with ``id`` as the primary key and a single-bucket LSM spec."""
|
||||
"""A table with ``id`` as the primary key and a single-bucket MemWAL spec."""
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _reader([1, 2, 3]))
|
||||
table.set_unenforced_primary_key("id")
|
||||
# num_buckets = 1: every row routes to the single bucket.
|
||||
table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 1))
|
||||
table.enable_memwal(MemWalWriteSpec.bucket("id", 1))
|
||||
return table
|
||||
|
||||
|
||||
@@ -78,7 +78,7 @@ def test_lsm_merge_insert_unsharded(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _reader([1, 2, 3]))
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
|
||||
table.enable_memwal(MemWalWriteSpec.unsharded())
|
||||
result = (
|
||||
table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
@@ -92,7 +92,7 @@ def test_lsm_merge_insert_identity(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _region_reader([(1, "us"), (2, "us")]))
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.set_lsm_write_spec(LsmWriteSpec.identity("region"))
|
||||
table.enable_memwal(MemWalWriteSpec.identity("region"))
|
||||
# All rows share one identity value, so they route to one shard.
|
||||
result = (
|
||||
table.merge_insert([])
|
||||
@@ -103,31 +103,31 @@ def test_lsm_merge_insert_identity(tmp_path):
|
||||
assert result.num_rows == 2
|
||||
|
||||
|
||||
def test_lsm_merge_insert_use_lsm_false(tmp_path):
|
||||
def test_lsm_merge_insert_use_memwal_false(tmp_path):
|
||||
table = _bucket_table(tmp_path) # rows id = 1, 2, 3
|
||||
# use_lsm(False) opts out: the standard path runs and commits even with a spec.
|
||||
# use_memwal(False) opts out: the standard path runs and commits even with a spec.
|
||||
result = (
|
||||
table.merge_insert("id")
|
||||
.when_not_matched_insert_all()
|
||||
.use_lsm(False)
|
||||
.use_memwal(False)
|
||||
.execute(_reader([3, 4, 5]))
|
||||
)
|
||||
assert result.num_inserted_rows == 2
|
||||
assert table.count_rows() == 5
|
||||
|
||||
|
||||
def test_lsm_merge_insert_use_lsm_true_without_spec_errors(tmp_path):
|
||||
# A table with a primary key but no LSM write spec installed.
|
||||
def test_lsm_merge_insert_use_memwal_true_without_spec_errors(tmp_path):
|
||||
# A table with a primary key but no MemWAL write spec installed.
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _reader([1, 2, 3]))
|
||||
table.set_unenforced_primary_key("id")
|
||||
# use_lsm(True) demands MemWAL routing; without a spec it errors.
|
||||
with pytest.raises(Exception, match="use_lsm"):
|
||||
# use_memwal(True) demands MemWAL routing; without a spec it errors.
|
||||
with pytest.raises(Exception, match="use_memwal"):
|
||||
(
|
||||
table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.use_lsm(True)
|
||||
.use_memwal(True)
|
||||
.execute(_reader([3, 4, 5]))
|
||||
)
|
||||
|
||||
@@ -145,7 +145,7 @@ def test_lsm_merge_insert_validate_single_shard_off(tmp_path):
|
||||
|
||||
|
||||
def test_lsm_merge_insert_no_spec_uses_standard_path(tmp_path):
|
||||
# A table with a primary key but no LSM write spec installed.
|
||||
# A table with a primary key but no MemWAL write spec installed.
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _reader([1, 2, 3]))
|
||||
table.set_unenforced_primary_key("id")
|
||||
@@ -186,7 +186,7 @@ def test_lsm_close_writers(tmp_path):
|
||||
.when_not_matched_insert_all()
|
||||
.execute(_reader([7, 8]))
|
||||
)
|
||||
table.close_lsm_writers()
|
||||
table.close_memwal_writers()
|
||||
# The writer reopens lazily on the next merge_insert.
|
||||
result = (
|
||||
table.merge_insert([])
|
||||
@@ -204,14 +204,14 @@ async def test_async_lsm_merge_insert(tmp_path):
|
||||
)
|
||||
table = await db.create_table("t", _reader([1, 2, 3]))
|
||||
await table.set_unenforced_primary_key("id")
|
||||
await table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 1))
|
||||
await table.enable_memwal(MemWalWriteSpec.bucket("id", 1))
|
||||
|
||||
builder = (
|
||||
table.merge_insert([]).when_matched_update_all().when_not_matched_insert_all()
|
||||
)
|
||||
result = await builder.execute(_reader([3, 4, 5]))
|
||||
assert result.num_rows == 3
|
||||
await table.close_lsm_writers()
|
||||
await table.close_memwal_writers()
|
||||
|
||||
|
||||
def _lsm_upsert(table, ids):
|
||||
@@ -228,7 +228,7 @@ def test_lsm_read_sees_active_memtable(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _reader([1, 2, 3])) # base ids 1,2,3
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
|
||||
table.enable_memwal(MemWalWriteSpec.unsharded())
|
||||
|
||||
_lsm_upsert(table, [4, 5]) # active memtable only, not committed to base
|
||||
|
||||
@@ -236,8 +236,8 @@ def test_lsm_read_sees_active_memtable(tmp_path):
|
||||
lsm = table.search().to_arrow()
|
||||
assert sorted(lsm["id"].to_pylist()) == [1, 2, 3, 4, 5]
|
||||
|
||||
# use_lsm(False) bypasses the MemWAL and reads the base table only.
|
||||
base_only = table.search().use_lsm(False).to_arrow()
|
||||
# use_memwal(False) bypasses the MemWAL and reads the base table only.
|
||||
base_only = table.search().use_memwal(False).to_arrow()
|
||||
assert sorted(base_only["id"].to_pylist()) == [1, 2, 3]
|
||||
|
||||
|
||||
@@ -245,7 +245,7 @@ def test_lsm_read_dedup_newest_wins(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _reader([1, 2, 3])) # id 2 -> value 1
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
|
||||
table.enable_memwal(MemWalWriteSpec.unsharded())
|
||||
|
||||
_lsm_upsert(table, [2, 3, 4]) # ids 2,3,4 -> values 0,1,2
|
||||
|
||||
@@ -258,22 +258,22 @@ def test_lsm_read_dedup_newest_wins(tmp_path):
|
||||
def test_lsm_read_without_spec_reads_base(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _reader([1, 2, 3]))
|
||||
table.set_unenforced_primary_key("id") # no LSM write spec
|
||||
table.set_unenforced_primary_key("id") # no MemWAL write spec
|
||||
|
||||
# No spec: default read and use_lsm(False) both read the base table, no error.
|
||||
# No spec: default read and use_memwal(False) both read the base table, no error.
|
||||
assert sorted(table.search().to_arrow()["id"].to_pylist()) == [1, 2, 3]
|
||||
assert sorted(table.search().use_lsm(False).to_arrow()["id"].to_pylist()) == [
|
||||
assert sorted(table.search().use_memwal(False).to_arrow()["id"].to_pylist()) == [
|
||||
1,
|
||||
2,
|
||||
3,
|
||||
]
|
||||
|
||||
|
||||
def test_lsm_read_unsupported_shape_errors_without_use_lsm_false(tmp_path):
|
||||
def test_lsm_read_unsupported_shape_errors_without_use_memwal_false(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _reader([1, 2, 3]))
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
|
||||
table.enable_memwal(MemWalWriteSpec.unsharded())
|
||||
_lsm_upsert(table, [4])
|
||||
|
||||
# with_row_id is unsupported by the LSM scanner; on a MemWAL table the default
|
||||
@@ -281,8 +281,8 @@ def test_lsm_read_unsupported_shape_errors_without_use_lsm_false(tmp_path):
|
||||
with pytest.raises(Exception):
|
||||
table.search().with_row_id(True).to_arrow()
|
||||
|
||||
# use_lsm(False) is the escape hatch: it reads the base table only.
|
||||
base = table.search().with_row_id(True).use_lsm(False).to_arrow()
|
||||
# use_memwal(False) is the escape hatch: it reads the base table only.
|
||||
base = table.search().with_row_id(True).use_memwal(False).to_arrow()
|
||||
assert sorted(base["id"].to_pylist()) == [1, 2, 3]
|
||||
|
||||
|
||||
@@ -293,7 +293,7 @@ async def test_async_lsm_read(tmp_path):
|
||||
)
|
||||
table = await db.create_table("t", _reader([1, 2, 3]))
|
||||
await table.set_unenforced_primary_key("id")
|
||||
await table.set_lsm_write_spec(LsmWriteSpec.unsharded())
|
||||
await table.enable_memwal(MemWalWriteSpec.unsharded())
|
||||
|
||||
builder = (
|
||||
table.merge_insert([]).when_matched_update_all().when_not_matched_insert_all()
|
||||
@@ -331,7 +331,7 @@ def _vector_reader(rows):
|
||||
def _vector_table(tmp_path):
|
||||
"""Base table whose vector column is indexed so its rows are visible to the LSM
|
||||
vector scanner (the base arm uses ``fast_search`` — indexed data only), plus an
|
||||
unsharded LSM spec that maintains that index for the memtable.
|
||||
unsharded MemWAL spec that maintains that index for the memtable.
|
||||
|
||||
Rows 1,2 are category ``a``, row 3 is ``b``, and 4..60 are filler ``c`` that
|
||||
give the tiny IVF index enough data to train.
|
||||
@@ -353,8 +353,8 @@ def _vector_table(tmp_path):
|
||||
"vector", config=IvfPq(num_partitions=1, num_sub_vectors=2, num_bits=4)
|
||||
)
|
||||
index_name = table.list_indices()[0].name
|
||||
table.set_lsm_write_spec(
|
||||
LsmWriteSpec.unsharded().with_maintained_indexes([index_name])
|
||||
table.enable_memwal(
|
||||
MemWalWriteSpec.unsharded().with_maintained_indexes([index_name])
|
||||
)
|
||||
return table
|
||||
|
||||
@@ -379,9 +379,9 @@ def test_lsm_read_vector_sees_memtable(tmp_path):
|
||||
assert {1, 2, 3} <= ids # indexed base rows
|
||||
assert 1000 in ids # in-flight memtable row
|
||||
|
||||
# use_lsm(False) bypasses the MemWAL, so the in-flight row is not visible.
|
||||
# use_memwal(False) bypasses the MemWAL, so the in-flight row is not visible.
|
||||
base_ids = set(
|
||||
table.search(query).use_lsm(False).limit(100).to_arrow()["id"].to_pylist()
|
||||
table.search(query).use_memwal(False).limit(100).to_arrow()["id"].to_pylist()
|
||||
)
|
||||
assert {1, 2, 3} <= base_ids
|
||||
assert 1000 not in base_ids
|
||||
@@ -453,8 +453,8 @@ def test_lsm_read_fts_sees_memtable(tmp_path):
|
||||
# Native FTS index (tantivy is not compatible with the LSM memtable index).
|
||||
table.create_index("text", config=FTS())
|
||||
index_name = table.list_indices()[0].name
|
||||
table.set_lsm_write_spec(
|
||||
LsmWriteSpec.unsharded().with_maintained_indexes([index_name])
|
||||
table.enable_memwal(
|
||||
MemWalWriteSpec.unsharded().with_maintained_indexes([index_name])
|
||||
)
|
||||
|
||||
# in-flight doc 4 lands in the memtable's maintained FTS index.
|
||||
@@ -498,8 +498,8 @@ def test_lsm_read_vector_unsupported_knobs_error(tmp_path):
|
||||
with pytest.raises(Exception, match="use_index"):
|
||||
table.search(query).bypass_vector_index().to_arrow()
|
||||
|
||||
# use_lsm(False) is the escape hatch: the base-only standard path honors them.
|
||||
base = table.search(query).distance_range(0.0, 100.0).use_lsm(False).to_arrow()
|
||||
# use_memwal(False) is the escape hatch: the base-only standard path honors them.
|
||||
base = table.search(query).distance_range(0.0, 100.0).use_memwal(False).to_arrow()
|
||||
assert 1000 not in set(base["id"].to_pylist())
|
||||
|
||||
|
||||
@@ -544,7 +544,7 @@ def test_lsm_read_fts_unmaintained_index_errors(tmp_path):
|
||||
table.create_index("text", config=FTS())
|
||||
# No maintained indexes: the active memtable FTS arm cannot serve un-compacted
|
||||
# docs, so the search would silently omit them — reject instead.
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([]))
|
||||
table.enable_memwal(MemWalWriteSpec.unsharded().with_maintained_indexes([]))
|
||||
with pytest.raises(Exception, match="maintained"):
|
||||
table.search("fox", query_type="fts", fts_columns="text").to_arrow()
|
||||
|
||||
@@ -553,7 +553,7 @@ def test_lsm_read_time_travel_errors(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _reader([1, 2, 3]))
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
|
||||
table.enable_memwal(MemWalWriteSpec.unsharded())
|
||||
pinned = table.version
|
||||
table.add(_reader([4, 5])) # standard add commits a newer version
|
||||
table.checkout(pinned) # detached head at the historical version
|
||||
@@ -562,8 +562,8 @@ def test_lsm_read_time_travel_errors(tmp_path):
|
||||
# historical version is rejected.
|
||||
with pytest.raises(Exception, match="time-travel"):
|
||||
table.search().to_arrow()
|
||||
# use_lsm(False) reads the base table at the pinned version.
|
||||
base = table.search().use_lsm(False).to_arrow()
|
||||
# use_memwal(False) reads the base table at the pinned version.
|
||||
base = table.search().use_memwal(False).to_arrow()
|
||||
assert sorted(base["id"].to_pylist()) == [1, 2, 3]
|
||||
|
||||
|
||||
@@ -571,14 +571,14 @@ def test_lsm_read_take_row_ids_errors(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _reader([1, 2, 3]))
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
|
||||
table.enable_memwal(MemWalWriteSpec.unsharded())
|
||||
_lsm_upsert(table, [4])
|
||||
# take-by-row-id auto-routes through the LSM scanner, which has no stable _rowid,
|
||||
# so it hard-errors instead of failing with an opaque column-not-found error.
|
||||
with pytest.raises(Exception, match="row id"):
|
||||
table.take_row_ids([0, 1]).to_arrow()
|
||||
# use_lsm(False) is the escape hatch: it reads the base table.
|
||||
base = table.take_row_ids([0, 1]).use_lsm(False).to_arrow()
|
||||
# use_memwal(False) is the escape hatch: it reads the base table.
|
||||
base = table.take_row_ids([0, 1]).use_memwal(False).to_arrow()
|
||||
assert base.num_rows == 2
|
||||
|
||||
|
||||
@@ -588,8 +588,8 @@ def test_lsm_read_fts_postfilter_errors(tmp_path):
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.create_index("text", config=FTS())
|
||||
index_name = table.list_indices()[0].name
|
||||
table.set_lsm_write_spec(
|
||||
LsmWriteSpec.unsharded().with_maintained_indexes([index_name])
|
||||
table.enable_memwal(
|
||||
MemWalWriteSpec.unsharded().with_maintained_indexes([index_name])
|
||||
)
|
||||
# The LSM scanner always prefilters; postfilter on FTS changes result semantics,
|
||||
# so it hard-errors (previously only the vector arm rejected it).
|
||||
@@ -608,9 +608,7 @@ def test_lsm_read_fts_multiple_same_type_indexes_errors(tmp_path):
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.create_index("text", config=FTS(), name="fts_a")
|
||||
table.create_index("text", config=FTS(), name="fts_b", replace=False)
|
||||
table.set_lsm_write_spec(
|
||||
LsmWriteSpec.unsharded().with_maintained_indexes(["fts_a"])
|
||||
)
|
||||
table.enable_memwal(MemWalWriteSpec.unsharded().with_maintained_indexes(["fts_a"]))
|
||||
# Two FTS indexes on the column: the base planner's chosen index is ambiguous, so
|
||||
# the scanner can't pick a catch-up watermark and rejects rather than risk
|
||||
# dropping rows the actually-used index has not caught up to.
|
||||
@@ -631,7 +629,7 @@ def test_lsm_read_vector_unmaintained_index_errors(tmp_path):
|
||||
)
|
||||
# Spec with NO maintained indexes: the base vector index's catch-up is untracked,
|
||||
# so the scanner rejects rather than risk dropping compacted-but-unindexed rows.
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([]))
|
||||
table.enable_memwal(MemWalWriteSpec.unsharded().with_maintained_indexes([]))
|
||||
with pytest.raises(Exception, match="maintained"):
|
||||
table.search([1.0] * VECTOR_DIM).to_arrow()
|
||||
|
||||
@@ -644,7 +642,7 @@ def test_lsm_read_fts_optimized_index_not_rejected(tmp_path):
|
||||
table.add(_fts_reader([(i, "lazy fox") for i in range(6, 11)]))
|
||||
table.optimize() # may split the FTS index into multiple physical segments
|
||||
name = table.list_indices()[0].name
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([name]))
|
||||
table.enable_memwal(MemWalWriteSpec.unsharded().with_maintained_indexes([name]))
|
||||
# Multiple physical segments of one logical index must not be miscounted as
|
||||
# multiple indexes and rejected.
|
||||
ids = set(
|
||||
|
||||
@@ -6,11 +6,17 @@
|
||||
import tempfile
|
||||
import shutil
|
||||
import importlib
|
||||
import http.server
|
||||
import json
|
||||
import threading
|
||||
from urllib.parse import parse_qs, unquote, urlparse
|
||||
import pytest
|
||||
import pyarrow as pa
|
||||
import lancedb
|
||||
from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
|
||||
from lancedb.namespace import _MAX_QUERY_K
|
||||
from lancedb.listing import AsyncListing, Listing
|
||||
from lancedb.namespace import _MAX_QUERY_K, _query_to_namespace_request
|
||||
from lancedb.query import Query
|
||||
from lancedb.table import AsyncTable, LanceTable
|
||||
|
||||
|
||||
@@ -72,6 +78,132 @@ def _namespace_lance_table(namespace_client: _NamespaceClient) -> LanceTable:
|
||||
return table
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("namespace_path", [[], ["test_ns"]])
|
||||
async def test_namespace_table_names_defaults_and_limits(tmp_path, namespace_path):
|
||||
db = lancedb.connect_namespace("dir", {"root": str(tmp_path)})
|
||||
if namespace_path:
|
||||
db.create_namespace(namespace_path)
|
||||
names = [f"t{idx:03}" for idx in range(12)]
|
||||
schema = pa.schema([pa.field("id", pa.int64())])
|
||||
for name in names:
|
||||
db.create_table(name, schema=schema, namespace_path=namespace_path)
|
||||
|
||||
assert list(db.table_names(namespace_path=namespace_path)) == names
|
||||
assert list(db.table_names(limit=None, namespace_path=namespace_path)) == names
|
||||
assert list(db.table_names(limit=0, namespace_path=namespace_path)) == []
|
||||
assert list(db.table_names(limit=2, namespace_path=namespace_path)) == names[:2]
|
||||
assert (
|
||||
list(
|
||||
db.table_names(page_token=names[1], limit=2, namespace_path=namespace_path)
|
||||
)
|
||||
== names[2:4]
|
||||
)
|
||||
|
||||
adb = lancedb.connect_namespace_async("dir", {"root": str(tmp_path)})
|
||||
assert list(await adb.table_names(namespace_path=namespace_path)) == names
|
||||
assert (
|
||||
list(await adb.table_names(limit=None, namespace_path=namespace_path)) == names
|
||||
)
|
||||
assert list(await adb.table_names(limit=0, namespace_path=namespace_path)) == []
|
||||
assert (
|
||||
list(await adb.table_names(limit=2, namespace_path=namespace_path)) == names[:2]
|
||||
)
|
||||
assert (
|
||||
list(
|
||||
await adb.table_names(
|
||||
page_token=names[1], limit=2, namespace_path=namespace_path
|
||||
)
|
||||
)
|
||||
== names[2:4]
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def paginated_rest_namespace():
|
||||
names = [f"t{idx:03}" for idx in reversed(range(12))]
|
||||
requests = []
|
||||
|
||||
class Handler(http.server.BaseHTTPRequestHandler):
|
||||
def do_GET(self):
|
||||
url = urlparse(self.path)
|
||||
query = parse_qs(url.query)
|
||||
requests.append((unquote(url.path), query))
|
||||
token = query.get("page_token", [None])[0]
|
||||
if token not in {None, "opaque-next"}:
|
||||
self.send_error(400, "Unknown page token")
|
||||
return
|
||||
start = 0 if token is None else 10
|
||||
body = {"tables": names[start : start + 10]}
|
||||
if start == 0:
|
||||
body["page_token"] = "opaque-next"
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.end_headers()
|
||||
self.wfile.write(json.dumps(body).encode())
|
||||
|
||||
def log_message(self, *args):
|
||||
pass
|
||||
|
||||
with http.server.HTTPServer(("127.0.0.1", 0), Handler) as server:
|
||||
thread = threading.Thread(target=server.serve_forever)
|
||||
thread.start()
|
||||
try:
|
||||
yield {"uri": f"http://127.0.0.1:{server.server_port}"}, requests
|
||||
finally:
|
||||
server.shutdown()
|
||||
thread.join()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("namespace_path", [[], ["n"]])
|
||||
async def test_rest_namespace_table_names_collects_pages(
|
||||
paginated_rest_namespace, namespace_path
|
||||
):
|
||||
properties, requests = paginated_rest_namespace
|
||||
names = [f"t{idx:03}" for idx in range(12)]
|
||||
db = lancedb.connect_namespace("rest", properties)
|
||||
first = db.list_tables(namespace_path=namespace_path)
|
||||
second = db.list_tables(namespace_path=namespace_path, page_token=first.page_token)
|
||||
assert len(first.tables) == 10
|
||||
assert first.page_token == "opaque-next"
|
||||
assert sorted(first.tables + second.tables) == names
|
||||
assert second.page_token is None
|
||||
|
||||
assert list(db.table_names(namespace_path=namespace_path)) == names
|
||||
assert list(db.table_names(limit=None, namespace_path=namespace_path)) == names
|
||||
assert list(db.table_names(limit=11, namespace_path=namespace_path)) == names[:11]
|
||||
assert (
|
||||
list(db.table_names(page_token="t005", limit=2, namespace_path=namespace_path))
|
||||
== names[6:8]
|
||||
)
|
||||
|
||||
adb = lancedb.connect_namespace_async("rest", properties)
|
||||
assert list(await adb.table_names(namespace_path=namespace_path)) == names
|
||||
assert (
|
||||
list(await adb.table_names(limit=None, namespace_path=namespace_path)) == names
|
||||
)
|
||||
assert (
|
||||
list(await adb.table_names(limit=11, namespace_path=namespace_path))
|
||||
== names[:11]
|
||||
)
|
||||
assert (
|
||||
list(
|
||||
await adb.table_names(
|
||||
page_token="t005", limit=2, namespace_path=namespace_path
|
||||
)
|
||||
)
|
||||
== names[6:8]
|
||||
)
|
||||
|
||||
request_count = len(requests)
|
||||
assert list(db.table_names(limit=0, namespace_path=namespace_path)) == []
|
||||
assert list(await adb.table_names(limit=0, namespace_path=namespace_path)) == []
|
||||
assert len(requests) == request_count
|
||||
identifier = "$" if not namespace_path else "n"
|
||||
assert all(path == f"/v1/namespace/{identifier}/table/list" for path, _ in requests)
|
||||
|
||||
|
||||
class TestNamespaceConnection:
|
||||
"""Test namespace-based LanceDB connection using DirectoryNamespace."""
|
||||
|
||||
@@ -104,6 +236,18 @@ class TestNamespaceConnection:
|
||||
assert isinstance(db, lancedb.LanceNamespaceDBConnection)
|
||||
assert len(list(db.table_names())) == 0
|
||||
|
||||
def test_list_materialized_views(self):
|
||||
db = lancedb.connect_namespace("dir", {"root": self.temp_dir})
|
||||
db.create_table("source", [{"id": 1}])
|
||||
db.create_materialized_view("view", "source", with_no_data=True)
|
||||
|
||||
names = db.list_materialized_views(page_limit=1)
|
||||
assert isinstance(names, Listing)
|
||||
assert names.num_page_results() == 0
|
||||
assert list(names) == ["view"]
|
||||
assert names.num_page_results() == 0
|
||||
assert names.page_token() is None
|
||||
|
||||
def test_sync_builtin_namespace_uses_rust_without_python_client(self, monkeypatch):
|
||||
"""Built-in sync namespace connections should not construct or call the
|
||||
Python namespace client for normal namespace/table management."""
|
||||
@@ -605,6 +749,18 @@ class TestAsyncNamespaceConnection:
|
||||
table_names = await db.table_names()
|
||||
assert len(list(table_names)) == 0
|
||||
|
||||
async def test_list_materialized_views(self):
|
||||
db = lancedb.connect_namespace_async("dir", {"root": self.temp_dir})
|
||||
await db.create_table("source", [{"id": 1}])
|
||||
await db.create_materialized_view("view", "source", with_no_data=True)
|
||||
|
||||
names = db.list_materialized_views(page_limit=1)
|
||||
assert isinstance(names, AsyncListing)
|
||||
assert names.num_page_results() == 0
|
||||
assert [name async for name in names] == ["view"]
|
||||
assert names.num_page_results() == 0
|
||||
assert names.page_token() is None
|
||||
|
||||
async def test_async_builtin_namespace_uses_rust_without_python_client(
|
||||
self, monkeypatch
|
||||
):
|
||||
@@ -956,6 +1112,29 @@ class TestPushdownOperations:
|
||||
]
|
||||
assert all(r.k <= 2**31 - 1 for r in namespace_client.requests)
|
||||
|
||||
def test_probe_fields_pass_through_independently(self):
|
||||
request = _query_to_namespace_request(
|
||||
["geneva", "hist"],
|
||||
Query(
|
||||
vector=[1.0, 2.0],
|
||||
nprobes=20,
|
||||
minimum_nprobes=3,
|
||||
maximum_nprobes=10,
|
||||
),
|
||||
)
|
||||
|
||||
assert request.nprobes == 20
|
||||
assert request.minimum_nprobes == 3
|
||||
assert request.maximum_nprobes == 10
|
||||
|
||||
request = _query_to_namespace_request(
|
||||
["geneva", "hist"],
|
||||
Query(vector=[1.0, 2.0], nprobes=5, maximum_nprobes=0),
|
||||
)
|
||||
assert request.nprobes == 5
|
||||
assert request.maximum_nprobes == 0
|
||||
assert request.to_dict()["maximum_nprobes"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestAsyncPushdownOperations:
|
||||
@@ -1030,3 +1209,15 @@ def test_local_table_to_arrow_and_to_pandas_are_unchanged(tmp_path):
|
||||
|
||||
assert table.to_arrow().column("id").to_pylist() == [1, 2]
|
||||
assert table.to_pandas()["id"].tolist() == [1, 2]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("effort", [0.0, 0.5, 1.0])
|
||||
def test_namespace_request_rejects_search_effort(effort):
|
||||
from lancedb.query import Query
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="namespace QueryTable does not support search_effort"
|
||||
):
|
||||
_query_to_namespace_request(
|
||||
["test"], Query(vector=[1.0, 0.0], search_effort=effort)
|
||||
)
|
||||
|
||||
@@ -9,6 +9,7 @@ from typing import List, Optional, Tuple
|
||||
import pyarrow as pa
|
||||
import pydantic
|
||||
import pytest
|
||||
from lancedb.conftest import MockTextEmbeddingFunction
|
||||
from lancedb.pydantic import (
|
||||
LanceModel,
|
||||
MultiVector,
|
||||
@@ -425,6 +426,25 @@ def test_bare_vector_raises_clear_error():
|
||||
exec("class TestModel(LanceModel):\n vector: Vector", namespace)
|
||||
|
||||
|
||||
def test_embedding_vector_list_annotation():
|
||||
embedding = MockTextEmbeddingFunction.create()
|
||||
|
||||
class StaticTypingModel(LanceModel):
|
||||
text: str = embedding.SourceField()
|
||||
vector: list[float] = embedding.VectorField()
|
||||
|
||||
schema = pydantic_to_schema(StaticTypingModel)
|
||||
assert schema == pa.schema(
|
||||
[
|
||||
pa.field("text", pa.utf8(), False),
|
||||
pa.field("vector", pa.list_(pa.float32(), embedding.ndims()), True),
|
||||
]
|
||||
)
|
||||
|
||||
model = StaticTypingModel(text="hello", vector=[0.0] * embedding.ndims())
|
||||
assert model.vector == [0.0] * embedding.ndims()
|
||||
|
||||
|
||||
def test_fixed_size_list_field():
|
||||
class TestModel(pydantic.BaseModel):
|
||||
vec: Vector(16)
|
||||
|
||||
@@ -24,6 +24,7 @@ from lancedb.pydantic import LanceModel, Vector
|
||||
from lancedb.query import (
|
||||
AsyncFTSQuery,
|
||||
AsyncHybridQuery,
|
||||
AsyncQuery,
|
||||
AsyncQueryBase,
|
||||
AsyncVectorQuery,
|
||||
ColumnOrdering,
|
||||
@@ -110,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
|
||||
@@ -502,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):
|
||||
@@ -974,6 +1194,19 @@ def test_query_builder_with_filter(table):
|
||||
assert all(np.array(rs[0]["vector"]) == [3, 4])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("query_type", ["vector", "hybrid"])
|
||||
@pytest.mark.parametrize("nprobes", [0, -1])
|
||||
def test_nprobes_nonpositive_sync(table, query_type, nprobes):
|
||||
if query_type == "hybrid":
|
||||
query = table.search(query_type="hybrid").vector([0, 0]).text("a")
|
||||
else:
|
||||
query = table.search([0, 0])
|
||||
with pytest.raises(
|
||||
ValueError, match="^Invalid input, nprobes must be greater than 0$"
|
||||
):
|
||||
query.nprobes(nprobes)
|
||||
|
||||
|
||||
def test_invalid_nprobes_sync(table):
|
||||
with pytest.raises(ValueError, match="minimum_nprobes must be greater than 0"):
|
||||
LanceVectorQueryBuilder(table, [0, 0], "vector").minimum_nprobes(0).to_list()
|
||||
@@ -981,12 +1214,22 @@ def test_invalid_nprobes_sync(table):
|
||||
ValueError,
|
||||
match="maximum_nprobes must be greater than or equal to minimum_nprobes",
|
||||
):
|
||||
LanceVectorQueryBuilder(table, [0, 0], "vector").maximum_nprobes(5).to_list()
|
||||
(
|
||||
LanceVectorQueryBuilder(table, [0, 0], "vector")
|
||||
.minimum_nprobes(10)
|
||||
.maximum_nprobes(5)
|
||||
.to_list()
|
||||
)
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="minimum_nprobes must be less than or equal to maximum_nprobes",
|
||||
match="maximum_nprobes must be greater than or equal to minimum_nprobes",
|
||||
):
|
||||
LanceVectorQueryBuilder(table, [0, 0], "vector").minimum_nprobes(100).to_list()
|
||||
(
|
||||
LanceVectorQueryBuilder(table, [0, 0], "vector")
|
||||
.maximum_nprobes(5)
|
||||
.minimum_nprobes(100)
|
||||
.to_list()
|
||||
)
|
||||
|
||||
|
||||
def test_nprobes_works_sync(table):
|
||||
@@ -1005,6 +1248,18 @@ def test_multiple_nprobes_calls_works_sync(table):
|
||||
).minimum_nprobes(20).to_list()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("hybrid", [False, True])
|
||||
async def test_nprobes_zero_async(table_async: AsyncTable, hybrid):
|
||||
query = table_async.query().nearest_to([0, 0])
|
||||
if hybrid:
|
||||
query = query.nearest_to_text("dog")
|
||||
with pytest.raises(
|
||||
ValueError, match="^Invalid input, nprobes must be greater than 0$"
|
||||
):
|
||||
query.nprobes(0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_nprobes_async(table_async: AsyncTable):
|
||||
with pytest.raises(ValueError, match="minimum_nprobes must be greater than 0"):
|
||||
@@ -1013,12 +1268,22 @@ async def test_invalid_nprobes_async(table_async: AsyncTable):
|
||||
ValueError,
|
||||
match="maximum_nprobes must be greater than or equal to minimum_nprobes",
|
||||
):
|
||||
await table_async.vector_search([0, 0]).maximum_nprobes(5).to_list()
|
||||
await (
|
||||
table_async.vector_search([0, 0])
|
||||
.minimum_nprobes(10)
|
||||
.maximum_nprobes(5)
|
||||
.to_list()
|
||||
)
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="minimum_nprobes must be less than or equal to maximum_nprobes",
|
||||
):
|
||||
await table_async.vector_search([0, 0]).minimum_nprobes(100).to_list()
|
||||
await (
|
||||
table_async.vector_search([0, 0])
|
||||
.maximum_nprobes(5)
|
||||
.minimum_nprobes(100)
|
||||
.to_list()
|
||||
)
|
||||
|
||||
|
||||
def test_query_builder_with_prefilter(table):
|
||||
@@ -1279,6 +1544,27 @@ async def test_query_to_polars_async(table_async: AsyncTable):
|
||||
assert df.shape == (0, num_columns)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("search_kwargs", [{}, {"query": None}])
|
||||
@pytest.mark.parametrize("with_vector", [False, True])
|
||||
async def test_async_search_without_query(mem_db_async, search_kwargs, with_vector):
|
||||
data = [{"id": i} for i in range(10)]
|
||||
if with_vector:
|
||||
for row in data:
|
||||
row["vector"] = [float(row["id"]), 1.0]
|
||||
table = await mem_db_async.create_table("test", data)
|
||||
|
||||
query = await table.search(**search_kwargs)
|
||||
assert isinstance(query, AsyncQuery)
|
||||
assert await query.limit(3).to_arrow() == await table.query().limit(3).to_arrow()
|
||||
|
||||
query = await table.search(**search_kwargs)
|
||||
assert await query.where("id >= 8").select(["id"]).limit(3).to_list() == [
|
||||
{"id": 8},
|
||||
{"id": 9},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_none_query(table_async: AsyncTable):
|
||||
with pytest.raises(ValueError):
|
||||
@@ -1588,8 +1874,9 @@ def test_query_serialization_sync(table: lancedb.table.Table):
|
||||
q,
|
||||
vector_column="vector",
|
||||
vector=[5.0, 6.0],
|
||||
minimum_nprobes=10,
|
||||
maximum_nprobes=10,
|
||||
nprobes=10,
|
||||
minimum_nprobes=None,
|
||||
maximum_nprobes=None,
|
||||
refine_factor=5,
|
||||
)
|
||||
|
||||
@@ -1607,8 +1894,19 @@ def test_query_serialization_sync(table: lancedb.table.Table):
|
||||
q,
|
||||
vector_column="vector",
|
||||
vector=[5.0, 6.0],
|
||||
minimum_nprobes=50,
|
||||
maximum_nprobes=50,
|
||||
nprobes=50,
|
||||
minimum_nprobes=None,
|
||||
maximum_nprobes=None,
|
||||
)
|
||||
|
||||
q = table.search([5.0, 6.0]).minimum_nprobes(5).nprobes(50).to_query_object()
|
||||
check_set_props(
|
||||
q,
|
||||
vector_column="vector",
|
||||
vector=[5.0, 6.0],
|
||||
nprobes=50,
|
||||
minimum_nprobes=5,
|
||||
maximum_nprobes=None,
|
||||
)
|
||||
|
||||
q = table.search([5.0, 6.0]).maximum_nprobes(10).to_query_object()
|
||||
@@ -1620,6 +1918,15 @@ def test_query_serialization_sync(table: lancedb.table.Table):
|
||||
minimum_nprobes=None,
|
||||
)
|
||||
|
||||
q = table.search([5.0, 6.0]).nprobes(5).maximum_nprobes(0).to_query_object()
|
||||
check_set_props(
|
||||
q,
|
||||
vector_column="vector",
|
||||
vector=[5.0, 6.0],
|
||||
nprobes=5,
|
||||
maximum_nprobes=0,
|
||||
)
|
||||
|
||||
q = table.search([5.0, 6.0]).distance_range(0.0, 1.0).to_query_object()
|
||||
check_set_props(
|
||||
q, vector_column="vector", vector=[5.0, 6.0], lower_bound=0.0, upper_bound=1.0
|
||||
@@ -1669,8 +1976,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
limit=10,
|
||||
vector=sample_vector,
|
||||
postfilter=False,
|
||||
minimum_nprobes=20,
|
||||
maximum_nprobes=20,
|
||||
minimum_nprobes=None,
|
||||
maximum_nprobes=None,
|
||||
with_row_id=False,
|
||||
bypass_vector_index=False,
|
||||
)
|
||||
@@ -1680,8 +1987,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
q,
|
||||
vector=sample_vector,
|
||||
postfilter=False,
|
||||
minimum_nprobes=20,
|
||||
maximum_nprobes=20,
|
||||
minimum_nprobes=None,
|
||||
maximum_nprobes=None,
|
||||
with_row_id=False,
|
||||
bypass_vector_index=False,
|
||||
limit=10,
|
||||
@@ -1692,8 +1999,9 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
q,
|
||||
vector=sample_vector,
|
||||
postfilter=False,
|
||||
minimum_nprobes=50,
|
||||
maximum_nprobes=50,
|
||||
nprobes=50,
|
||||
minimum_nprobes=None,
|
||||
maximum_nprobes=None,
|
||||
with_row_id=False,
|
||||
bypass_vector_index=False,
|
||||
limit=10,
|
||||
@@ -1712,8 +2020,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
filter="id = 1",
|
||||
postfilter=True,
|
||||
vector=sample_vector,
|
||||
minimum_nprobes=20,
|
||||
maximum_nprobes=20,
|
||||
minimum_nprobes=None,
|
||||
maximum_nprobes=None,
|
||||
with_row_id=False,
|
||||
bypass_vector_index=False,
|
||||
)
|
||||
@@ -1727,8 +2035,9 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
check_set_props(
|
||||
q,
|
||||
vector=sample_vector,
|
||||
minimum_nprobes=10,
|
||||
maximum_nprobes=10,
|
||||
nprobes=10,
|
||||
minimum_nprobes=None,
|
||||
maximum_nprobes=None,
|
||||
refine_factor=5,
|
||||
postfilter=False,
|
||||
with_row_id=False,
|
||||
@@ -1741,7 +2050,24 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
q,
|
||||
vector=sample_vector,
|
||||
minimum_nprobes=5,
|
||||
maximum_nprobes=20,
|
||||
maximum_nprobes=None,
|
||||
postfilter=False,
|
||||
with_row_id=False,
|
||||
bypass_vector_index=False,
|
||||
limit=10,
|
||||
)
|
||||
|
||||
q = (
|
||||
(await table_async.search([5.0, 6.0]))
|
||||
.nprobes(5)
|
||||
.maximum_nprobes(0)
|
||||
.to_query_object()
|
||||
)
|
||||
check_set_props(
|
||||
q,
|
||||
vector=sample_vector,
|
||||
nprobes=5,
|
||||
maximum_nprobes=0,
|
||||
postfilter=False,
|
||||
with_row_id=False,
|
||||
bypass_vector_index=False,
|
||||
@@ -1759,8 +2085,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
lower_bound=0.0,
|
||||
upper_bound=1.0,
|
||||
postfilter=False,
|
||||
minimum_nprobes=20,
|
||||
maximum_nprobes=20,
|
||||
minimum_nprobes=None,
|
||||
maximum_nprobes=None,
|
||||
with_row_id=False,
|
||||
bypass_vector_index=False,
|
||||
limit=10,
|
||||
@@ -1772,8 +2098,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
distance_type="cosine",
|
||||
vector=sample_vector,
|
||||
postfilter=False,
|
||||
minimum_nprobes=20,
|
||||
maximum_nprobes=20,
|
||||
minimum_nprobes=None,
|
||||
maximum_nprobes=None,
|
||||
with_row_id=False,
|
||||
bypass_vector_index=False,
|
||||
limit=10,
|
||||
@@ -1785,8 +2111,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
ef=7,
|
||||
vector=sample_vector,
|
||||
postfilter=False,
|
||||
minimum_nprobes=20,
|
||||
maximum_nprobes=20,
|
||||
minimum_nprobes=None,
|
||||
maximum_nprobes=None,
|
||||
with_row_id=False,
|
||||
bypass_vector_index=False,
|
||||
limit=10,
|
||||
@@ -1798,8 +2124,8 @@ async def test_query_serialization_async(table_async: AsyncTable):
|
||||
bypass_vector_index=True,
|
||||
vector=sample_vector,
|
||||
postfilter=False,
|
||||
minimum_nprobes=20,
|
||||
maximum_nprobes=20,
|
||||
minimum_nprobes=None,
|
||||
maximum_nprobes=None,
|
||||
with_row_id=False,
|
||||
limit=10,
|
||||
)
|
||||
@@ -2091,6 +2417,12 @@ def test_ensure_vector_query_nested_empty_list():
|
||||
ensure_vector_query([[]])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("query", [[], np.array([], dtype=np.float32)])
|
||||
def test_vector_query_builder_empty_vector(table, query):
|
||||
with pytest.raises(ValueError, match="^Query vector must not be empty$"):
|
||||
LanceVectorQueryBuilder(table, query, "vector")
|
||||
|
||||
|
||||
def test_fast_search(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
|
||||
@@ -2152,3 +2484,127 @@ def test_blob_v2_to_batches_row_id(tmp_db):
|
||||
assert "_rowid" in hits.column_names
|
||||
blobs = table.fetch_blobs("blob", hits)
|
||||
assert [blobs[i].as_py() for i in range(len(blobs))] == [b"one", b"two"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("metric", ["l2", "cosine", "dot"])
|
||||
@pytest.mark.parametrize(
|
||||
"effort,probes", [(None, 4), (0, 1), (0.25, 2), (0.5, 4), (0.75, 8), (1, 16)]
|
||||
)
|
||||
def test_search_effort_probe_budget(tmp_path, monkeypatch, metric, effort, probes):
|
||||
import re
|
||||
from lancedb.index import IvfFlat
|
||||
|
||||
monkeypatch.setenv("LANCE_AUTO_PROBE_MARGIN", "0")
|
||||
monkeypatch.setenv("LANCE_AUTO_MIN_INITIAL_NPROBES", "4")
|
||||
monkeypatch.setenv("LANCE_AUTO_MAX_INITIAL_NPROBES", "4")
|
||||
rng = np.random.default_rng(42)
|
||||
vectors = np.repeat(np.eye(16, dtype=np.float32), 32, axis=0)
|
||||
vectors += rng.normal(0, 0.001, vectors.shape).astype(np.float32)
|
||||
data = pa.table(
|
||||
{
|
||||
"id": np.arange(512),
|
||||
"text": ["dog"] * 512,
|
||||
"vector": pa.FixedSizeListArray.from_arrays(pa.array(vectors.ravel()), 16),
|
||||
}
|
||||
)
|
||||
table = lancedb.connect(tmp_path).create_table("effort", data.slice(0, 256))
|
||||
table.add(data.slice(256))
|
||||
table.create_index(
|
||||
"vector", config=IvfFlat(distance_type=metric, num_partitions=16)
|
||||
)
|
||||
vector = np.eye(16, dtype=np.float32)[0]
|
||||
|
||||
def search(value=effort):
|
||||
return table.search(vector).distance_type(metric).search_effort(value).limit(8)
|
||||
|
||||
def partitions(query):
|
||||
vector_plan = query.analyze_plan().split("FTS Search Plan:", 1)[0]
|
||||
values = re.findall(r"partitions_searched=(\d+)", vector_plan)
|
||||
assert values
|
||||
return sum(map(int, values))
|
||||
|
||||
assert partitions(search()) == probes
|
||||
exact = set(
|
||||
table.search(vector)
|
||||
.distance_type(metric)
|
||||
.bypass_vector_index()
|
||||
.limit(8)
|
||||
.to_arrow()["id"]
|
||||
.to_pylist()
|
||||
)
|
||||
actual = set(search().to_arrow()["id"].to_pylist())
|
||||
assert len(actual & exact) / 8 >= 0.5
|
||||
expected = 1 if effort is None else probes
|
||||
assert partitions(search().nprobes(1)) == expected
|
||||
assert (
|
||||
partitions(
|
||||
table.search(vector)
|
||||
.distance_type(metric)
|
||||
.nprobes(1)
|
||||
.search_effort(effort)
|
||||
.limit(8)
|
||||
)
|
||||
== expected
|
||||
)
|
||||
assert partitions(search().nprobes(1).search_effort(None)) == 1
|
||||
assert partitions(search().maximum_nprobes(2)) <= 2
|
||||
assert partitions(search().minimum_nprobes(3).maximum_nprobes(3)) == 3
|
||||
assert partitions(search(0).nprobes(16).maximum_nprobes(0)) == 1
|
||||
if effort is None:
|
||||
assert search().to_arrow() == search(0.5).to_arrow()
|
||||
filtered = search(0).maximum_nprobes(2).where("id % 8 = 0", prefilter=True)
|
||||
assert partitions(filtered) <= 2
|
||||
assert all(row["id"] % 8 == 0 for row in filtered.to_list())
|
||||
|
||||
if metric == "l2" and effort == 0.5:
|
||||
table.create_index("text", config=FTS())
|
||||
hybrid = (
|
||||
table.search(query_type="hybrid")
|
||||
.vector(vector)
|
||||
.text("dog")
|
||||
.nprobes(1)
|
||||
.search_effort(0.5)
|
||||
.limit(8)
|
||||
)
|
||||
assert partitions(hybrid) == 4
|
||||
assert partitions(hybrid.search_effort(None)) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"effort", [-0.1, 1.1, float("nan"), float("inf"), -float("inf")]
|
||||
)
|
||||
def test_search_effort_invalid_sync(table, effort):
|
||||
with pytest.raises(ValueError, match="search_effort must be finite and in"):
|
||||
table.search([0, 0]).search_effort(effort).to_list()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("hybrid", [False, True])
|
||||
@pytest.mark.parametrize("effort", [None, 0, 0.25, 0.5, 0.75, 1])
|
||||
async def test_search_effort_async_roundtrip(table_async, hybrid, effort):
|
||||
query = table_async.vector_search([0, 0])
|
||||
if hybrid:
|
||||
query = query.nearest_to_text("dog")
|
||||
query = query.nprobes(5).search_effort(effort).minimum_nprobes(2).maximum_nprobes(3)
|
||||
request = query.to_query_object()
|
||||
assert request.search_effort == effort
|
||||
assert request.nprobes == 5
|
||||
assert request.minimum_nprobes == 2
|
||||
assert request.maximum_nprobes == 3
|
||||
assert query.search_effort(None).to_query_object().search_effort is None
|
||||
query.search_effort(effort)
|
||||
if not hybrid:
|
||||
assert len(await query.limit(1).to_list()) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("hybrid", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"effort", [-0.1, 1.1, float("nan"), float("inf"), -float("inf")]
|
||||
)
|
||||
async def test_search_effort_invalid_async(table_async, hybrid, effort):
|
||||
query = table_async.vector_search([0, 0])
|
||||
if hybrid:
|
||||
query = query.nearest_to_text("dog")
|
||||
with pytest.raises(ValueError, match="search_effort must be finite and in"):
|
||||
query.search_effort(effort)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,8 +1,10 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import threading
|
||||
from uuid import UUID
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import pyarrow as pa
|
||||
|
||||
@@ -11,7 +13,7 @@ from lancedb import _lancedb
|
||||
from lancedb.arrow import AsyncRecordBatchReader
|
||||
from lancedb.db import AsyncConnection
|
||||
from lancedb.remote.db import RemoteDBConnection
|
||||
from lancedb.sql import AsyncQuery, Query
|
||||
from lancedb.sql import AsyncQuery, Query, to_parameter_batch
|
||||
|
||||
NIL_QUERY_ID = UUID(int=0)
|
||||
|
||||
@@ -24,12 +26,20 @@ class FakeNativeQuery:
|
||||
|
||||
|
||||
class FakeNativeConnection:
|
||||
async def execute_query_async(self, query, *, default_namespace_path=None):
|
||||
def __init__(self):
|
||||
self.parameters = []
|
||||
|
||||
async def execute_query_async(
|
||||
self, query, *, default_namespace_path=None, parameters=None
|
||||
):
|
||||
self.parameters.append(parameters)
|
||||
return FakeNativeQuery()
|
||||
|
||||
|
||||
class FakeAsyncConnection:
|
||||
async def execute_query_async(self, query, *, default_namespace_path=None):
|
||||
async def execute_query_async(
|
||||
self, query, *, default_namespace_path=None, parameters=None
|
||||
):
|
||||
return AsyncQuery(FakeNativeQuery())
|
||||
|
||||
|
||||
@@ -160,3 +170,100 @@ def test_execute_query_async_rejects_invalid_namespace_components(
|
||||
remote_connection().execute_query_async(
|
||||
"SELECT 1", default_namespace_path=default_namespace_path
|
||||
)
|
||||
|
||||
|
||||
def test_positional_parameters_keep_their_types():
|
||||
vector = np.array([0.1, 0.2, 0.3], dtype=np.float32)
|
||||
batch = to_parameter_batch(
|
||||
[42, 1.5, "a", np.float32(0.25), vector, pa.scalar(7, pa.int16())]
|
||||
)
|
||||
assert batch.num_rows == 1
|
||||
assert batch.schema.names == ["1", "2", "3", "4", "5", "6"]
|
||||
assert batch.schema.types == [
|
||||
pa.int64(),
|
||||
pa.float64(),
|
||||
pa.string(),
|
||||
pa.float32(),
|
||||
pa.list_(pa.float32(), 3),
|
||||
pa.int16(),
|
||||
]
|
||||
sent = batch.column(4).flatten().to_numpy()
|
||||
assert sent.dtype == np.float32
|
||||
assert sent.tobytes() == vector.tobytes()
|
||||
|
||||
|
||||
def test_named_parameters_and_batches():
|
||||
half = pa.array(np.array([1.0, 2.0], dtype=np.float16))
|
||||
batch = to_parameter_batch({"vector": half, "k": 3})
|
||||
assert batch.schema.names == ["vector", "k"]
|
||||
assert batch.schema.field("vector").type == pa.list_(pa.float16(), 2)
|
||||
|
||||
row = pa.record_batch({"id": [1]})
|
||||
assert to_parameter_batch(row) is row
|
||||
assert to_parameter_batch(pa.table({"id": [1]})).equals(row)
|
||||
assert to_parameter_batch(None) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"parameters",
|
||||
[pa.record_batch({"id": [1, 2]}), pa.table({"id": pa.array([], pa.int64())})],
|
||||
)
|
||||
def test_parameters_are_exactly_one_row(parameters):
|
||||
with pytest.raises(ValueError, match="exactly one row"):
|
||||
to_parameter_batch(parameters)
|
||||
|
||||
|
||||
def test_parameters_must_be_a_collection():
|
||||
with pytest.raises(TypeError, match="query parameters"):
|
||||
to_parameter_batch("SELECT 1")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_execute_query_sends_parameters_as_a_batch():
|
||||
native = FakeNativeConnection()
|
||||
connection = AsyncConnection(native)
|
||||
await connection.execute_query("SELECT $1", parameters=[1])
|
||||
await connection.execute_query_async("SELECT 1")
|
||||
sent, unparameterized = native.parameters
|
||||
assert sent.equals(pa.record_batch({"1": [1]}))
|
||||
assert unparameterized is None
|
||||
|
||||
|
||||
def test_parameters_reach_the_server_as_the_values_that_were_sent():
|
||||
"""A round trip through the bindings against a Flight server that answers
|
||||
each exchange with the parameter row it received."""
|
||||
flight = pytest.importorskip("pyarrow.flight")
|
||||
|
||||
class EchoServer(flight.FlightServerBase):
|
||||
def __init__(self):
|
||||
super().__init__("grpc://127.0.0.1:0")
|
||||
self.commands = []
|
||||
self.lock = threading.Lock()
|
||||
|
||||
def do_exchange(self, context, descriptor, reader, writer):
|
||||
parameters = reader.read_all()
|
||||
with self.lock:
|
||||
self.commands.append(descriptor.command)
|
||||
writer.begin(parameters.schema)
|
||||
writer.write_table(parameters)
|
||||
writer.close()
|
||||
|
||||
vector = np.array([0.1, 1e-45, -0.0], dtype=np.float32)
|
||||
with EchoServer() as server:
|
||||
db = lancedb.connect(
|
||||
"db://analytics",
|
||||
api_key="test-key",
|
||||
host_override="http://localhost:10024",
|
||||
sql_host_override=f"grpc://127.0.0.1:{server.port}",
|
||||
)
|
||||
echoed = db.execute_query(
|
||||
"SELECT echo", parameters={"vector": vector, "k": 10}
|
||||
).read_all()
|
||||
(command,) = server.commands
|
||||
|
||||
assert b"CommandStatementQuery" in command
|
||||
assert b"SELECT echo" in command
|
||||
assert echoed.schema.field("vector").type == pa.list_(pa.float32(), 3)
|
||||
received = echoed.column("vector").combine_chunks().flatten().to_numpy()
|
||||
assert received.view(np.uint32).tolist() == vector.view(np.uint32).tolist()
|
||||
assert echoed.column("k").to_pylist() == [10]
|
||||
|
||||
@@ -12,7 +12,7 @@ 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 Any, List
|
||||
from unittest.mock import patch
|
||||
@@ -113,6 +113,32 @@ def test_basic(mem_db: DBConnection):
|
||||
assert table.to_arrow() == expected_data
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_v2", [False, True])
|
||||
def test_migrate_v2_manifest_paths(tmp_path, enable_v2):
|
||||
db = lancedb.connect(
|
||||
tmp_path,
|
||||
storage_options={"new_table_enable_v2_manifest_paths": str(enable_v2).lower()},
|
||||
)
|
||||
table = db.create_table("calls", [{"id": 1, "vector": [1.0, 1.0]}])
|
||||
table.add([{"id": 2, "vector": [2.0, 2.0]}])
|
||||
expected_data = table.to_arrow()
|
||||
expected_versions = table.list_versions()
|
||||
assert table.uses_v2_manifest_paths() == enable_v2
|
||||
|
||||
# Migration is also safe to repeat on a table already using v2 paths.
|
||||
for _ in range(2):
|
||||
table.migrate_v2_manifest_paths()
|
||||
assert table.uses_v2_manifest_paths()
|
||||
reopened = db.open_table("calls")
|
||||
assert reopened.uses_v2_manifest_paths()
|
||||
assert reopened.to_arrow() == expected_data
|
||||
assert reopened.list_versions() == expected_versions
|
||||
|
||||
manifests = list((tmp_path / "calls.lance" / "_versions").glob("*.manifest"))
|
||||
assert len(manifests) == len(expected_versions)
|
||||
assert all(len(path.stem) == 20 and path.stem.isdigit() for path in manifests)
|
||||
|
||||
|
||||
def test_search_preserves_nulls_from_sliced_arrow_table(mem_db: DBConnection):
|
||||
data = pa.table(
|
||||
{
|
||||
@@ -967,6 +993,63 @@ def test_polars(mem_db: DBConnection):
|
||||
assert len(filtered_result) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_versions_timestamp_precision():
|
||||
# 2026-10-03T01:03:49.274Z in nanoseconds. Float math turns .274000 into
|
||||
# .273999, so make sure the conversion is exact.
|
||||
ts_nanos = 1790989429274000000
|
||||
|
||||
class FakeInner:
|
||||
async def list_versions(self):
|
||||
return [{"version": 1, "timestamp": ts_nanos, "metadata": {}}]
|
||||
|
||||
table = table_module.AsyncTable(FakeInner())
|
||||
versions = await table.list_versions()
|
||||
|
||||
expected = datetime.fromtimestamp(ts_nanos // 1_000_000_000)
|
||||
assert versions[0]["timestamp"] == expected + timedelta(microseconds=274000)
|
||||
assert versions[0]["timestamp"].microsecond == 274000
|
||||
|
||||
|
||||
def test_list_versions_timestamp_precision_sync(mem_db: DBConnection):
|
||||
# The sync LanceTable.list_versions delegates to AsyncTable.list_versions,
|
||||
# which is the path the parity report hit.
|
||||
ts_nanos = 1790989429274000000
|
||||
|
||||
class FakeInner:
|
||||
async def list_versions(self):
|
||||
return [{"version": 1, "timestamp": ts_nanos, "metadata": {}}]
|
||||
|
||||
table = mem_db.create_table("ts_precision", data=[{"id": 1}])
|
||||
table._table = table_module.AsyncTable(FakeInner())
|
||||
versions = table.list_versions()
|
||||
|
||||
expected = datetime.fromtimestamp(ts_nanos // 1_000_000_000)
|
||||
assert versions[0]["timestamp"] == expected + timedelta(microseconds=274000)
|
||||
assert versions[0]["timestamp"].microsecond == 274000
|
||||
|
||||
|
||||
def test_list_versions_timestamp_precision_remote():
|
||||
# RemoteTable.list_versions also delegates to AsyncTable.list_versions.
|
||||
from lancedb.remote.table import RemoteTable
|
||||
|
||||
ts_nanos = 1790989429274000000
|
||||
|
||||
class FakeInner:
|
||||
def name(self):
|
||||
return "ts_precision"
|
||||
|
||||
async def list_versions(self):
|
||||
return [{"version": 1, "timestamp": ts_nanos, "metadata": {}}]
|
||||
|
||||
table = RemoteTable(table_module.AsyncTable(FakeInner()), "dev")
|
||||
versions = table.list_versions()
|
||||
|
||||
expected = datetime.fromtimestamp(ts_nanos // 1_000_000_000)
|
||||
assert versions[0]["timestamp"] == expected + timedelta(microseconds=274000)
|
||||
assert versions[0]["timestamp"].microsecond == 274000
|
||||
|
||||
|
||||
def test_versioning(mem_db: DBConnection):
|
||||
table = mem_db.create_table(
|
||||
"test",
|
||||
@@ -2529,6 +2612,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",
|
||||
@@ -3751,6 +3879,24 @@ def test_empty_query(mem_db: DBConnection):
|
||||
assert df.num_rows == 42
|
||||
|
||||
|
||||
@pytest.mark.parametrize("query", [[], np.array([], dtype=np.float32)])
|
||||
@pytest.mark.parametrize("vector_column_name", [None, "vector"])
|
||||
@pytest.mark.parametrize("query_type", ["auto", "vector"])
|
||||
@pytest.mark.parametrize("multiple_vector_columns", [False, True])
|
||||
def test_search_empty_vector(
|
||||
mem_db, query, vector_column_name, query_type, multiple_vector_columns
|
||||
):
|
||||
fields = [pa.field("vector", pa.list_(pa.float32(), 8))]
|
||||
if multiple_vector_columns:
|
||||
fields.append(pa.field("vec2", pa.list_(pa.float32(), 4)))
|
||||
table = mem_db.create_table("empty_vector_query", schema=pa.schema(fields))
|
||||
|
||||
with pytest.raises(ValueError, match="^Query vector must not be empty$"):
|
||||
table.search(
|
||||
query, vector_column_name=vector_column_name, query_type=query_type
|
||||
).limit(3).to_arrow()
|
||||
|
||||
|
||||
def test_search_with_schema_inf_single_vector(mem_db: DBConnection):
|
||||
class MyTable(LanceModel):
|
||||
text: str
|
||||
@@ -3870,6 +4016,26 @@ def test_search_nested_vector_no_candidates(mem_db: DBConnection):
|
||||
table.search([0.0, 1.0]).limit(1).to_arrow()
|
||||
|
||||
|
||||
def test_search_single_vector_column_wrong_query_dim(mem_db: DBConnection):
|
||||
"""A query vector whose dimension doesn't match the (only) vector column
|
||||
must still resolve that column, surfacing a dimension-mismatch error from
|
||||
query execution rather than a misleading "no vector column" error."""
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.uint32(), False),
|
||||
pa.field("vector_col", pa.list_(pa.float32(), 32), False),
|
||||
]
|
||||
)
|
||||
table = mem_db.create_table(
|
||||
"wrong_query_dim",
|
||||
data=[{"id": 1, "vector_col": [0.1] * 32}],
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="dim"):
|
||||
table.search([0.01] * 120).limit(1).to_arrow()
|
||||
|
||||
|
||||
def test_compact_cleanup(tmp_db: DBConnection):
|
||||
pytest.importorskip("lance")
|
||||
table = tmp_db.create_table(
|
||||
@@ -4144,14 +4310,74 @@ async def test_add_columns_with_schema(mem_db_async: AsyncConnection):
|
||||
)
|
||||
|
||||
|
||||
def test_alter_columns(mem_db: DBConnection):
|
||||
@pytest.mark.parametrize("rename_key", ["rename", "name"])
|
||||
def test_alter_columns(mem_db: DBConnection, rename_key):
|
||||
data = pa.table({"id": [0, 1]})
|
||||
table = mem_db.create_table("my_table", data=data)
|
||||
alter_columns_res = table.alter_columns({"path": "id", "rename": "new_id"})
|
||||
alter_columns_res = table.alter_columns({"path": "id", rename_key: "new_id"})
|
||||
assert alter_columns_res.version == 2
|
||||
assert table.to_arrow().column_names == ["new_id"]
|
||||
|
||||
|
||||
INVALID_COLUMN_ALTERATIONS = [
|
||||
(({"path": "id"},), "One of rename, nullable or data_type"),
|
||||
(
|
||||
({"path": "id", "nulable": False},), # spellchecker:disable-line
|
||||
"Unknown column alteration key 'nulable'", # spellchecker:disable-line
|
||||
),
|
||||
(
|
||||
(
|
||||
{
|
||||
"path": "id",
|
||||
"rename": "new_id",
|
||||
"nulable": False, # spellchecker:disable-line
|
||||
},
|
||||
),
|
||||
"Unknown column alteration key 'nulable'", # spellchecker:disable-line
|
||||
),
|
||||
(
|
||||
({"path": "id", "rename": "new_id"}, {"path": "id"}),
|
||||
"One of rename, nullable or data_type",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_alter_columns_nullable_false(mem_db: DBConnection):
|
||||
table = mem_db.create_table("my_table", data=pa.table({"id": [0, 1]}))
|
||||
result = table.alter_columns({"path": "id", "nullable": False})
|
||||
assert result.version == 2
|
||||
assert not table.schema.field("id").nullable
|
||||
|
||||
|
||||
@pytest.mark.parametrize("alterations, match", INVALID_COLUMN_ALTERATIONS)
|
||||
def test_alter_columns_rejects_invalid(mem_db: DBConnection, alterations, match):
|
||||
table = mem_db.create_table("my_table", data=pa.table({"id": [0, 1]}))
|
||||
initial_version = table.version
|
||||
initial_schema = table.schema
|
||||
|
||||
with pytest.raises(ValueError, match=match):
|
||||
table.alter_columns(*alterations)
|
||||
|
||||
assert table.version == initial_version
|
||||
assert table.schema == initial_schema
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("alterations, match", INVALID_COLUMN_ALTERATIONS)
|
||||
async def test_alter_columns_rejects_invalid_async(
|
||||
mem_db_async: AsyncConnection, alterations, match
|
||||
):
|
||||
table = await mem_db_async.create_table("my_table", data=pa.table({"id": [0, 1]}))
|
||||
initial_version = await table.version()
|
||||
initial_schema = await table.schema()
|
||||
|
||||
with pytest.raises(ValueError, match=match):
|
||||
await table.alter_columns(*alterations)
|
||||
|
||||
assert await table.version() == initial_version
|
||||
assert await table.schema() == initial_schema
|
||||
|
||||
|
||||
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)
|
||||
@@ -4177,10 +4403,11 @@ def test_update_field_metadata(mem_db: DBConnection):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_alter_columns_async(mem_db_async: AsyncConnection):
|
||||
@pytest.mark.parametrize("rename_key", ["rename", "name"])
|
||||
async def test_alter_columns_async(mem_db_async: AsyncConnection, rename_key):
|
||||
data = pa.table({"id": [0, 1]})
|
||||
table = await mem_db_async.create_table("my_table", data=data)
|
||||
alter_columns_res = await table.alter_columns({"path": "id", "rename": "new_id"})
|
||||
alter_columns_res = await table.alter_columns({"path": "id", rename_key: "new_id"})
|
||||
assert alter_columns_res.version == 2
|
||||
assert (await table.to_arrow()).column_names == ["new_id"]
|
||||
alter_columns_res = await table.alter_columns(
|
||||
@@ -4507,22 +4734,25 @@ def test_stats(mem_db: DBConnection):
|
||||
"my_table",
|
||||
data=[{"text": "foo", "id": 0}, {"text": "bar", "id": 1}],
|
||||
)
|
||||
assert len(table) == 2
|
||||
table.add([{"text": "baz", "id": 1}])
|
||||
assert len(table) == 3
|
||||
stats = table.stats()
|
||||
print(f"{stats=}")
|
||||
assert stats == {
|
||||
# Full on-disk size of the data file, footer and metadata included.
|
||||
"total_bytes": 637,
|
||||
"num_rows": 2,
|
||||
# Full on-disk size of the data files, footer and metadata included:
|
||||
# 637 bytes for the two-row fragment plus 415 for the one-row fragment.
|
||||
"total_bytes": 1052,
|
||||
"num_rows": 3,
|
||||
"num_deleted_rows": 0,
|
||||
"num_indices": 0,
|
||||
"fragment_stats": {
|
||||
"num_fragments": 1,
|
||||
"num_small_fragments": 1,
|
||||
"num_fragments": 2,
|
||||
"num_small_fragments": 2,
|
||||
"lengths": {
|
||||
"min": 2,
|
||||
"min": 1,
|
||||
"max": 2,
|
||||
"mean": 2,
|
||||
"p25": 2,
|
||||
"mean": 1,
|
||||
"p25": 1,
|
||||
"p50": 2,
|
||||
"p75": 2,
|
||||
"p99": 2,
|
||||
@@ -4530,12 +4760,19 @@ def test_stats(mem_db: DBConnection):
|
||||
},
|
||||
}
|
||||
|
||||
# Both rows with id = 1 are deleted, but that empties the second fragment,
|
||||
# which is dropped outright rather than kept with a deletion file, so only
|
||||
# the row marked in the surviving fragment is counted.
|
||||
table.delete("id = 1")
|
||||
# Index files count toward total_bytes too (only deletion files and
|
||||
# manifests are excluded).
|
||||
table.create_index("id", config=BTree())
|
||||
stats_with_index = table.stats()
|
||||
assert stats_with_index["num_indices"] == 1
|
||||
assert stats_with_index["total_bytes"] > stats["total_bytes"]
|
||||
assert stats_with_index["num_rows"] == 1
|
||||
assert stats_with_index["num_deleted_rows"] == 1
|
||||
assert stats_with_index["fragment_stats"]["num_fragments"] == 1
|
||||
|
||||
|
||||
def test_create_table_empty_list_with_schema(mem_db: DBConnection):
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
|
||||
import os
|
||||
import pathlib
|
||||
from decimal import Decimal, localcontext
|
||||
from typing import Optional
|
||||
|
||||
import lance
|
||||
@@ -27,7 +28,13 @@ import pandas as pd
|
||||
import polars as pl
|
||||
import pytest
|
||||
import lancedb
|
||||
from lancedb.util import flatten_columns, get_uri_scheme, join_uri, value_to_sql
|
||||
from lancedb.util import (
|
||||
flatten_columns,
|
||||
get_uri_scheme,
|
||||
inf_vector_column_query,
|
||||
join_uri,
|
||||
value_to_sql,
|
||||
)
|
||||
from utils import exception_output
|
||||
|
||||
|
||||
@@ -157,6 +164,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')"
|
||||
@@ -951,6 +993,39 @@ def test_cast_to_target_schema_coerces_binary_to_blob_v2():
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"string_type", [pa.string(), pa.large_string(), pa.string_view()]
|
||||
)
|
||||
def test_cast_to_target_schema_coerces_string_to_blob_v2_uri(string_type):
|
||||
data = pa.table({"image": pa.array(["file:///payload", None], type=string_type)})
|
||||
target = pa.schema([lancedb.blob("image")])
|
||||
|
||||
output = _cast_to_target_schema(data.to_reader(), target).read_all()
|
||||
|
||||
image = output["image"].chunk(0)
|
||||
assert type(image.type) is lancedb.BlobType
|
||||
assert image.storage.to_pylist() == [
|
||||
{"data": None, "uri": "file:///payload", "position": None, "size": None},
|
||||
None,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"source_type, value",
|
||||
[
|
||||
(pa.binary(4), b"abcd"),
|
||||
(pa.list_(pa.uint8()), [1, 2]),
|
||||
(pa.int64(), 42),
|
||||
],
|
||||
)
|
||||
def test_cast_to_target_schema_rejects_unsupported_blob_input(source_type, value):
|
||||
data = pa.table({"image": pa.array([value], type=source_type)})
|
||||
target = pa.schema([lancedb.blob("image")])
|
||||
|
||||
with pytest.raises(ValueError, match="cannot coerce column 'image' with type"):
|
||||
_cast_to_target_schema(data.to_reader(), target)
|
||||
|
||||
|
||||
def test_cast_to_target_schema_coerces_binary_to_metadata_blob_struct():
|
||||
storage = lancedb.blob("image").type.storage_type
|
||||
target = pa.schema(
|
||||
@@ -1182,3 +1257,41 @@ def test_infer_vector_column_raises_clear_error(tmp_path):
|
||||
# Plain vector search on a table with no vector column should raise
|
||||
# a clear ValueError, not a cryptic TypeError.
|
||||
table.search([1.0, 2.0]).to_list()
|
||||
|
||||
|
||||
def test_inf_vector_column_query_dim_mismatch_still_resolves():
|
||||
"""Regression: a query dimension that matches no column must not hide an
|
||||
otherwise-unambiguous vector column behind "no vector column in the data".
|
||||
|
||||
`dim` exists to disambiguate between *multiple* vector columns of
|
||||
different sizes. When the schema has exactly one vector column, that
|
||||
column is still the only sensible one to search; a dimension mismatch is
|
||||
a query-execution concern, not a missing-column one. This previously
|
||||
broke remote search (#4409 wired up client-side inference there), which
|
||||
expects the server to reject the mismatched dimension itself.
|
||||
"""
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.uint32(), False),
|
||||
pa.field("vector_col", pa.list_(pa.float32(), 32), False),
|
||||
pa.field("s", pa.string(), False),
|
||||
]
|
||||
)
|
||||
|
||||
assert inf_vector_column_query(schema, dim=32) == "vector_col"
|
||||
assert inf_vector_column_query(schema, dim=120) == "vector_col"
|
||||
assert inf_vector_column_query(schema) == "vector_col"
|
||||
|
||||
|
||||
def test_inf_vector_column_query_dim_mismatch_with_multiple_columns():
|
||||
"""With more than one vector column, a dimension that matches none of
|
||||
them is still ambiguous and must raise, not silently pick one."""
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("vector1", pa.list_(pa.float32(), 10), False),
|
||||
pa.field("vector2", pa.list_(pa.float32(), 20), False),
|
||||
]
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="more than one vector column"):
|
||||
inf_vector_column_query(schema, dim=99)
|
||||
|
||||
@@ -81,15 +81,15 @@ def test_base_table_hybrid_still_joins_on_row_ids():
|
||||
assert builder._vector_query._columns is None
|
||||
|
||||
|
||||
def test_wal_hybrid_with_use_lsm_false_joins_on_row_ids():
|
||||
builder = _wal_hybrid().use_lsm(False)
|
||||
def test_wal_hybrid_with_use_memwal_false_joins_on_row_ids():
|
||||
builder = _wal_hybrid().use_memwal(False)
|
||||
builder._create_query_builders()
|
||||
assert builder._pk_fusion is None
|
||||
|
||||
|
||||
def test_wal_hybrid_refuses_an_explicit_with_row_id():
|
||||
builder = _wal_hybrid().with_row_id(True)
|
||||
with pytest.raises(NotImplementedError, match="use_lsm"):
|
||||
with pytest.raises(NotImplementedError, match="use_memwal"):
|
||||
builder._create_query_builders()
|
||||
|
||||
|
||||
@@ -119,7 +119,9 @@ def test_wal_hybrid_fuses_on_the_key_and_drops_it_again():
|
||||
limit=10,
|
||||
with_row_ids=True,
|
||||
)
|
||||
results = builder._finish_hybrid_results(combined)
|
||||
results = builder._finish_hybrid_results(
|
||||
combined, blob_mode="lazy", blob_sources={}
|
||||
)
|
||||
|
||||
assert "_rowid" not in results.column_names, "the surrogate is internal"
|
||||
assert "id" not in results.column_names, "the caller selected only `text`"
|
||||
@@ -139,7 +141,7 @@ def test_wal_hybrid_refuses_row_id_even_with_return_score_all(row_id_first: bool
|
||||
else:
|
||||
builder = builder.rerank(reranker).with_row_id(True)
|
||||
|
||||
with pytest.raises(NotImplementedError, match="use_lsm"):
|
||||
with pytest.raises(NotImplementedError, match="use_memwal"):
|
||||
builder._create_query_builders()
|
||||
|
||||
|
||||
@@ -197,7 +199,7 @@ def test_fallback_explains_the_refusal_to_a_caller_who_asked_for_row_ids():
|
||||
def run():
|
||||
raise _REFUSAL
|
||||
|
||||
with pytest.raises(NotImplementedError, match="use_lsm") as info:
|
||||
with pytest.raises(NotImplementedError, match="use_memwal") as info:
|
||||
with_pk_fallback(
|
||||
run, mock.Mock(), fused=lambda: False, caller_requested_row_id=True
|
||||
)
|
||||
|
||||
@@ -73,10 +73,10 @@ class MockPermutationServer:
|
||||
},
|
||||
},
|
||||
)
|
||||
if path == f"/v1/table/{self.name}/get_lsm_write_spec/":
|
||||
if path == f"/v1/table/{self.name}/get_memwal_spec/":
|
||||
self._read_body(request)
|
||||
# Null spec: this table has no LSM write path.
|
||||
return self._json(request, {"lsm_write_spec": None})
|
||||
# Null spec: this table has no MemWAL write path.
|
||||
return self._json(request, {"memwal_spec": None})
|
||||
if path == f"/v1/table/{self.name}/count_rows/":
|
||||
self._read_body(request)
|
||||
return self._json(request, self.num_rows)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
+52
-15
@@ -1,12 +1,13 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use lancedb::catalog::{
|
||||
CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest, ListDatabasesRequest,
|
||||
};
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use futures::{StreamExt, future::poll_fn};
|
||||
|
||||
use lancedb::catalog::{CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest};
|
||||
use pyo3::exceptions::{PyStopAsyncIteration, PyValueError};
|
||||
use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pyfunction, pymethods};
|
||||
|
||||
use crate::connection::{Connection, PyClientConfig};
|
||||
@@ -20,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()
|
||||
@@ -67,19 +75,48 @@ impl Catalog {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (*, limit=None, page_token=None))]
|
||||
fn list_databases<'py>(
|
||||
self_: PyRef<'py, Self>,
|
||||
limit: Option<u32>,
|
||||
page_token: Option<String>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
#[pyo3(signature = (*, page_token=None, page_limit=None))]
|
||||
fn list_databases(&self, page_token: Option<String>, page_limit: Option<u32>) -> DatabaseNames {
|
||||
DatabaseNames {
|
||||
inner: Arc::new(Mutex::new(
|
||||
self.inner.list_databases(page_token, page_limit),
|
||||
)),
|
||||
next_lock: Arc::new(tokio::sync::Mutex::new(())),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass]
|
||||
pub struct DatabaseNames {
|
||||
inner: Arc<Mutex<lancedb::catalog::DatabaseNames>>,
|
||||
next_lock: Arc<tokio::sync::Mutex<()>>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl DatabaseNames {
|
||||
fn num_page_results(&self, py: Python<'_>) -> usize {
|
||||
py.detach(|| self.inner.lock().unwrap().num_page_results())
|
||||
}
|
||||
|
||||
fn page_token(&self, py: Python<'_>) -> Option<String> {
|
||||
py.detach(|| self.inner.lock().unwrap().page_token().map(str::to_owned))
|
||||
}
|
||||
|
||||
fn __aiter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> {
|
||||
self_
|
||||
}
|
||||
|
||||
fn __anext__(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
let next_lock = self_.next_lock.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let mut request = ListDatabasesRequest::default();
|
||||
request.limit = limit;
|
||||
request.page_token = page_token;
|
||||
let response = inner.list_databases(request).await.infer_error()?;
|
||||
Ok((response.databases, response.page_token))
|
||||
// Serialize advances while allowing synchronous state inspection
|
||||
// between polls, including while a REST request is pending.
|
||||
let _guard = next_lock.lock().await;
|
||||
poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx))
|
||||
.await
|
||||
.ok_or_else(|| PyStopAsyncIteration::new_err(""))?
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
+86
-67
@@ -14,9 +14,10 @@ use crate::{
|
||||
table::Table,
|
||||
};
|
||||
use arrow::{
|
||||
array::RecordBatch,
|
||||
datatypes::Schema,
|
||||
ffi_stream::ArrowArrayStreamReader,
|
||||
pyarrow::{FromPyArrow, ToPyArrow},
|
||||
pyarrow::{FromPyArrow, PyArrowType, ToPyArrow},
|
||||
};
|
||||
use lancedb::{
|
||||
connection::Connection as LanceConnection,
|
||||
@@ -148,18 +149,22 @@ impl Connection {
|
||||
self.get_inner().map(|inner| inner.uri().to_string())
|
||||
}
|
||||
|
||||
#[pyo3(signature = (query, *, default_namespace_path=None))]
|
||||
#[pyo3(signature = (query, *, default_namespace_path=None, parameters=None))]
|
||||
pub fn execute_query_async<'a>(
|
||||
self_: PyRef<'a, Self>,
|
||||
query: String,
|
||||
default_namespace_path: Option<Bound<'_, PyAny>>,
|
||||
parameters: Option<PyArrowType<RecordBatch>>,
|
||||
) -> PyResult<Bound<'a, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
let default_namespace_path = parse_default_namespace_path(default_namespace_path)?;
|
||||
future_into_py(self_.py(), async move {
|
||||
let operation = inner
|
||||
let mut operation = inner
|
||||
.execute_query_async(query)
|
||||
.default_namespace_path(default_namespace_path);
|
||||
if let Some(PyArrowType(parameters)) = parameters {
|
||||
operation = operation.parameters(parameters);
|
||||
}
|
||||
operation
|
||||
.execute()
|
||||
.await
|
||||
@@ -407,8 +412,7 @@ impl Connection {
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false, *, vector_source_json=None))]
|
||||
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))]
|
||||
pub fn create_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
@@ -417,17 +421,10 @@ impl Connection {
|
||||
filter: Option<String>,
|
||||
limit: Option<u64>,
|
||||
with_no_data: bool,
|
||||
vector_source_json: Option<String>,
|
||||
) -> 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(source) = vector_source_json {
|
||||
builder = builder.vector_source(
|
||||
serde_json::from_str(&source)
|
||||
.map_err(|e| PyValueError::new_err(e.to_string()))?,
|
||||
);
|
||||
}
|
||||
if let Some(projections) = projections {
|
||||
builder = builder.select(projections);
|
||||
}
|
||||
@@ -443,8 +440,7 @@ impl Connection {
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false, *, vector_source_json=None))]
|
||||
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))]
|
||||
pub fn create_materialized_view_async(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
@@ -453,17 +449,10 @@ impl Connection {
|
||||
filter: Option<String>,
|
||||
limit: Option<u64>,
|
||||
with_no_data: bool,
|
||||
vector_source_json: Option<String>,
|
||||
) -> 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(source) = vector_source_json {
|
||||
builder = builder.vector_source(
|
||||
serde_json::from_str(&source)
|
||||
.map_err(|e| PyValueError::new_err(e.to_string()))?,
|
||||
);
|
||||
}
|
||||
if let Some(projections) = projections {
|
||||
builder = builder.select(projections);
|
||||
}
|
||||
@@ -480,12 +469,16 @@ impl Connection {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_materialized_views(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let views = inner.list_materialized_views().await.infer_error()?;
|
||||
Ok(views)
|
||||
})
|
||||
#[pyo3(signature = (*, page_token=None, page_limit=None))]
|
||||
pub fn list_materialized_views(
|
||||
&self,
|
||||
page_token: Option<String>,
|
||||
page_limit: Option<u32>,
|
||||
) -> PyResult<crate::listing::NameListing> {
|
||||
Ok(crate::listing::NameListing::new(
|
||||
self.get_inner()?
|
||||
.list_materialized_views(crate::listing::options(page_token, page_limit)),
|
||||
))
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, namespace_path=None))]
|
||||
@@ -754,31 +747,37 @@ impl Connection {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (request_json, namespace_path=None))]
|
||||
pub fn create_function_async(
|
||||
self_: PyRef<'_, Self>,
|
||||
request_json: String,
|
||||
namespace_path: Option<Vec<String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
let request = lancedb::function::FunctionRegistrationRequest::from_json(&request_json)
|
||||
.infer_error()?;
|
||||
let namespace_path = namespace_path.unwrap_or_default();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.create_function_async(request)
|
||||
.create_function_async(request, &namespace_path)
|
||||
.await
|
||||
.infer_error()
|
||||
.map(crate::job::Job::new_typed)
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, version, namespace_path=None))]
|
||||
pub fn get_function(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
version: String,
|
||||
namespace_path: Option<Vec<String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
let namespace_path = namespace_path.unwrap_or_default();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.get_function(name, version)
|
||||
.get_function(name, version, &namespace_path)
|
||||
.await
|
||||
.infer_error()?
|
||||
.to_canonical_json()
|
||||
@@ -786,39 +785,50 @@ impl Connection {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_functions(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.list_functions()
|
||||
.await
|
||||
.infer_error()?
|
||||
.into_iter()
|
||||
.map(|function| function.to_canonical_json().infer_error())
|
||||
.collect::<PyResult<Vec<_>>>()
|
||||
})
|
||||
#[pyo3(signature = (namespace_path=None, *, page_token=None, page_limit=None))]
|
||||
pub fn list_functions(
|
||||
&self,
|
||||
namespace_path: Option<Vec<String>>,
|
||||
page_token: Option<String>,
|
||||
page_limit: Option<u32>,
|
||||
) -> PyResult<crate::listing::FunctionListing> {
|
||||
Ok(crate::listing::FunctionListing::new(
|
||||
self.get_inner()?.list_functions(
|
||||
&namespace_path.unwrap_or_default(),
|
||||
crate::listing::options(page_token, page_limit),
|
||||
),
|
||||
))
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, version, namespace_path=None))]
|
||||
pub fn drop_function(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
version: String,
|
||||
namespace_path: Option<Vec<String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
let namespace_path = namespace_path.unwrap_or_default();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.drop_function(name, version).await.infer_error()
|
||||
inner
|
||||
.drop_function(name, version, &namespace_path)
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, version, namespace_path=None))]
|
||||
pub fn drop_function_async(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
version: String,
|
||||
namespace_path: Option<Vec<String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
let namespace_path = namespace_path.unwrap_or_default();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.drop_function_async(name, version)
|
||||
.drop_function_async(name, version, &namespace_path)
|
||||
.await
|
||||
.infer_error()
|
||||
.map(|(dropped, job)| (dropped, crate::job::Job::new(job)))
|
||||
@@ -859,16 +869,19 @@ impl Connection {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (namespace_path=None))]
|
||||
#[pyo3(signature = (namespace_path=None, *, page_token=None, page_limit=None))]
|
||||
pub fn list_secrets(
|
||||
self_: PyRef<'_, Self>,
|
||||
&self,
|
||||
namespace_path: Option<Vec<String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
let namespace_path = namespace_path.unwrap_or_default();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.list_secrets(&namespace_path).await.infer_error()
|
||||
})
|
||||
page_token: Option<String>,
|
||||
page_limit: Option<u32>,
|
||||
) -> PyResult<crate::listing::NameListing> {
|
||||
Ok(crate::listing::NameListing::new(
|
||||
self.get_inner()?.list_secrets(
|
||||
&namespace_path.unwrap_or_default(),
|
||||
crate::listing::options(page_token, page_limit),
|
||||
),
|
||||
))
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, namespace_path=None))]
|
||||
@@ -969,27 +982,31 @@ impl Connection {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (namespace_path=None))]
|
||||
#[pyo3(signature = (namespace_path=None, *, page_token=None, page_limit=None))]
|
||||
pub fn list_views(
|
||||
self_: PyRef<'_, Self>,
|
||||
&self,
|
||||
namespace_path: Option<Vec<String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
let namespace_path = namespace_path.unwrap_or_default();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.list_views(&namespace_path).await.infer_error()
|
||||
})
|
||||
page_token: Option<String>,
|
||||
page_limit: Option<u32>,
|
||||
) -> PyResult<crate::listing::NameListing> {
|
||||
Ok(crate::listing::NameListing::new(
|
||||
self.get_inner()?.list_views(
|
||||
&namespace_path.unwrap_or_default(),
|
||||
crate::listing::options(page_token, page_limit),
|
||||
),
|
||||
))
|
||||
}
|
||||
|
||||
pub fn list_jobs(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let jobs = inner.list_jobs().await.infer_error()?;
|
||||
Ok(jobs
|
||||
.into_iter()
|
||||
.map(crate::job::JobInfo::from)
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
#[pyo3(signature = (*, page_token=None, page_limit=None))]
|
||||
pub fn list_jobs(
|
||||
&self,
|
||||
page_token: Option<String>,
|
||||
page_limit: Option<u32>,
|
||||
) -> PyResult<crate::listing::JobListing> {
|
||||
Ok(crate::listing::JobListing::new(
|
||||
self.get_inner()?
|
||||
.list_jobs(crate::listing::options(page_token, page_limit)),
|
||||
))
|
||||
}
|
||||
|
||||
pub fn cancel_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
|
||||
@@ -1203,6 +1220,7 @@ pub struct PyClientConfig {
|
||||
user_agent: String,
|
||||
retry_config: Option<PyClientRetryConfig>,
|
||||
timeout_config: Option<PyClientTimeoutConfig>,
|
||||
blob_request_concurrency: Option<usize>,
|
||||
extra_headers: Option<HashMap<String, String>>,
|
||||
id_delimiter: Option<String>,
|
||||
tls_config: Option<PyClientTlsConfig>,
|
||||
@@ -1288,6 +1306,7 @@ impl From<PyClientConfig> for lancedb::remote::ClientConfig {
|
||||
user_agent: value.user_agent,
|
||||
retry_config: value.retry_config.map(Into::into).unwrap_or_default(),
|
||||
timeout_config: value.timeout_config.map(Into::into).unwrap_or_default(),
|
||||
blob_request_concurrency: value.blob_request_concurrency,
|
||||
extra_headers: value.extra_headers.unwrap_or_default(),
|
||||
id_delimiter: value.id_delimiter,
|
||||
tls_config: value.tls_config.map(Into::into),
|
||||
|
||||
+8
-4
@@ -54,8 +54,6 @@ pub fn extract_index_params(source: &Option<Bound<'_, PyAny>>) -> PyResult<Lance
|
||||
.with_position(params.with_position)
|
||||
.lower_case(params.lower_case)
|
||||
.max_token_length(params.max_token_length)
|
||||
.remove_stop_words(params.remove_stop_words)
|
||||
.stem(params.stem)
|
||||
.ascii_folding(params.ascii_folding)
|
||||
.ngram_min_length(params.ngram_min_length)
|
||||
.ngram_max_length(params.ngram_max_length)
|
||||
@@ -65,6 +63,12 @@ pub fn extract_index_params(source: &Option<Bound<'_, PyAny>>) -> PyResult<Lance
|
||||
DocumentGranularity::try_from(params.document_granularity.as_str())
|
||||
.map_err(|err| PyValueError::new_err(err.to_string()))?,
|
||||
);
|
||||
if let Some(stem) = params.stem {
|
||||
inner_opts = inner_opts.stem(stem);
|
||||
}
|
||||
if let Some(remove_stop_words) = params.remove_stop_words {
|
||||
inner_opts = inner_opts.remove_stop_words(remove_stop_words);
|
||||
}
|
||||
if let Some(memory_limit) = params.memory_limit {
|
||||
inner_opts = inner_opts.memory_limit_mb(memory_limit);
|
||||
}
|
||||
@@ -215,8 +219,8 @@ struct FtsParams {
|
||||
language: String,
|
||||
max_token_length: Option<usize>,
|
||||
lower_case: bool,
|
||||
stem: bool,
|
||||
remove_stop_words: bool,
|
||||
stem: Option<bool>,
|
||||
remove_stop_words: Option<bool>,
|
||||
custom_stop_words: Option<Vec<String>>,
|
||||
ascii_folding: bool,
|
||||
ngram_min_length: u32,
|
||||
|
||||
+13
-2
@@ -16,12 +16,13 @@ use query::{FTSQuery, HybridQuery, Query, VectorQuery};
|
||||
use session::Session;
|
||||
use table::{
|
||||
AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, FtsToken,
|
||||
FunctionErrorFragment, FunctionErrorRecord, FunctionErrors, LsmWriteSpec, MergeResult,
|
||||
FunctionErrorFragment, FunctionErrorRecord, FunctionErrors, MemWalWriteSpec, MergeResult,
|
||||
PyBlobFile, RefreshColumnResult, RefreshMaterializedViewResult, Table,
|
||||
UpdateFieldMetadataResult, UpdateResult,
|
||||
};
|
||||
|
||||
pub mod arrow;
|
||||
pub mod authz;
|
||||
pub mod catalog;
|
||||
pub mod connection;
|
||||
pub mod error;
|
||||
@@ -29,6 +30,7 @@ pub mod expr;
|
||||
pub mod header;
|
||||
pub mod index;
|
||||
pub mod job;
|
||||
mod listing;
|
||||
pub mod namespace;
|
||||
pub mod oauth;
|
||||
pub mod otel;
|
||||
@@ -64,6 +66,15 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
env_logger::init_from_env(env);
|
||||
m.add_class::<Connection>()?;
|
||||
m.add_class::<catalog::Catalog>()?;
|
||||
m.add_class::<catalog::DatabaseNames>()?;
|
||||
m.add_class::<listing::NameListing>()?;
|
||||
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>()?;
|
||||
@@ -93,7 +104,7 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<UpdateFieldMetadataResult>()?;
|
||||
m.add_class::<AddResult>()?;
|
||||
m.add_class::<MergeResult>()?;
|
||||
m.add_class::<LsmWriteSpec>()?;
|
||||
m.add_class::<MemWalWriteSpec>()?;
|
||||
m.add_class::<DeleteResult>()?;
|
||||
m.add_class::<DropColumnsResult>()?;
|
||||
m.add_class::<UpdateResult>()?;
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use crate::error::PythonErrorExt;
|
||||
use crate::runtime::future_into_py;
|
||||
use futures::{StreamExt, future::poll_fn};
|
||||
use lancedb::listing::{Listing, ListingOptions};
|
||||
use pyo3::exceptions::PyStopAsyncIteration;
|
||||
use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pymethods};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
pub fn options(page_token: Option<String>, page_limit: Option<u32>) -> ListingOptions {
|
||||
let mut options = ListingOptions::default();
|
||||
options.page_token = page_token;
|
||||
options.page_limit = page_limit;
|
||||
options
|
||||
}
|
||||
|
||||
#[pyclass]
|
||||
pub struct NameListing {
|
||||
inner: Arc<Mutex<Listing<String>>>,
|
||||
next_lock: Arc<tokio::sync::Mutex<()>>,
|
||||
}
|
||||
impl NameListing {
|
||||
pub fn new(inner: Listing<String>) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(Mutex::new(inner)),
|
||||
next_lock: Arc::new(tokio::sync::Mutex::new(())),
|
||||
}
|
||||
}
|
||||
}
|
||||
#[pymethods]
|
||||
impl NameListing {
|
||||
fn num_page_results(&self, py: Python<'_>) -> usize {
|
||||
py.detach(|| self.inner.lock().unwrap().num_page_results())
|
||||
}
|
||||
fn page_token(&self, py: Python<'_>) -> Option<String> {
|
||||
py.detach(|| self.inner.lock().unwrap().page_token().map(str::to_owned))
|
||||
}
|
||||
fn __aiter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> {
|
||||
self_
|
||||
}
|
||||
fn __anext__(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
let next_lock = self_.next_lock.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let _guard = next_lock.lock().await;
|
||||
let value = poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx))
|
||||
.await
|
||||
.ok_or_else(|| PyStopAsyncIteration::new_err(""))?
|
||||
.infer_error()?;
|
||||
Ok(value)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass]
|
||||
pub struct FunctionListing {
|
||||
inner: Arc<Mutex<Listing<lancedb::function::FunctionVersion>>>,
|
||||
next_lock: Arc<tokio::sync::Mutex<()>>,
|
||||
}
|
||||
impl FunctionListing {
|
||||
pub fn new(inner: Listing<lancedb::function::FunctionVersion>) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(Mutex::new(inner)),
|
||||
next_lock: Arc::new(tokio::sync::Mutex::new(())),
|
||||
}
|
||||
}
|
||||
}
|
||||
#[pymethods]
|
||||
impl FunctionListing {
|
||||
fn num_page_results(&self, py: Python<'_>) -> usize {
|
||||
py.detach(|| self.inner.lock().unwrap().num_page_results())
|
||||
}
|
||||
fn page_token(&self, py: Python<'_>) -> Option<String> {
|
||||
py.detach(|| self.inner.lock().unwrap().page_token().map(str::to_owned))
|
||||
}
|
||||
fn __aiter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> {
|
||||
self_
|
||||
}
|
||||
fn __anext__(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
let next_lock = self_.next_lock.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let _guard = next_lock.lock().await;
|
||||
let value = poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx))
|
||||
.await
|
||||
.ok_or_else(|| PyStopAsyncIteration::new_err(""))?
|
||||
.infer_error()?;
|
||||
value.to_canonical_json().infer_error()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass]
|
||||
pub struct JobListing {
|
||||
inner: Arc<Mutex<Listing<lancedb::database::JobInfo>>>,
|
||||
next_lock: Arc<tokio::sync::Mutex<()>>,
|
||||
}
|
||||
impl JobListing {
|
||||
pub fn new(inner: Listing<lancedb::database::JobInfo>) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(Mutex::new(inner)),
|
||||
next_lock: Arc::new(tokio::sync::Mutex::new(())),
|
||||
}
|
||||
}
|
||||
}
|
||||
#[pymethods]
|
||||
impl JobListing {
|
||||
fn num_page_results(&self, py: Python<'_>) -> usize {
|
||||
py.detach(|| self.inner.lock().unwrap().num_page_results())
|
||||
}
|
||||
fn page_token(&self, py: Python<'_>) -> Option<String> {
|
||||
py.detach(|| self.inner.lock().unwrap().page_token().map(str::to_owned))
|
||||
}
|
||||
fn __aiter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> {
|
||||
self_
|
||||
}
|
||||
fn __anext__(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
let next_lock = self_.next_lock.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let _guard = next_lock.lock().await;
|
||||
let value = poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx))
|
||||
.await
|
||||
.ok_or_else(|| PyStopAsyncIteration::new_err(""))?
|
||||
.infer_error()?;
|
||||
Ok(crate::job::JobInfo::from(value))
|
||||
})
|
||||
}
|
||||
}
|
||||
+58
-28
@@ -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)
|
||||
@@ -180,6 +180,9 @@ impl<'py> IntoPyObject<'py> for PyLanceDB<FtsQuery> {
|
||||
.expect("Failed to import namespace");
|
||||
|
||||
match self.0 {
|
||||
FtsQuery::CombinedFields(_) => Err(PyValueError::new_err(
|
||||
"Combined-fields FTS queries are not yet supported in Python",
|
||||
)),
|
||||
FtsQuery::Match(query) => {
|
||||
let kwargs = PyDict::new(py);
|
||||
kwargs.set_item("boost", query.boost)?;
|
||||
@@ -196,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);
|
||||
@@ -330,13 +333,13 @@ pub struct PyQueryRequest {
|
||||
pub select_source_columns: Option<HashMap<String, String>>,
|
||||
pub fast_search: Option<bool>,
|
||||
pub with_row_id: Option<bool>,
|
||||
pub use_lsm: Option<bool>,
|
||||
pub use_memwal: Option<bool>,
|
||||
pub column: Option<String>,
|
||||
pub query_vector: Option<PyQueryVectors>,
|
||||
pub nprobes: Option<usize>,
|
||||
pub minimum_nprobes: Option<usize>,
|
||||
// None means user did not set it and default should be used (currently 20)
|
||||
// Some(0) means user set it to None and there is no limit
|
||||
pub maximum_nprobes: Option<usize>,
|
||||
pub search_effort: Option<f64>,
|
||||
pub lower_bound: Option<f32>,
|
||||
pub upper_bound: Option<f32>,
|
||||
pub ef: Option<usize>,
|
||||
@@ -363,11 +366,13 @@ impl From<AnyQuery> for PyQueryRequest {
|
||||
select: PySelect(query_request.select),
|
||||
fast_search: Some(query_request.fast_search),
|
||||
with_row_id: Some(query_request.with_row_id),
|
||||
use_lsm: query_request.use_lsm,
|
||||
use_memwal: query_request.use_memwal,
|
||||
column: None,
|
||||
query_vector: None,
|
||||
nprobes: None,
|
||||
minimum_nprobes: None,
|
||||
maximum_nprobes: None,
|
||||
search_effort: None,
|
||||
lower_bound: None,
|
||||
upper_bound: None,
|
||||
ef: None,
|
||||
@@ -390,14 +395,13 @@ impl From<AnyQuery> for PyQueryRequest {
|
||||
select: PySelect(vector_query.base.select),
|
||||
fast_search: Some(vector_query.base.fast_search),
|
||||
with_row_id: Some(vector_query.base.with_row_id),
|
||||
use_lsm: vector_query.base.use_lsm,
|
||||
use_memwal: vector_query.base.use_memwal,
|
||||
column: vector_query.column,
|
||||
query_vector: Some(PyQueryVectors(vector_query.query_vector)),
|
||||
minimum_nprobes: Some(vector_query.minimum_nprobes),
|
||||
maximum_nprobes: match vector_query.maximum_nprobes {
|
||||
None => Some(0),
|
||||
Some(value) => Some(value),
|
||||
},
|
||||
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,
|
||||
@@ -536,8 +540,8 @@ impl Query {
|
||||
self.inner = self.inner.clone().fast_search();
|
||||
}
|
||||
|
||||
pub fn use_lsm(&mut self, enable: bool) {
|
||||
self.inner = self.inner.clone().use_lsm(enable);
|
||||
pub fn use_memwal(&mut self, enable: bool) {
|
||||
self.inner = self.inner.clone().use_memwal(enable);
|
||||
}
|
||||
|
||||
pub fn with_row_id(&mut self) {
|
||||
@@ -702,8 +706,8 @@ impl TakeQuery {
|
||||
self.inner = self.inner.clone().with_row_id();
|
||||
}
|
||||
|
||||
pub fn use_lsm(&mut self, enable: bool) {
|
||||
self.inner = self.inner.clone().use_lsm(enable);
|
||||
pub fn use_memwal(&mut self, enable: bool) {
|
||||
self.inner = self.inner.clone().use_memwal(enable);
|
||||
}
|
||||
|
||||
#[pyo3(signature = ())]
|
||||
@@ -815,8 +819,8 @@ impl FTSQuery {
|
||||
self.inner = self.inner.clone().fast_search();
|
||||
}
|
||||
|
||||
pub fn use_lsm(&mut self, enable: bool) {
|
||||
self.inner = self.inner.clone().use_lsm(enable);
|
||||
pub fn use_memwal(&mut self, enable: bool) {
|
||||
self.inner = self.inner.clone().use_memwal(enable);
|
||||
}
|
||||
|
||||
pub fn with_row_id(&mut self) {
|
||||
@@ -966,8 +970,8 @@ impl VectorQuery {
|
||||
self.inner = self.inner.clone().fast_search();
|
||||
}
|
||||
|
||||
pub fn use_lsm(&mut self, enable: bool) {
|
||||
self.inner = self.inner.clone().use_lsm(enable);
|
||||
pub fn use_memwal(&mut self, enable: bool) {
|
||||
self.inner = self.inner.clone().use_memwal(enable);
|
||||
}
|
||||
|
||||
pub fn with_row_id(&mut self) {
|
||||
@@ -992,8 +996,18 @@ impl VectorQuery {
|
||||
self.inner = self.inner.clone().refine_factor(refine_factor);
|
||||
}
|
||||
|
||||
pub fn nprobes(&mut self, nprobe: u32) {
|
||||
self.inner = self.inner.clone().nprobes(nprobe as usize);
|
||||
pub fn nprobes(&mut self, nprobe: u32) -> PyResult<()> {
|
||||
self.inner = self.inner.clone().nprobes(nprobe as usize).infer_error()?;
|
||||
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<()> {
|
||||
@@ -1164,9 +1178,9 @@ impl HybridQuery {
|
||||
self.inner_fts.postfilter();
|
||||
}
|
||||
|
||||
pub fn use_lsm(&mut self, enable: bool) {
|
||||
self.inner_vec.use_lsm(enable);
|
||||
self.inner_fts.use_lsm(enable);
|
||||
pub fn use_memwal(&mut self, enable: bool) {
|
||||
self.inner_vec.use_memwal(enable);
|
||||
self.inner_fts.use_memwal(enable);
|
||||
}
|
||||
|
||||
pub fn add_query_vector(&mut self, vector: Bound<'_, PyAny>) -> PyResult<()> {
|
||||
@@ -1185,8 +1199,20 @@ impl HybridQuery {
|
||||
self.inner_vec.refine_factor(refine_factor);
|
||||
}
|
||||
|
||||
pub fn nprobes(&mut self, nprobe: u32) {
|
||||
self.inner_vec.nprobes(nprobe);
|
||||
pub fn nprobes(&mut self, nprobe: u32) -> PyResult<()> {
|
||||
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) {
|
||||
@@ -1233,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;
|
||||
|
||||
+175
-112
@@ -41,16 +41,16 @@ use pyo3::{
|
||||
|
||||
mod scannable;
|
||||
|
||||
/// Convert `LsmStats` to a Python dict, preserving the per-bucket list.
|
||||
/// Convert `MemWalStats` to a Python dict, preserving the per-table-shard list.
|
||||
///
|
||||
/// Deliberately not flattened to a table-level summary: a table is N
|
||||
/// buckets on one node, and the per-bucket detail is the reason the
|
||||
/// endpoint exists — flattening hides the single hot bucket someone opened
|
||||
/// table shards on one node, and the per-shard detail is the reason the
|
||||
/// endpoint exists — flattening hides the single hot table shard someone opened
|
||||
/// it to find.
|
||||
fn lsm_stats_to_py(py: Python<'_>, stats: &lancedb::table::LsmStats) -> PyResult<Py<PyDict>> {
|
||||
fn memwal_stats_to_py(py: Python<'_>, stats: &lancedb::table::MemWalStats) -> PyResult<Py<PyDict>> {
|
||||
let out = PyDict::new(py);
|
||||
let buckets = PyList::empty(py);
|
||||
for b in &stats.buckets {
|
||||
let table_shards = PyList::empty(py);
|
||||
for b in &stats.table_shards {
|
||||
let e = PyDict::new(py);
|
||||
e.set_item("shard_id", &b.shard_id)?;
|
||||
e.set_item("status", &b.status)?;
|
||||
@@ -66,15 +66,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(
|
||||
@@ -96,9 +96,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())
|
||||
}
|
||||
|
||||
@@ -399,20 +399,20 @@ fn fmt_maintained(maintained: &Option<Vec<String>>) -> String {
|
||||
/// Constructed via the `bucket(...)`, `identity(...)`, or `unsharded()`
|
||||
/// classmethods, then optionally chain `with_maintained_indexes(...)` and
|
||||
/// `with_writer_config_defaults(...)`. A fresh spec maintains every index the
|
||||
/// MemWAL supports, resolved on install.
|
||||
/// table has, including ones created later.
|
||||
#[pyclass(module = "lancedb._lancedb", from_py_object)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct LsmWriteSpec {
|
||||
inner: lancedb::table::LsmWriteSpec,
|
||||
pub struct MemWalWriteSpec {
|
||||
inner: lancedb::table::MemWalWriteSpec,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl LsmWriteSpec {
|
||||
impl MemWalWriteSpec {
|
||||
/// Hash-bucket sharding by the unenforced primary key column.
|
||||
#[staticmethod]
|
||||
pub fn bucket(column: String, num_buckets: u32) -> Self {
|
||||
Self {
|
||||
inner: lancedb::table::LsmWriteSpec::bucket(column, num_buckets),
|
||||
inner: lancedb::table::MemWalWriteSpec::bucket(column, num_buckets),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -426,7 +426,7 @@ impl LsmWriteSpec {
|
||||
#[staticmethod]
|
||||
pub fn identity(column: String) -> Self {
|
||||
Self {
|
||||
inner: lancedb::table::LsmWriteSpec::identity(column),
|
||||
inner: lancedb::table::MemWalWriteSpec::identity(column),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -435,13 +435,13 @@ impl LsmWriteSpec {
|
||||
#[staticmethod]
|
||||
pub fn unsharded() -> Self {
|
||||
Self {
|
||||
inner: lancedb::table::LsmWriteSpec::unsharded(),
|
||||
inner: lancedb::table::MemWalWriteSpec::unsharded(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Set which indexes the MemWAL maintains. `None` (the default)
|
||||
/// resolves every supported index on install; a list is verbatim,
|
||||
/// and an empty list maintains nothing.
|
||||
/// Set which indexes the MemWAL 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))]
|
||||
pub fn with_maintained_indexes(&self, indexes: Option<Vec<String>>) -> Self {
|
||||
Self {
|
||||
@@ -459,33 +459,33 @@ impl LsmWriteSpec {
|
||||
|
||||
pub fn __repr__(&self) -> String {
|
||||
match &self.inner {
|
||||
lancedb::table::LsmWriteSpec::Bucket {
|
||||
lancedb::table::MemWalWriteSpec::Bucket {
|
||||
column,
|
||||
num_buckets,
|
||||
maintained_indexes,
|
||||
writer_config_defaults,
|
||||
} => format!(
|
||||
"LsmWriteSpec.bucket(column={:?}, num_buckets={}, maintained_indexes={}, writer_config_defaults={:?})",
|
||||
"MemWalWriteSpec.bucket(column={:?}, num_buckets={}, maintained_indexes={}, writer_config_defaults={:?})",
|
||||
column,
|
||||
num_buckets,
|
||||
fmt_maintained(maintained_indexes),
|
||||
writer_config_defaults,
|
||||
),
|
||||
lancedb::table::LsmWriteSpec::Identity {
|
||||
lancedb::table::MemWalWriteSpec::Identity {
|
||||
column,
|
||||
maintained_indexes,
|
||||
writer_config_defaults,
|
||||
} => format!(
|
||||
"LsmWriteSpec.identity(column={:?}, maintained_indexes={}, writer_config_defaults={:?})",
|
||||
"MemWalWriteSpec.identity(column={:?}, maintained_indexes={}, writer_config_defaults={:?})",
|
||||
column,
|
||||
fmt_maintained(maintained_indexes),
|
||||
writer_config_defaults,
|
||||
),
|
||||
lancedb::table::LsmWriteSpec::Unsharded {
|
||||
lancedb::table::MemWalWriteSpec::Unsharded {
|
||||
maintained_indexes,
|
||||
writer_config_defaults,
|
||||
} => format!(
|
||||
"LsmWriteSpec.unsharded(maintained_indexes={}, writer_config_defaults={:?})",
|
||||
"MemWalWriteSpec.unsharded(maintained_indexes={}, writer_config_defaults={:?})",
|
||||
fmt_maintained(maintained_indexes),
|
||||
writer_config_defaults,
|
||||
),
|
||||
@@ -497,9 +497,9 @@ impl LsmWriteSpec {
|
||||
#[getter]
|
||||
pub fn spec_type(&self) -> &'static str {
|
||||
match &self.inner {
|
||||
lancedb::table::LsmWriteSpec::Bucket { .. } => "bucket",
|
||||
lancedb::table::LsmWriteSpec::Identity { .. } => "identity",
|
||||
lancedb::table::LsmWriteSpec::Unsharded { .. } => "unsharded",
|
||||
lancedb::table::MemWalWriteSpec::Bucket { .. } => "bucket",
|
||||
lancedb::table::MemWalWriteSpec::Identity { .. } => "identity",
|
||||
lancedb::table::MemWalWriteSpec::Unsharded { .. } => "unsharded",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -507,9 +507,9 @@ impl LsmWriteSpec {
|
||||
#[getter]
|
||||
pub fn column(&self) -> Option<String> {
|
||||
match &self.inner {
|
||||
lancedb::table::LsmWriteSpec::Bucket { column, .. }
|
||||
| lancedb::table::LsmWriteSpec::Identity { column, .. } => Some(column.clone()),
|
||||
lancedb::table::LsmWriteSpec::Unsharded { .. } => None,
|
||||
lancedb::table::MemWalWriteSpec::Bucket { column, .. }
|
||||
| lancedb::table::MemWalWriteSpec::Identity { column, .. } => Some(column.clone()),
|
||||
lancedb::table::MemWalWriteSpec::Unsharded { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -517,7 +517,7 @@ impl LsmWriteSpec {
|
||||
#[getter]
|
||||
pub fn num_buckets(&self) -> Option<u32> {
|
||||
match &self.inner {
|
||||
lancedb::table::LsmWriteSpec::Bucket { num_buckets, .. } => Some(*num_buckets),
|
||||
lancedb::table::MemWalWriteSpec::Bucket { num_buckets, .. } => Some(*num_buckets),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
@@ -535,14 +535,14 @@ impl LsmWriteSpec {
|
||||
}
|
||||
}
|
||||
|
||||
impl From<LsmWriteSpec> for lancedb::table::LsmWriteSpec {
|
||||
fn from(spec: LsmWriteSpec) -> Self {
|
||||
impl From<MemWalWriteSpec> for lancedb::table::MemWalWriteSpec {
|
||||
fn from(spec: MemWalWriteSpec) -> Self {
|
||||
spec.inner
|
||||
}
|
||||
}
|
||||
|
||||
impl From<lancedb::table::LsmWriteSpec> for LsmWriteSpec {
|
||||
fn from(inner: lancedb::table::LsmWriteSpec) -> Self {
|
||||
impl From<lancedb::table::MemWalWriteSpec> for MemWalWriteSpec {
|
||||
fn from(inner: lancedb::table::MemWalWriteSpec) -> Self {
|
||||
Self { inner }
|
||||
}
|
||||
}
|
||||
@@ -724,8 +724,6 @@ impl From<lancedb::RefreshMaterializedViewResult> for RefreshMaterializedViewRes
|
||||
fn from(result: lancedb::RefreshMaterializedViewResult) -> Self {
|
||||
let mode = match result.mode {
|
||||
lancedb::RefreshMode::Rebuild => "rebuild",
|
||||
lancedb::RefreshMode::Incremental => "incremental",
|
||||
lancedb::RefreshMode::NoOp => "no_op",
|
||||
};
|
||||
Self {
|
||||
mode: mode.to_string(),
|
||||
@@ -831,6 +829,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>> {
|
||||
@@ -839,7 +845,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())
|
||||
}
|
||||
|
||||
@@ -847,10 +853,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()))
|
||||
})
|
||||
}
|
||||
@@ -873,7 +876,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 {
|
||||
@@ -886,7 +889,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> {
|
||||
@@ -895,7 +898,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 {
|
||||
@@ -912,17 +915,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())
|
||||
}
|
||||
}
|
||||
@@ -957,8 +982,8 @@ impl From<LanceDbFtsToken> for FtsToken {
|
||||
language = "English".to_string(),
|
||||
max_token_length = Some(40),
|
||||
lower_case = true,
|
||||
stem = true,
|
||||
remove_stop_words = true,
|
||||
stem = None,
|
||||
remove_stop_words = None,
|
||||
custom_stop_words = None,
|
||||
ascii_folding = true,
|
||||
ngram_min_length = 3,
|
||||
@@ -972,15 +997,15 @@ pub fn tokenize(
|
||||
language: String,
|
||||
max_token_length: Option<u32>,
|
||||
lower_case: bool,
|
||||
stem: bool,
|
||||
remove_stop_words: bool,
|
||||
stem: Option<bool>,
|
||||
remove_stop_words: Option<bool>,
|
||||
custom_stop_words: Option<Vec<String>>,
|
||||
ascii_folding: bool,
|
||||
ngram_min_length: u32,
|
||||
ngram_max_length: u32,
|
||||
prefix_only: bool,
|
||||
) -> PyResult<Vec<FtsToken>> {
|
||||
let params = FtsIndexBuilder::default()
|
||||
let mut params = FtsIndexBuilder::default()
|
||||
.base_tokenizer(base_tokenizer)
|
||||
.language(&language)
|
||||
.map_err(|_| {
|
||||
@@ -991,13 +1016,17 @@ pub fn tokenize(
|
||||
})?
|
||||
.max_token_length(max_token_length.map(|value| value as usize))
|
||||
.lower_case(lower_case)
|
||||
.stem(stem)
|
||||
.remove_stop_words(remove_stop_words)
|
||||
.ascii_folding(ascii_folding)
|
||||
.ngram_min_length(ngram_min_length)
|
||||
.ngram_max_length(ngram_max_length)
|
||||
.ngram_prefix_only(prefix_only)
|
||||
.custom_stop_words(custom_stop_words);
|
||||
if let Some(stem) = stem {
|
||||
params = params.stem(stem);
|
||||
}
|
||||
if let Some(remove_stop_words) = remove_stop_words {
|
||||
params = params.remove_stop_words(remove_stop_words);
|
||||
}
|
||||
let tokens = lancedb_tokenize(&query, ¶ms).infer_error()?;
|
||||
Ok(tokens.into_iter().map(FtsToken::from).collect())
|
||||
}
|
||||
@@ -1047,6 +1076,11 @@ impl Table {
|
||||
self.inner.is_some()
|
||||
}
|
||||
|
||||
/// Whether this table has a local Lance dataset that the client may open.
|
||||
pub fn _is_native(&self) -> PyResult<bool> {
|
||||
Ok(self.inner_ref()?.as_native().is_some())
|
||||
}
|
||||
|
||||
/// Closes the table, releasing any resources associated with it.
|
||||
pub fn close(&mut self) {
|
||||
self.inner.take();
|
||||
@@ -1400,6 +1434,7 @@ impl Table {
|
||||
let dict = PyDict::new(py);
|
||||
dict.set_item("total_bytes", stats.total_bytes)?;
|
||||
dict.set_item("num_rows", stats.num_rows)?;
|
||||
dict.set_item("num_deleted_rows", stats.num_deleted_rows)?;
|
||||
dict.set_item("num_indices", stats.num_indices)?;
|
||||
|
||||
let fragment_stats = PyDict::new(py);
|
||||
@@ -1586,14 +1621,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
|
||||
@@ -1719,8 +1770,8 @@ impl Table {
|
||||
if let Some(use_index) = parameters.use_index {
|
||||
builder.use_index(use_index);
|
||||
}
|
||||
if let Some(use_lsm) = parameters.use_lsm {
|
||||
builder.use_lsm(use_lsm);
|
||||
if let Some(use_memwal) = parameters.use_memwal {
|
||||
builder.use_memwal(use_memwal);
|
||||
}
|
||||
if let Some(validate_single_shard) = parameters.validate_single_shard {
|
||||
builder.validate_single_shard(validate_single_shard);
|
||||
@@ -1745,29 +1796,29 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn set_lsm_write_spec<'a>(
|
||||
pub fn enable_memwal<'a>(
|
||||
self_: PyRef<'a, Self>,
|
||||
spec: LsmWriteSpec,
|
||||
spec: MemWalWriteSpec,
|
||||
) -> PyResult<Bound<'a, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
let native_spec = lancedb::table::LsmWriteSpec::from(spec);
|
||||
let native_spec = lancedb::table::MemWalWriteSpec::from(spec);
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.set_lsm_write_spec(native_spec).await.infer_error()
|
||||
inner.enable_memwal(native_spec).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn unset_lsm_write_spec(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
pub fn disable_memwal(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.unset_lsm_write_spec().await.infer_error()
|
||||
inner.disable_memwal().await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn get_lsm_write_spec(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
pub fn get_memwal_spec(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let spec = inner.get_lsm_write_spec().await.infer_error()?;
|
||||
Ok(spec.map(LsmWriteSpec::from))
|
||||
let spec = inner.get_memwal_spec().await.infer_error()?;
|
||||
Ok(spec.map(MemWalWriteSpec::from))
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1786,55 +1837,54 @@ impl Table {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Converge the table's LSM write path into its base table.
|
||||
/// Converge the table's MemWAL write path into its base table.
|
||||
///
|
||||
/// Best-effort: with writes flowing, new rows may land after the last
|
||||
/// pass. Errors if the table stops making progress.
|
||||
pub fn checkpoint_lsm(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
pub fn checkpoint_memwal(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.checkpoint_lsm().await.infer_error()
|
||||
inner.checkpoint_memwal().await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
/// Seal every bucket's active memtable into L0.
|
||||
pub fn flush_lsm(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(
|
||||
self_.py(),
|
||||
async move { inner.flush_lsm().await.infer_error() },
|
||||
)
|
||||
}
|
||||
|
||||
/// Trigger a background L0 → base pass per bucket. Returns once the
|
||||
/// passes are dispatched, not once they finish — watch `get_lsm_stats`.
|
||||
pub fn compact_lsm(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
/// Freeze every table shard's active memtable into an SSTable.
|
||||
pub fn flush_memwal(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.compact_lsm().await.infer_error()
|
||||
inner.flush_memwal().await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
/// Live LSM state, or `None` when the LSM write path is not enabled.
|
||||
#[pyo3(signature = (include_generation_rows=false))]
|
||||
pub fn get_lsm_stats(
|
||||
/// Trigger a background SSTable compaction pass per table shard. Returns once the
|
||||
/// passes are dispatched, not once they finish — watch `get_memwal_stats`.
|
||||
pub fn compact_memwal(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.compact_memwal().await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
/// Live MemWAL state, or `None` when the MemWAL write path is not enabled.
|
||||
#[pyo3(signature = (include_sstable_rows=false))]
|
||||
pub fn get_memwal_stats(
|
||||
self_: PyRef<'_, Self>,
|
||||
include_generation_rows: bool,
|
||||
include_sstable_rows: bool,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let stats = inner
|
||||
.get_lsm_stats(include_generation_rows)
|
||||
.get_memwal_stats(include_sstable_rows)
|
||||
.await
|
||||
.infer_error()?;
|
||||
Python::attach(|py| stats.map(|s| lsm_stats_to_py(py, &s)).transpose())
|
||||
Python::attach(|py| stats.map(|s| memwal_stats_to_py(py, &s)).transpose())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn close_lsm_writers(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
pub fn close_memwal_writers(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.close_lsm_writers().await.infer_error()
|
||||
inner.close_memwal_writers().await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1843,7 +1893,10 @@ impl Table {
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.as_native()
|
||||
.ok_or_else(|| PyValueError::new_err("This cannot be run on a remote table"))?
|
||||
.ok_or_else(|| lancedb::Error::NotSupported {
|
||||
message: "uses_v2_manifest_paths is not supported for remote tables.".into(),
|
||||
})
|
||||
.infer_error()?
|
||||
.uses_v2_manifest_paths()
|
||||
.await
|
||||
.infer_error()
|
||||
@@ -1855,7 +1908,10 @@ impl Table {
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.as_native()
|
||||
.ok_or_else(|| PyValueError::new_err("This cannot be run on a remote table"))?
|
||||
.ok_or_else(|| lancedb::Error::NotSupported {
|
||||
message: "migrate_manifest_paths_v2 is not supported for remote tables.".into(),
|
||||
})
|
||||
.infer_error()?
|
||||
.migrate_manifest_paths_v2()
|
||||
.await
|
||||
.infer_error()
|
||||
@@ -1951,10 +2007,9 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (full=false, source_version=None))]
|
||||
#[pyo3(signature = (source_version=None))]
|
||||
pub fn refresh_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
full: bool,
|
||||
source_version: Option<u64>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
@@ -1962,7 +2017,7 @@ impl Table {
|
||||
let view = lancedb::MaterializedView::from_table(inner)
|
||||
.await
|
||||
.infer_error()?;
|
||||
let mut builder = view.refresh().full(full);
|
||||
let mut builder = view.refresh();
|
||||
if let Some(version) = source_version {
|
||||
builder = builder.source_version(version);
|
||||
}
|
||||
@@ -1971,10 +2026,9 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (full=false, source_version=None))]
|
||||
#[pyo3(signature = (source_version=None))]
|
||||
pub fn refresh_materialized_view_async(
|
||||
self_: PyRef<'_, Self>,
|
||||
full: bool,
|
||||
source_version: Option<u64>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
@@ -1982,7 +2036,7 @@ impl Table {
|
||||
let view = lancedb::MaterializedView::from_table(inner)
|
||||
.await
|
||||
.infer_error()?;
|
||||
let mut builder = view.refresh().full(full);
|
||||
let mut builder = view.refresh();
|
||||
if let Some(version) = source_version {
|
||||
builder = builder.source_version(version);
|
||||
}
|
||||
@@ -1997,11 +2051,9 @@ impl Table {
|
||||
let view = lancedb::MaterializedView::from_table(inner)
|
||||
.await
|
||||
.infer_error()?;
|
||||
view.definition().to_json().map_err(|err| {
|
||||
PyRuntimeError::new_err(format!(
|
||||
"failed to serialize materialized-view definition: {err}"
|
||||
))
|
||||
})
|
||||
Ok(lancedb::materialized_view::definition_metadata_from_sql(
|
||||
view.definition_sql(),
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
@@ -2031,6 +2083,17 @@ impl Table {
|
||||
let alterations = alterations
|
||||
.iter()
|
||||
.map(|alteration| {
|
||||
for key in alteration.keys().iter() {
|
||||
let key = key.extract::<String>()?;
|
||||
if !matches!(
|
||||
key.as_str(),
|
||||
"path" | "rename" | "name" | "nullable" | "data_type"
|
||||
) {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"Unknown column alteration key '{key}'"
|
||||
)));
|
||||
}
|
||||
}
|
||||
let path = alteration
|
||||
.get_item("path")?
|
||||
.ok_or_else(|| PyValueError::new_err("Missing path"))?
|
||||
@@ -2153,7 +2216,7 @@ pub struct MergeInsertParams {
|
||||
when_not_matched_by_source_condition_expr: Option<PyExpr>,
|
||||
timeout: Option<std::time::Duration>,
|
||||
use_index: Option<bool>,
|
||||
use_lsm: Option<bool>,
|
||||
use_memwal: Option<bool>,
|
||||
validate_single_shard: Option<bool>,
|
||||
}
|
||||
|
||||
|
||||
Generated
+42
-16
@@ -1849,19 +1849,19 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace"
|
||||
version = "0.8.6"
|
||||
version = "0.13.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "lance-namespace-urllib3-client" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/af/12/f7ab93b29be3edbf5fc3610714bf2d06088e7f4524bfb38dfd6852458b08/lance_namespace-0.8.6.tar.gz", hash = "sha256:18232e721c8188145f4ec9389cc2dfbeeabf54a619d94885ea1b3375bee9f4af", size = 11529, upload-time = "2026-06-12T17:36:41.651Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/c0/e7/d5d46594678ee479c0eda830c47b2f5c46133bd100e84e9aa6e01306eca7/lance_namespace-0.13.0.tar.gz", hash = "sha256:24554a0997bdb39595c6e4cb3ac6722069f6cd3bd1a74e7acbba7bfaf774a40d", size = 11636, upload-time = "2026-09-14T02:47:03.871523Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a0/1b/5b1668ee2dc8910965f390640359112a31157092fcf8e000b89c79b58708/lance_namespace-0.8.6-py3-none-any.whl", hash = "sha256:571eae34f9aad70e5b05020416c2860889b9ec82993ccd0eb015e7b39c3ea309", size = 13383, upload-time = "2026-06-12T17:36:43.456Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e7/47/cfa33cca1ba7c749fd2918cfc5c8ded788f378cc6424d23e4fead5a14125/lance_namespace-0.13.0-py3-none-any.whl", hash = "sha256:438c7b17aef421c21c138196e715f2510d62f07e865372047f87c1e75e618c7a", size = 13507, upload-time = "2026-09-14T02:47:01.765655Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace-urllib3-client"
|
||||
version = "0.8.6"
|
||||
version = "0.13.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "pydantic" },
|
||||
@@ -1869,9 +1869,9 @@ dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
{ name = "urllib3" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/c7/80/fb224b4a89c1c1638cde949cb6cce6c3aca7759effbfea46a3d9c3960b21/lance_namespace_urllib3_client-0.8.6.tar.gz", hash = "sha256:b6fb1d306e74a7576e5309919020be744527de484a63dbf5eed10f8b368548df", size = 228772, upload-time = "2026-06-12T17:36:42.609Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/69/25/9aaa4a5e8999693fb0f227c2c0c4b97bd8f0539066408cdff0b89d41b2d5/lance_namespace_urllib3_client-0.13.0.tar.gz", hash = "sha256:1e8a79c6e4e6277033597fd76aa0e2f33d909ca1436c935936e9e569221f43ef", size = 238348, upload-time = "2026-09-14T02:47:04.521878Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/c5/90/1e27de15cd1b16785a1c7312beb0a59e75c8344a815f600f58173a565bd1/lance_namespace_urllib3_client-0.8.6-py3-none-any.whl", hash = "sha256:9d78249c3fb15aa3d15d668f78f04a275af3d08d800a7027492f37996ac4968b", size = 369950, upload-time = "2026-06-12T17:36:40.438Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/60/4c/7b8f0712a7fe1b8655b711342bc4989696635a79e7027ee56ba3cae23e99/lance_namespace_urllib3_client-0.13.0-py3-none-any.whl", hash = "sha256:fb361eb4f6c7f2d1f9e92809609657b73e38241393e9c4516e12d6b6ee643a0a", size = 406848, upload-time = "2026-09-14T02:47:02.715954Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1900,6 +1900,7 @@ clip = [
|
||||
{ name = "torch" },
|
||||
]
|
||||
dev = [
|
||||
{ name = "maturin" },
|
||||
{ name = "pre-commit" },
|
||||
{ name = "pyright" },
|
||||
{ name = "ruff" },
|
||||
@@ -1980,7 +1981,8 @@ requires-dist = [
|
||||
{ name = "huggingface-hub", marker = "extra == 'embeddings'", specifier = ">=0.19.0" },
|
||||
{ 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.3.2" },
|
||||
{ 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 = "==9.0.0" },
|
||||
{ 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,8 +3938,8 @@ crypto = [
|
||||
|
||||
[[package]]
|
||||
name = "pylance"
|
||||
version = "9.0.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
version = "14.0.0b12"
|
||||
source = { registry = "https://pypi.fury.io/lance-format/" }
|
||||
dependencies = [
|
||||
{ name = "lance-namespace" },
|
||||
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
|
||||
@@ -3922,12 +3948,12 @@ dependencies = [
|
||||
{ name = "pyarrow" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/e9/be/45733acd64801991852aac8e658601fd8fc12f76ceb81e57fca690896b90/pylance-9.0.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:8257213501d3298c5b6a344d60938e4bbe4de9f00cd3265371a56d1dc3dd15ca", size = 68377982, upload-time = "2026-07-24T16:53:45.247Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d0/3e/1ef707cb215cc7268c63ad84a91344ad6313d3984b343eeb19a9b708698e/pylance-9.0.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:804eedfa1fda2e703cca8580c76f0b44a1b849b725a8908c06f8acfca811732f", size = 71844362, upload-time = "2026-07-24T16:56:07.6Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c8/fb/a499e5c53ddb75c7de44100fd2bb1f7cc735966200d20292eed7af9ef552/pylance-9.0.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0a0b75595e3766c1d5f4c90abdc52337b4de60f1a52d123a4f8c4e5bcbbbfa8f", size = 75663088, upload-time = "2026-07-24T17:10:31.283Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e9/80/0714e09f64a68dbdf62558955e737a7436df763b9834a6aba506861b5352/pylance-9.0.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:d2f69c5c390ae3710c35a429905fe15f769951777fafb1b72a359e035ec121f7", size = 71866858, upload-time = "2026-07-24T16:56:35.937Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/4b/3c/78d3a6d6ca0d843b7c3ac0c30d9cd2cf4635b0c2cabe6ec66583d5bfc1a1/pylance-9.0.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:836268a7832d62f3d5ccbe1c3fca239621971d90297bcfff14a70b3cb6842aa8", size = 75642656, upload-time = "2026-07-24T17:13:11.879Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/cf/c1/dc9c9a31e171530ec0add024d922488c046437d32a7087c29e254a7eacc7/pylance-9.0.0-cp310-abi3-win_amd64.whl", hash = "sha256:96441d27a5ed3805300388ccf8f31835cd280c98751f80b6a1a11dcd6808fc43", size = 81668288, upload-time = "2026-07-24T17:05:38.707Z" },
|
||||
{ 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]]
|
||||
|
||||
Reference in New Issue
Block a user