Merge remote-tracking branch 'origin/main' into gatekeeper/fix-2923-1

This commit is contained in:
Gatefixer
2026-08-08 11:45:28 +00:00
78 changed files with 4162 additions and 495 deletions
+3 -3
View File
@@ -1,6 +1,6 @@
[package]
name = "lancedb-python"
version = "0.37.1-beta.0"
version = "0.37.1-beta.1"
publish = false
edition.workspace = true
description = "Python bindings for LanceDB"
@@ -26,7 +26,7 @@ lance-namespace-impls.workspace = true
lance-io.workspace = true
env_logger.workspace = true
log.workspace = true
pyo3 = { version = "0.28", features = ["extension-module", "abi3-py39", "chrono"] }
pyo3 = { version = "0.28", features = ["extension-module", "abi3-py310", "chrono"] }
chrono = { version = "0.4", default-features = false, features = ["clock"] }
pyo3-async-runtimes = { version = "0.28", features = [
"attributes",
@@ -43,7 +43,7 @@ libc = "0.2"
[build-dependencies]
pyo3-build-config = { version = "0.28", features = [
"extension-module",
"abi3-py39",
"abi3-py310",
] }
[features]
+2 -1
View File
@@ -60,7 +60,7 @@ tests = [
"pytest-asyncio>=0.21",
"duckdb>=0.9.0",
"pytz>=2023.3",
"polars>=0.19, <=1.3.0",
"polars>=0.19, <=1.32.3",
"pyarrow<25",
"pyarrow-stubs>=16.0",
"pylance==9.0.0rc1",
@@ -140,6 +140,7 @@ include = [
"python/lancedb/remote/errors.py",
"python/lancedb/embeddings/__init__.py",
"python/lancedb/_lancedb.pyi",
"python/type_tests/connect.py",
]
exclude = ["python/tests/"]
pythonVersion = "3.13"
+11 -4
View File
@@ -355,6 +355,10 @@ class Table:
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 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: ...
@property
def tags(self) -> Tags: ...
@@ -649,9 +653,10 @@ class LsmWriteSpec:
def identity(column: str) -> "LsmWriteSpec": ...
@staticmethod
def unsharded() -> "LsmWriteSpec": ...
def with_maintained_indexes(self, indexes: List[str]) -> "LsmWriteSpec":
"""Return a copy of this spec asking the MemWAL to keep the named
indexes up to date as rows are appended."""
def with_maintained_indexes(self, indexes: Optional[List[str]]) -> "LsmWriteSpec":
"""Set which indexes the MemWAL keeps up to date. None resolves every
index on the table at install, failing if one cannot be maintained;
a list is verbatim, empty means none."""
...
def with_writer_config_defaults(self, defaults: Dict[str, str]) -> "LsmWriteSpec":
"""Return a copy of this spec recording the given default
@@ -666,7 +671,9 @@ class LsmWriteSpec:
@property
def num_buckets(self) -> Optional[int]: ...
@property
def maintained_indexes(self) -> List[str]: ...
def maintained_indexes(self) -> Optional[List[str]]:
"""Indexes the MemWAL keeps up to date, or None for every supported one."""
...
@property
def writer_config_defaults(self) -> Dict[str, str]: ...
+17 -3
View File
@@ -707,6 +707,9 @@ class LanceDBConnection(DBConnection):
self._namespace_client_properties = namespace_client_properties
if _inner is not None:
self._conn = _inner
# Native-derived wrappers resolve this in their async reconstruction
# path so construction never synchronously re-enters LOOP.
self._read_consistency_interval = read_consistency_interval
self._cached_namespace_client = None
return
@@ -756,11 +759,14 @@ class LanceDBConnection(DBConnection):
# storage_options. Also, this class really shouldn't be holding any state
# beyond _conn.
self._conn = AsyncConnection(LOOP.run(do_connect()))
# Keep property access synchronous so debugger introspection cannot wait on
# the background loop while that thread is suspended at a breakpoint.
self._read_consistency_interval = read_consistency_interval
self._cached_namespace_client: Optional[LanceNamespace] = None
@property
def read_consistency_interval(self) -> Optional[timedelta]:
return LOOP.run(self._conn.get_read_consistency_interval())
return self._read_consistency_interval
@property
def session(self) -> Optional[Session]:
@@ -771,8 +777,16 @@ class LanceDBConnection(DBConnection):
return self._conn.uri
@classmethod
def from_inner(cls, inner: LanceDbConnection):
return cls(None, _inner=inner)
def from_inner(
cls,
inner: LanceDbConnection,
read_consistency_interval: Optional[timedelta],
):
return cls(
None,
read_consistency_interval=read_consistency_interval,
_inner=inner,
)
def __repr__(self) -> str:
return f"{self.__class__.__name__}(uri={self._conn.uri!r})"
@@ -101,8 +101,7 @@ class InstructorEmbeddingFunction(TextEmbeddingFunction):
@weak_lru(maxsize=1)
def ndims(self):
model = self.get_model()
return model.encode("foo").shape[0]
return len(self.generate_embeddings([[self.source_instruction, "foo"]])[0])
def compute_query_embeddings(self, query: str, *args, **kwargs) -> List[np.array]:
return self.generate_embeddings([[self.query_instruction, query]])
+4 -3
View File
@@ -87,12 +87,13 @@ class JinaEmbeddings(EmbeddingFunction):
if isinstance(image, bytes):
image_dict = {"image": base64.b64encode(image).decode("utf-8")}
elif isinstance(image, (str, Path)):
parsed = urlparse.urlparse(image)
# TODO handle drive letter on windows.
parsed = urlparse(str(image))
PIL_Image = attempt_import_or_raise("PIL.Image", "pillow")
if parsed.scheme == "file":
pil_image = PIL_Image.open(parsed.path)
elif parsed.scheme == "":
elif parsed.scheme == "" or (os.name == "nt" and len(parsed.scheme) == 1):
# A Windows drive letter parses as a one-character scheme
# ("C:\\img.png" -> scheme="c"), so treat it as a local path.
pil_image = PIL_Image.open(image if os.name == "nt" else parsed.path)
elif parsed.scheme.startswith("http"):
pil_image = PIL_Image.open(io.BytesIO(url_retrieve(image)))
+1 -1
View File
@@ -226,7 +226,7 @@ class PermutationBuilder:
async def do_execute():
inner_tbl = await self._async.execute()
return LanceTable.from_inner(inner_tbl)
return await LanceTable.from_inner(inner_tbl)
return LOOP.run(do_execute())
+1
View File
@@ -0,0 +1 @@
+10
View File
@@ -153,6 +153,16 @@ def Vector(
return FixedSizeList
def _raise_bare_vector_error(*_args):
raise TypeError("Vector must be parameterized with a dimension, e.g. Vector(128).")
# Pydantic v1 and v2 otherwise treat the bare Vector factory as a field validator
# and inspect its signature, which produces misleading errors about internal types.
setattr(Vector, "__get_validators__", _raise_bare_vector_error)
setattr(Vector, "__get_pydantic_core_schema__", _raise_bare_vector_error)
def MultiVector(
dim: int, value_type: pa.DataType = pa.float32(), nullable: bool = True
) -> Type:
+126 -12
View File
@@ -108,6 +108,11 @@ def _should_push_down_query_table(
return namespace_client is not None and "QueryTable" in pushdown_operations
def _polars_predicate_pushdown_barrier(frame: Any) -> Any:
"""Return a Polars frame unchanged while blocking predicate pushdown."""
return frame
_MODEL_BACKED_TOKENIZER_PREFIXES = ("jieba", "lindera")
_MODEL_BACKED_TOKENIZER_ERRORS = (
"unknown base tokenizer",
@@ -864,12 +869,18 @@ class Table(ABC):
"""
raise NotImplementedError
def to_polars(self, **kwargs) -> "pl.DataFrame":
"""Return the table as a polars.DataFrame.
def to_polars(self, **kwargs) -> "pl.LazyFrame":
"""Return the table as a Polars LazyFrame.
Note
----
The Polars streaming engine is not supported because it does not currently
implement Python PyArrow dataset scans. Use the default engine when collecting
this LazyFrame.
Returns
-------
polars.DataFrame
polars.LazyFrame
"""
raise NotImplementedError
@@ -2182,11 +2193,15 @@ class LanceTable(Table):
return self.name
@classmethod
def from_inner(cls, tbl: LanceDBTable):
from .db import LanceDBConnection
async def from_inner(cls, tbl: LanceDBTable):
from .db import AsyncConnection, LanceDBConnection
async_tbl = AsyncTable(tbl)
conn = LanceDBConnection.from_inner(tbl.database())
inner_conn = tbl.database()
read_consistency_interval = await AsyncConnection(
inner_conn
).get_read_consistency_interval()
conn = LanceDBConnection.from_inner(inner_conn, read_consistency_interval)
return cls(
conn,
async_tbl.name,
@@ -2565,6 +2580,9 @@ class LanceTable(Table):
2. Currently we've disabled push-down of the filters from polars
because polars pushdown into pyarrow uses pyarrow compute
expressions rather than SQl strings (which LanceDB supports)
3. The Polars streaming engine is not supported because it does not
currently implement Python PyArrow dataset scans. Use the default
engine when collecting this LazyFrame.
Returns
-------
@@ -2573,8 +2591,12 @@ class LanceTable(Table):
from lancedb.integrations.pyarrow import PyarrowDatasetAdapter
dataset = PyarrowDatasetAdapter(self)
return pl.scan_pyarrow_dataset(
dataset, allow_pyarrow_filter=False, batch_size=batch_size
# Polars 1.32's non-PyArrow callback path passes batch_size twice. Keep
# the compatible PyArrow path, but block predicates because this adapter
# cannot translate PyArrow expressions into LanceDB filters.
return pl.scan_pyarrow_dataset(dataset, batch_size=batch_size).map_batches(
_polars_predicate_pushdown_barrier,
predicate_pushdown=False,
)
# New unified API overload
@@ -3954,6 +3976,28 @@ class LanceTable(Table):
[`AsyncTable.get_lsm_write_spec`][lancedb.AsyncTable.get_lsm_write_spec]."""
return LOOP.run(self._table.get_lsm_write_spec())
def checkpoint_lsm(self) -> None:
"""Synchronous version of
[`AsyncTable.checkpoint_lsm`][lancedb.AsyncTable.checkpoint_lsm]."""
return LOOP.run(self._table.checkpoint_lsm())
def flush_lsm(self) -> None:
"""Synchronous version of
[`AsyncTable.flush_lsm`][lancedb.AsyncTable.flush_lsm]."""
return LOOP.run(self._table.flush_lsm())
def compact_lsm(self) -> None:
"""Synchronous version of
[`AsyncTable.compact_lsm`][lancedb.AsyncTable.compact_lsm]."""
return LOOP.run(self._table.compact_lsm())
def get_lsm_stats(self, *, include_generation_rows: bool = False) -> Optional[dict]:
"""Synchronous version of
[`AsyncTable.get_lsm_stats`][lancedb.AsyncTable.get_lsm_stats]."""
return LOOP.run(
self._table.get_lsm_stats(include_generation_rows=include_generation_rows)
)
def close_lsm_writers(self) -> None:
"""Close cached MemWAL shard writers. See
[`AsyncTable.close_lsm_writers`][lancedb.AsyncTable.close_lsm_writers]."""
@@ -4632,6 +4676,13 @@ class AsyncTable:
via [`set_unenforced_primary_key`]; bucket sharding additionally
requires it to be the single column being bucketed.
By default the MemWAL maintains every index on the table, resolved
here — a snapshot, so an index created afterwards needs the spec unset
and set again. This fails if one cannot be maintained; name the set
with ``with_maintained_indexes`` to install anyway. That pins an exact
set (a still-building index is rejected, not omitted); ``[]`` maintains
none.
Parameters
----------
spec : LsmWriteSpec
@@ -4658,12 +4709,73 @@ class AsyncTable:
Returns ``None`` when the MemWAL LSM write path is not enabled (no
spec has been set, or it was removed with `unset_lsm_write_spec`).
The returned spec — including its ``maintained_indexes`` and
``writer_config_defaults`` — mirrors what was passed to
`set_lsm_write_spec`.
The returned spec mirrors what was passed to `set_lsm_write_spec`,
except that ``maintained_indexes`` always reports the concrete list
resolved when the spec was set — ``None`` never round-trips.
"""
return await self._inner.get_lsm_write_spec()
async def checkpoint_lsm(self) -> None:
"""Converge this table's LSM write path into its base table.
One flush, sealing every memtable into L0, then compaction triggers
until every generation that existed at that moment has reached base.
The loop runs client-side, reading progress from ``get_lsm_stats``.
Best-effort: generations created *while* it runs are deliberately not
waited on, which is what lets it terminate on a table taking writes.
Idempotent and safe on a cadence.
There is no deadline, and the caller owns that. It returns when the
target generations are gone, raises on a terminal server fault, and
otherwise waits however long the server takes. A slow table and a
stuck one are the same picture from the client: the compactor pool is
shared across every table on the node, so a checkpoint queued behind
unrelated work looks exactly like one that is merging. Wrap this in
``asyncio.wait_for`` for a wall-clock bound; abandoning it partway
costs nothing.
"""
return await self._inner.checkpoint_lsm()
async def flush_lsm(self) -> None:
"""Seal every bucket's active memtable into L0.
Does not touch the base table — moving L0 into base is
`compact_lsm`. On a node that has not claimed this table, this claims
it and replays its WAL log first.
"""
return await self._inner.flush_lsm()
async def compact_lsm(self) -> None:
"""Trigger a background L0 to base compaction pass per bucket.
Returns once the passes are dispatched, not once they finish: watch
``get_lsm_stats`` for progress, or use ``checkpoint_lsm`` to loop
until the current L0 has reached base.
"""
return await self._inner.compact_lsm()
async def get_lsm_stats(
self, *, include_generation_rows: bool = False
) -> Optional[dict]:
"""Read live per-bucket LSM state.
Answers "how far behind is my fresh tier", "which bucket is hot", and
"why is my fresh-tier vector search brute-force". Mutates no table
state, though on a node that has not claimed this table it claims it,
exactly as a read would.
Returns ``None`` only when the LSM write path is not enabled.
Parameters
----------
include_generation_rows
Report a row count per L0 generation. Off by default: each count
opens an uncached Lance dataset, and ``checkpoint_lsm`` polls this
needing only generation numbers.
"""
return await self._inner.get_lsm_stats(include_generation_rows)
async def close_lsm_writers(self) -> None:
"""Drain and close any cached MemWAL shard writers for this table.
@@ -6229,7 +6341,9 @@ class TableStatistics:
Attributes
----------
total_bytes: int
The total number of bytes in the table.
The total size, in bytes, of the table's data files, index files, and
overlay files. Read from the manifest, so this excludes deletion files
and manifests.
num_rows: int
The total number of rows in the table.
num_indices: int
+5
View File
@@ -395,6 +395,11 @@ def _(value: dict):
)
@value_to_sql.register(pa.Scalar)
def _(value: pa.Scalar):
return value_to_sql(value.as_py())
@value_to_sql.register(np.ndarray)
def _(value: np.ndarray):
return value_to_sql(value.tolist())
+31 -2
View File
@@ -2,9 +2,11 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import inspect
import re
import sys
from datetime import timedelta
from importlib import resources
import os
from types import SimpleNamespace
@@ -17,6 +19,10 @@ from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
from lancedb.pydantic import LanceModel, Vector
def test_package_includes_pep_561_marker():
assert resources.files(lancedb).joinpath("py.typed").is_file()
def test_basic(tmp_path):
db = lancedb.connect(tmp_path)
@@ -62,21 +68,44 @@ def test_basic(tmp_path):
assert db.open_table("test").name == db["test"].name
def test_sync_repr_does_not_use_background_loop(tmp_path, monkeypatch):
def test_sync_debugger_inspection_does_not_use_background_loop(tmp_path, monkeypatch):
from lancedb.background_loop import LOOP
db = lancedb.connect(tmp_path)
table = db.create_table("test", data=[{"id": 1}])
def fail_run(*args, **kwargs):
raise AssertionError("repr should not use the Python background loop")
raise AssertionError("debugger inspection should not use the background loop")
monkeypatch.setattr(LOOP, "run", fail_run)
# Debuggers enumerate and evaluate every exposed attribute when expanding a
# variable. This must remain safe while their breakpoint suspends LOOP's thread.
members = dict(inspect.getmembers(db))
assert members["uri"] == str(tmp_path)
assert members["read_consistency_interval"] is None
assert repr(db) == f"LanceDBConnection(uri={str(tmp_path)!r})"
assert repr(table) == f"LanceTable(name='test', _conn={db!r})"
def test_read_consistency_interval_does_not_use_background_loop(tmp_path, monkeypatch):
from lancedb.background_loop import LOOP
from lancedb.db import LanceDBConnection
consistency_interval = timedelta(seconds=5)
db = lancedb.connect(tmp_path, read_consistency_interval=consistency_interval)
db_from_inner = LanceDBConnection.from_inner(db._inner, consistency_interval)
def fail_run(*args, **kwargs):
raise AssertionError("properties should not use the Python background loop")
monkeypatch.setattr(LOOP, "run", fail_run)
assert db.read_consistency_interval == consistency_interval
assert db_from_inner.read_consistency_interval == consistency_interval
def test_ingest_pd(tmp_path):
db = lancedb.connect(tmp_path)
+51 -27
View File
@@ -64,6 +64,23 @@ def test_embedding_function(tmp_path):
assert np.allclose(actual, expected)
def test_instructor_ndims_uses_instruction():
instructor = get_registry().get("instructor").create()
model = MagicMock()
model.encode.return_value = np.zeros((1, 384))
with patch.object(type(instructor), "get_model", return_value=model):
assert instructor.ndims() == 384
model.encode.assert_called_once_with(
[[instructor.source_instruction, "foo"]],
batch_size=instructor.batch_size,
show_progress_bar=instructor.show_progress_bar,
normalize_embeddings=instructor.normalize_embeddings,
device=instructor.device,
)
def test_embedding_function_variables():
@register("variable-testing")
class VariableTestingFunction(TextEmbeddingFunction):
@@ -115,34 +132,16 @@ def test_embedding_function_variables():
assert func.safe_model_dump()["secret_key"] == "$var:secret"
def test_parse_functions_with_variables():
@register("variable-parsing-test")
class VariableParsingFunction(TextEmbeddingFunction):
api_key: str
base_url: Optional[str] = None
@staticmethod
def sensitive_keys():
return ["api_key"]
def ndims(self):
return 10
def generate_embeddings(self, texts):
# Mock implementation that just returns random embeddings
# In real usage, this would use the api_key to call an API
return [np.random.rand(self.ndims()).tolist() for _ in texts]
def test_openai_variables_survive_metadata_round_trip():
registry = EmbeddingFunctionRegistry.get_instance()
registry.set_var("test_api_key", "sk-test-key-12345")
registry.set_var("test_base_url", "https://api.example.com")
conf = EmbeddingFunctionConfig(
source_column="text",
vector_column="vector",
function=registry.get("variable-parsing-test").create(
api_key="$var:test_api_key", base_url="$var:test_base_url"
function=registry.get("openai").create(
api_key="$var:test_api_key", base_url="https://api.example.com"
),
)
@@ -150,7 +149,10 @@ def test_parse_functions_with_variables():
# Create a mock arrow table with the metadata
schema = pa.schema(
[pa.field("text", pa.string()), pa.field("vector", pa.list_(pa.float32(), 10))]
[
pa.field("text", pa.string()),
pa.field("vector", pa.list_(pa.float32(), 1536)),
]
)
table = pa.table({"text": [], "vector": []}, schema=schema)
table = table.replace_schema_metadata(metadata)
@@ -164,13 +166,15 @@ def test_parse_functions_with_variables():
assert parsed_func.api_key == "sk-test-key-12345"
assert parsed_func.base_url == "https://api.example.com"
embeddings = parsed_func.generate_embeddings(["test text"])
assert len(embeddings) == 1
assert len(embeddings[0]) == 10
assert parsed_func.safe_model_dump()["api_key"] == "$var:test_api_key"
with patch("lancedb.embeddings.openai.attempt_import_or_raise") as import_openai:
parsed_func._openai_client
import_openai.return_value.OpenAI.assert_called_once_with(
api_key="sk-test-key-12345", base_url="https://api.example.com"
)
def test_embedding_with_bad_results(tmp_path):
@register("null-embedding")
@@ -627,3 +631,23 @@ def test_url_retrieve_downloads_image():
image_bytes = url_retrieve(image_url)
img = Image.open(io.BytesIO(image_bytes))
assert img.size[0] > 0 and img.size[1] > 0
def test_jina_generate_image_input_dict_local_path(tmp_path):
"""
JinaEmbeddings._generate_image_input_dict must accept a local image path
(str or Path), not just bytes. Previously it crashed with
`AttributeError: 'function' object has no attribute 'urlparse'` on any
str/Path input because it called `urlparse.urlparse(image)` instead of
`urlparse(image)` (urlparse was imported as a function, not a module).
"""
Image = pytest.importorskip("PIL.Image")
from lancedb.embeddings.jinaai import JinaEmbeddings
image_path = tmp_path / "test.png"
Image.new("RGB", (4, 4), color="red").save(image_path, format="PNG")
for image in (str(image_path), image_path):
image_dict = JinaEmbeddings._generate_image_input_dict(image)
assert "image" in image_dict
assert isinstance(image_dict["image"], str) and len(image_dict["image"]) > 0
+81 -1
View File
@@ -12,7 +12,7 @@ import pyarrow.compute as pc
import pytest
import pytest_asyncio
from lancedb.index import FTS
from lancedb.index import BTree, FTS, IvfPq
from lancedb.table import AsyncTable, Table
@@ -99,6 +99,86 @@ async def test_async_hybrid_query_filters(table: AsyncTable):
assert result["text"].to_pylist() == ["cat", "b"]
@pytest.mark.asyncio
async def test_hybrid_query_with_stale_fixed_size_binary_prefilter(
tmpdir_factory,
):
tmp_path = str(tmpdir_factory.mktemp("stale_scalar_prefilter"))
db = await lancedb.connect_async(tmp_path)
def fixed_size_binary(value: int) -> bytes:
return value.to_bytes(16, byteorder="big")
num_rows = 1000
data = pa.table(
{
"space_id": pa.array(
[fixed_size_binary(i) for i in range(num_rows)],
type=pa.binary(16),
),
"text": ["book"] * num_rows,
"vector": pa.array(
[[float(i), float(i)] for i in range(num_rows)],
type=pa.list_(pa.float32(), 2),
),
}
)
table = await db.create_table("test", data)
await table.create_index(
"vector", config=IvfPq(num_partitions=4, num_sub_vectors=2)
)
await table.create_index("space_id", config=BTree())
await table.create_index("text", config=FTS(with_position=False))
# Advance the search indices without advancing the scalar index. This is the
# state that previously let hybrid search use an incomplete scalar prefilter.
await table.add(data)
lance_dataset = await table.to_lance()
lance_dataset.optimize.optimize_indices(index_names=["vector_idx", "text_idx"])
await table.checkout_latest()
scalar_stats = await table.index_stats("space_id_idx")
assert scalar_stats is not None
assert scalar_stats.num_indexed_rows == num_rows
assert scalar_stats.num_unindexed_rows == num_rows
for index_name in ["vector_idx", "text_idx"]:
search_stats = await table.index_stats(index_name)
assert search_stats is not None
assert search_stats.num_indexed_rows == num_rows * 2
assert search_stats.num_unindexed_rows == 0
matching_ids = [5, 10, 15, 20, 25, 30]
literals = [
f"arrow_cast(0x{fixed_size_binary(i).hex()}, 'FixedSizeBinary(16)')"
for i in matching_ids
]
predicate = f"space_id IN ({', '.join(literals)})"
expected_ids = sorted(fixed_size_binary(i) for i in matching_ids for _ in range(2))
vector_query = (
table.query().where(predicate).nearest_to([5.0, 5.0]).limit(num_rows * 2)
)
vector_results = await vector_query.to_arrow()
assert sorted(vector_results["space_id"].to_pylist()) == expected_ids
fts_query = (
table.query().where(predicate).nearest_to_text("book").limit(num_rows * 2)
)
fts_results = await fts_query.to_arrow()
assert sorted(fts_results["space_id"].to_pylist()) == expected_ids
hybrid_results = await (
table.query()
.where(predicate)
.nearest_to([5.0, 5.0])
.nearest_to_text("book")
.limit(num_rows * 2)
.to_arrow()
)
assert sorted(hybrid_results["space_id"].to_pylist()) == expected_ids
@pytest.mark.asyncio
async def test_async_hybrid_query_default_limit(table: AsyncTable):
# add 10 new rows
+33
View File
@@ -0,0 +1,33 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import re
import shutil
import subprocess
import sys
import lancedb._lancedb as _lancedb
import pytest
@pytest.mark.skipif(sys.platform != "linux", reason="ldd is Linux-specific")
def test_native_extension_does_not_link_openssl():
"""OpenSSL-linked wheels abort when imported on RHEL hosts in FIPS mode."""
ldd = shutil.which("ldd")
if ldd is None:
pytest.skip("ldd is not installed")
result = subprocess.run(
[ldd, _lancedb.__file__],
check=True,
capture_output=True,
text=True,
)
openssl_libraries = re.findall(
r"^\s*(lib(?:crypto|ssl)\S*)\s+=>", result.stdout, flags=re.MULTILINE
)
assert not openssl_libraries, (
"the LanceDB native extension must use rustls instead of linking OpenSSL: "
f"{openssl_libraries}"
)
+25
View File
@@ -372,6 +372,31 @@ async def test_create_vector_index(some_table: AsyncTable):
assert stats.num_indices == 1
@pytest.mark.asyncio
async def test_create_ivf_index_reports_unsplittable_partitions(db_async):
dim = 8
num_partitions = 300 # More than 256 selects hierarchical k-means.
base_vectors = [[float(row == column) for column in range(dim)] for row in range(5)]
vectors = pa.array(base_vectors * 200, pa.list_(pa.float32(), dim))
table = await db_async.create_table(
"unsplittable_partitions",
pa.table({"vector": vectors}),
)
error_pattern = (
rf"Cannot create {num_partitions} IVF partitions: k-means could only form"
)
with pytest.raises(RuntimeError, match=error_pattern):
await table.create_index(
"vector",
config=IvfFlat(
distance_type="dot",
num_partitions=num_partitions,
max_iterations=10,
),
)
@pytest.mark.asyncio
async def test_create_4bit_ivfpq_index(some_table: AsyncTable):
# Can create
+11 -4
View File
@@ -83,7 +83,9 @@ def test_lsm_write_spec_repr():
assert s.spec_type == "bucket"
assert s.column == "id"
assert s.num_buckets == 4
assert s.maintained_indexes == []
# A fresh spec defers its maintained set to install time.
assert s.maintained_indexes is None
assert s.with_maintained_indexes([]).maintained_indexes == []
assert "bucket" in repr(s)
assert "id" in repr(s)
assert "4" in repr(s)
@@ -169,18 +171,23 @@ def test_get_lsm_write_spec(tmp_path):
table.unset_lsm_write_spec()
assert table.get_lsm_write_spec() is None
# Identity round-trips (column recovered from the schema).
# 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()
assert spec.spec_type == "identity"
assert spec.column == "id"
assert spec.maintained_indexes == [idx_name]
table.unset_lsm_write_spec()
# Unsharded round-trips (no routing column).
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
# 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()
assert spec.spec_type == "unsharded"
assert spec.column is None
assert spec.maintained_indexes == []
@pytest.mark.asyncio
+2 -2
View File
@@ -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())
table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([]))
with pytest.raises(Exception, match="maintained"):
table.search("fox", query_type="fts", fts_columns="text").to_arrow()
@@ -631,7 +631,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())
table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([]))
with pytest.raises(Exception, match="maintained"):
table.search([1.0] * VECTOR_DIM).to_arrow()
@@ -0,0 +1,42 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import importlib
import re
import sys
from pathlib import Path
import pytest
def test_pyo3_abi_matches_minimum_supported_python():
project_dir = Path(__file__).parents[2]
pyproject = (project_dir / "pyproject.toml").read_text()
cargo_manifest = (project_dir / "Cargo.toml").read_text()
minimum_python = re.search(
r'^requires-python\s*=\s*">=(\d+)\.(\d+)"$', pyproject, re.MULTILINE
)
assert minimum_python is not None
major, minor = minimum_python.groups()
expected_abi = f"abi3-py{major}{minor}"
configured_abis = re.findall(r'"(abi3-py\d+)"', cargo_manifest)
assert configured_abis == [expected_abi, expected_abi], (
"the pyo3 runtime and build ABI features must both match requires-python"
)
@pytest.mark.skipif(sys.platform != "win32", reason="Windows wheel regression test")
def test_windows_wheel_tag_and_native_import():
project_dir = Path(__file__).parents[2]
wheels = list((project_dir.parent / "target" / "wheels").glob("lancedb-*.whl"))
if not wheels:
pytest.skip("no wheel artifact is available in this development environment")
assert len(wheels) == 1
assert wheels[0].name.endswith("-cp310-abi3-win_amd64.whl")
native_module = importlib.import_module("lancedb._lancedb")
assert Path(native_module.__file__).suffix == ".pyd"
+20
View File
@@ -6,6 +6,7 @@ import math
import pytest
from lancedb import DBConnection, Table, connect
from lancedb.background_loop import LOOP
from lancedb.permutation import Permutation, Permutations, permutation_builder
@@ -31,6 +32,25 @@ def test_split_random_ratios(mem_db):
assert 65 <= split_1_count <= 75 # ~70% ± tolerance
def test_execute_does_not_reenter_background_loop(tmp_path, monkeypatch):
import threading
db = connect(tmp_path)
tbl = db.create_table("test_table", pa.table({"x": range(10)}))
original_run = LOOP.run
def fail_on_reentry(future):
assert threading.current_thread() is not LOOP.thread
return original_run(future)
monkeypatch.setattr(LOOP, "run", fail_on_reentry)
permutation_tbl = permutation_builder(tbl).execute()
assert permutation_tbl.count_rows() == 10
assert permutation_tbl._conn.read_consistency_interval is None
def test_split_random_counts(mem_db):
"""Test random splitting with absolute counts."""
tbl = mem_db.create_table(
+11
View File
@@ -415,6 +415,17 @@ def test_nullable_vector():
assert schema == pa.schema([pa.field("vec", pa.list_(pa.float32(), 16), True)])
def test_bare_vector_raises_clear_error():
namespace = {
"__name__": "test_model_without_pyarrow",
"LanceModel": LanceModel,
"Vector": Vector,
}
with pytest.raises(TypeError, match=r"Vector must be parameterized.*Vector\(128\)"):
exec("class TestModel(LanceModel):\n vector: Vector", namespace)
def test_fixed_size_list_field():
class TestModel(pydantic.BaseModel):
vec: Vector(16)
+9
View File
@@ -570,6 +570,15 @@ def test_query_builder(table):
assert all(np.array(rs[0]["vector"]) == [1, 2])
def test_query_multiple_vectors(table):
results = table.search([np.array([1, 2]), np.array([4, 5])]).limit(1).to_list()
assert len(results) == 2
results_by_query = {result["query_index"]: result for result in results}
assert results_by_query[0]["id"] == 1
assert results_by_query[1]["id"] == 2
def test_with_row_id(table: lancedb.table.Table):
rs = table.search().with_row_id(True).to_arrow()
assert "_rowid" in rs.column_names
+6
View File
@@ -35,6 +35,12 @@ def make_mock_http_handler(handler):
return MockLanceDBHandler
@pytest.mark.parametrize("db_name", ["a" * 64, "invalid..database"])
def test_connect_rejects_invalid_cloud_dns_hostname(db_name):
with pytest.raises(ValueError, match="DNS labels must contain 1 to 63 bytes"):
lancedb.connect(f"db://{db_name}", api_key="fake")
@contextlib.contextmanager
def mock_lancedb_connection(handler):
with http.server.HTTPServer(
+302 -4
View File
@@ -2,10 +2,14 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import ctypes
import gc
import os
import sys
import threading
import warnings
import weakref
from concurrent.futures import ThreadPoolExecutor
from datetime import date, datetime, timedelta
from time import sleep
from typing import List
@@ -98,6 +102,30 @@ def test_basic(mem_db: DBConnection):
assert table.to_arrow() == expected_data
def test_search_preserves_nulls_from_sliced_arrow_table(mem_db: DBConnection):
data = pa.table(
{
"id": [0, 1, 2, 3, 4],
"score_cn": [None, 22, None, 5, 8],
"score_mt": [None, 42, None, 5, 8],
"vector": [
[20, 19, -1, -1],
[41, 38, 22, 42],
[10, 10, -1, -1],
[5, 5, 5, 5],
[8, 8, 8, 8],
],
}
).slice(1)
table = mem_db.create_table("sliced_nullable", data=data)
result = table.search([41, 38, 22, 42]).limit(1).to_arrow()
assert result["id"].to_pylist() == [1]
assert result["score_cn"].to_pylist() == [22]
assert result["score_mt"].to_pylist() == [42]
def test_table_to_pandas_default_matches_arrow(tmp_db: DBConnection):
pd = pytest.importorskip("pandas")
data = pa.table({"id": [1, 2], "text": ["one", "two"]})
@@ -434,6 +462,38 @@ def test_add(mem_db: DBConnection):
_add(table, schema)
def test_add_releases_arrow_buffers_without_gc(mem_db: DBConnection):
"""Regression test for https://github.com/lancedb/lancedb/issues/2512."""
schema = pa.schema([pa.field("x", pa.int64())])
table = mem_db.create_table("test_add_releases_arrow_buffers", schema=schema)
class BufferOwner:
def __init__(self, size: int):
self.memory = ctypes.create_string_buffer(size)
owner_refs = []
gc_was_enabled = gc.isenabled()
gc.disable()
try:
for _ in range(3):
size = 8 * 1024
owner = BufferOwner(size)
arrow_buffer = pa.foreign_buffer(
ctypes.addressof(owner.memory), size, owner
)
array = pa.Array.from_buffers(pa.int64(), 1024, [None, arrow_buffer])
batch = pa.RecordBatch.from_arrays([array], schema=schema)
owner_refs.append(weakref.ref(owner))
table.add(batch)
del batch, array, arrow_buffer, owner
assert all(owner_ref() is None for owner_ref in owner_refs)
finally:
if gc_was_enabled:
gc.enable()
def test_add_write_parallelism(mem_db: DBConnection):
schema = pa.schema([pa.field("id", pa.int64())])
table = mem_db.create_table("test", schema=schema)
@@ -869,6 +929,7 @@ def test_polars(mem_db: DBConnection):
# enter table to polars dataframe
result = table.to_polars()
assert isinstance(result, pl.LazyFrame)
assert np.allclose(result.collect()["vector"].to_list(), data["vector"])
# make sure filtering isn't broken
@@ -1785,6 +1846,27 @@ def test_add_with_empty_fixed_size_list_drops_bad_rows(mem_db: DBConnection):
assert np.allclose(data["embedding"].to_pylist()[0], np.array([0.1] * 16))
def test_add_nullable_fixed_size_list_with_none(mem_db: DBConnection):
"""Regression test for issue #2340."""
table = mem_db.create_table(
"test_nullable_fixed_size_list",
schema=pa.schema(
[
pa.field("id", pa.string()),
pa.field("feature", pa.list_(pa.float32(), 256)),
pa.field("tags", pa.list_(pa.string())),
]
),
)
table.add([{"id": "1", "feature": None, "tags": ["tag1", "tag2"]}])
result = table.to_arrow()
assert result.to_pylist() == [
{"id": "1", "feature": None, "tags": ["tag1", "tag2"]}
]
def test_add_nullable_struct_with_none(mem_db: DBConnection):
"""Regression test for issue #2654: a nullable struct column whose
first batch contains only None values must not crash in
@@ -1824,6 +1906,33 @@ def test_add_nullable_struct_with_none(mem_db: DBConnection):
assert result.column("data").to_pylist() == [{"x": 1.0}, None]
def test_read_mostly_null_list_v2_2_page_boundary(tmp_path):
# Regression test for #3194. This row/value count crosses a v2.2 structural
# encoding page boundary where Lance 3.0.0 sliced repetition/definition
# levels by row offset and decoded child arrays at different lengths.
num_rows = 64_885
num_values = 217
list_type = pa.list_(pa.float32())
source = pa.table(
{
"id": np.arange(num_rows, dtype=np.int64),
"coords": pa.array(
[[1.0, 2.0, 3.0, 4.0]] * num_values + [None] * (num_rows - num_values),
type=list_type,
),
}
)
db = lancedb.connect(
tmp_path,
storage_options={"new_table_data_storage_version": "2.2"},
)
table = db.create_table("test_sparse_nullable_list", data=source)
result = table.search().select(["id", "coords"]).limit(num_rows).to_arrow()
assert result.equals(source)
def test_add_with_integer_embeddings_preserves_casting(mem_db: DBConnection):
class Schema(LanceModel):
text: str
@@ -2109,6 +2218,45 @@ def test_merge(tmp_db: DBConnection, tmp_path):
table.merge(other_dataset, left_on="id")
@pytest.mark.parametrize("storage_version", ["legacy", "stable"])
def test_search_after_merge(tmp_path, storage_version):
pytest.importorskip("lance")
pd = pytest.importorskip("pandas")
db = lancedb.connect(
tmp_path,
storage_options={"new_table_data_storage_version": storage_version},
)
rng = np.random.default_rng(42)
row_count = 512
vectors = rng.standard_normal((row_count, 8)).astype(np.float32)
table = db.create_table(
"search_after_merge",
data=pd.DataFrame(
{
"id": [str(i) for i in range(row_count)],
"vector": list(vectors),
}
),
)
table.create_index("vector", config=IvfPq(num_partitions=1, num_sub_vectors=2))
links = pd.DataFrame(
{
"id": [str(i) for i in range(row_count // 2)],
"link": [f"https://example.com/{i}" for i in range(row_count // 2)],
}
)
table.merge(links, left_on="id")
query = table.search(vectors[-1]).refine_factor(50).limit(10)
assert "ANN" in query.explain_plan(verbose=True)
result = query.to_arrow()
links_by_id = dict(zip(result["id"].to_pylist(), result["link"].to_pylist()))
assert links_by_id[str(row_count - 1)] is None
def test_delete(mem_db: DBConnection):
table = mem_db.create_table(
"my_table",
@@ -2124,6 +2272,27 @@ def test_delete(mem_db: DBConnection):
assert table.to_arrow()["id"].to_pylist() == [1]
def test_concurrent_deletes_are_thread_safe(mem_db: DBConnection):
num_workers = 8
table = mem_db.create_table(
"my_table", data=[{"id": row_id} for row_id in range(num_workers)]
)
barrier = threading.Barrier(num_workers)
def delete(row_id: int):
barrier.wait()
return table.delete(f"id = {row_id}")
with ThreadPoolExecutor(max_workers=num_workers) as pool:
results = list(pool.map(delete, range(num_workers)))
assert all(result.num_deleted_rows == 1 for result in results)
assert sorted(result.version for result in results) == list(
range(2, num_workers + 2)
)
assert table.count_rows() == 0
def test_delete_expr(mem_db: DBConnection):
table = mem_db.create_table(
"my_table",
@@ -2174,6 +2343,20 @@ def test_update(mem_db: DBConnection):
assert np.allclose(v, np.array([[1.2, 1.9], [1.1, 1.1]]))
def test_update_with_arrow_scalar(mem_db: DBConnection):
schema = pa.schema({"id": pa.int64(), "vector": pa.list_(pa.float32(), 4)})
table = mem_db.create_table("my_table", schema=schema)
table.add([{"id": 1, "vector": [1.0, 2.0, 3.0, 4.0]}])
value = table.search().select(["vector"]).limit(1).to_arrow()["vector"][0]
assert isinstance(value, pa.FixedSizeListScalar)
result = table.update(where="id == 1", values={"vector": value})
assert result.rows_updated == 1
assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0, 3.0, 4.0]]
def test_update_types(mem_db: DBConnection):
table = mem_db.create_table(
"my_table",
@@ -2341,6 +2524,55 @@ def test_merge_insert(mem_db: DBConnection):
)
def test_merge_insert_nullable_pandas_into_pydantic_schema(mem_db: DBConnection):
# Regression test for https://github.com/lancedb/lancedb/issues/2366
pd = pytest.importorskip("pandas")
class Document(LanceModel):
id: int
title: str
content: str
table = mem_db.create_table("documents", schema=Document)
table.add(
pd.DataFrame(
{
"title": ["Old title", "Unchanged"],
"id": [2, 3],
"content": ["Old content", "Keep this"],
}
)
)
# Pandas produces nullable Arrow fields, in an order that differs from the
# non-nullable Pydantic schema. This is valid as long as the data has no nulls.
new_data = pd.DataFrame(
{
"title": ["Inserted", "Updated"],
"id": [1, 2],
"content": ["New row", "New content"],
}
)
result = (
table.merge_insert("id")
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(new_data)
)
assert result.num_inserted_rows == 1
assert result.num_updated_rows == 1
expected = pa.Table.from_pylist(
[
{"id": 1, "title": "Inserted", "content": "New row"},
{"id": 2, "title": "Updated", "content": "New content"},
{"id": 3, "title": "Unchanged", "content": "Keep this"},
],
schema=Document.to_arrow_schema(),
)
assert table.to_arrow().sort_by("id") == expected
def test_merge_insert_by_source_delete_expr(mem_db: DBConnection):
table = mem_db.create_table(
"my_table",
@@ -2441,6 +2673,36 @@ def test_merge_insert_subschema(mem_db: DBConnection, data_format):
assert table.to_arrow().sort_by("id") == expected
def test_repeated_partial_merge_insert_with_scalar_index(mem_db: DBConnection):
def make_batch(start: int) -> pa.Table:
return pa.table(
{
"id": [f"id-{i:04}" for i in range(start, start + 100)],
"category": ["A"] * 100,
"value_a": [float(i) for i in range(start, start + 100)],
"value_b": [float(i) / 10 for i in range(100)],
}
)
table = mem_db.create_table("my_table", data=make_batch(0))
table.add(make_batch(100))
table.add(make_batch(200))
table.create_index("id", config=BTree())
ids = [f"id-{i:04}" for i in range(100, 200)]
for value in (999.0, 888.0):
result = (
table.merge_insert("id")
.when_matched_update_all()
.execute(pa.table({"id": ids, "value_a": [value] * 100}))
)
assert result.num_updated_rows == 100
actual = table.to_arrow().sort_by("id")
assert actual.num_rows == 300
assert actual["value_a"].to_pylist()[100:200] == [888.0] * 100
@pytest.mark.asyncio
async def test_merge_insert_async(mem_db_async: AsyncConnection):
data = pa.table({"a": [1, 2, 3], "b": ["a", "b", "c"]})
@@ -2537,15 +2799,40 @@ def test_create_with_embedding_function(mem_db: DBConnection):
assert actual == expected
def test_create_f16_table_from_arrow_data(mem_db: DBConnection):
dimension = 32
num_rows = 512
values = pa.array(
np.random.default_rng(42)
.standard_normal(num_rows * dimension)
.astype(np.float16)
)
df = pa.table(
{
"text": [f"s-{i}" for i in range(num_rows)],
"vector": pa.FixedSizeListArray.from_arrays(values, dimension),
}
)
table = mem_db.create_table("f16_tbl", data=df)
assert table.schema.field("vector").type == pa.list_(pa.float16(), dimension)
table.create_index(num_partitions=2, num_sub_vectors=2)
query = df["vector"][2].as_py()
expected = table.search(query).limit(2).to_arrow()
assert "s-2" in expected["text"].to_pylist()
def test_create_f16_table(mem_db: DBConnection):
class MyTable(LanceModel):
text: str
vector: Vector(32, value_type=pa.float16())
rng = np.random.default_rng(42)
df = pa.table(
{
"text": [f"s-{i}" for i in range(512)],
"vector": [np.random.randn(32).astype(np.float16) for _ in range(512)],
"vector": [rng.standard_normal(32).astype(np.float16) for _ in range(512)],
}
)
table = mem_db.create_table(
@@ -3426,7 +3713,8 @@ def test_stats(mem_db: DBConnection):
stats = table.stats()
print(f"{stats=}")
assert stats == {
"total_bytes": 60,
# Full on-disk size of the data file, footer and metadata included.
"total_bytes": 633,
"num_rows": 2,
"num_indices": 0,
"fragment_stats": {
@@ -3444,6 +3732,13 @@ def test_stats(mem_db: DBConnection):
},
}
# 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"]
def test_create_table_empty_list_with_schema(mem_db: DBConnection):
"""Test creating table with empty list data and schema
@@ -3467,8 +3762,8 @@ def test_create_table_empty_list_no_schema_error(mem_db: DBConnection):
mem_db.create_table("test_empty_no_schema", data=[])
def test_add_table_with_empty_embeddings(tmp_path):
"""Test exact scenario from issue #1968
def test_create_table_without_data_with_vector_schema(tmp_path):
"""Test exact scenario from issue #1968.
Regression test for issue #1968:
https://github.com/lancedb/lancedb/issues/1968
@@ -3480,6 +3775,9 @@ def test_add_table_with_empty_embeddings(tmp_path):
embedding: Vector(16)
table = db.create_table("test", schema=MySchema)
assert table.count_rows() == 0
assert table.schema == MySchema.to_arrow_schema()
table.add(
[{"text": "bar", "embedding": [0.1] * 16}],
on_bad_vectors="drop",
@@ -75,6 +75,22 @@ class TestVoyageAIModelRegistration:
with pytest.raises(ValueError, match="not supported"):
func.ndims()
def test_voyage3_source_embeddings_use_text_api(self, mock_voyageai_client):
"""Regression test for text table data being sent to the multimodal API."""
mock_voyageai_client.tokenize.return_value = [["hello", "world"]]
mock_voyageai_client.embed.return_value.embeddings = [[0.1] * 1024]
registry = get_registry()
func = registry.get("voyageai").create(name="voyage-3")
embeddings = func.compute_source_embeddings("hello world")
assert embeddings == [[0.1] * 1024]
mock_voyageai_client.embed.assert_called_once_with(
texts=["hello world"], model="voyage-3", input_type="document"
)
mock_voyageai_client.multimodal_embed.assert_not_called()
@pytest.mark.parametrize(
"model_name",
[
+15
View File
@@ -0,0 +1,15 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
from typing import assert_type
import lancedb
from lancedb import AsyncConnection, DBConnection
def check_connect_type() -> None:
assert_type(lancedb.connect("memory://"), DBConnection)
async def check_connect_async_type() -> None:
assert_type(await lancedb.connect_async("memory://"), AsyncConnection)
+141 -16
View File
@@ -28,11 +28,72 @@ use pyo3::{
Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python,
exceptions::{PyRuntimeError, PyValueError},
pyclass, pyfunction, pymethods,
types::{IntoPyDict, PyAnyMethods, PyBytes, PyDict, PyDictMethods},
types::{IntoPyDict, PyAnyMethods, PyBytes, PyDict, PyDictMethods, PyList, PyListMethods},
};
mod scannable;
/// Convert `LsmStats` to a Python dict, preserving the per-bucket 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
/// it to find.
fn lsm_stats_to_py(py: Python<'_>, stats: &lancedb::table::LsmStats) -> PyResult<Py<PyDict>> {
let out = PyDict::new(py);
let buckets = PyList::empty(py);
for b in &stats.buckets {
let e = PyDict::new(py);
e.set_item("shard_id", &b.shard_id)?;
e.set_item("status", &b.status)?;
e.set_item("writer_epoch", b.writer_epoch)?;
e.set_item("manifest_version", b.manifest_version)?;
e.set_item("current_generation", b.current_generation)?;
e.set_item(
"replay_after_wal_entry_position",
b.replay_after_wal_entry_position,
)?;
e.set_item(
"wal_entry_position_last_seen",
b.wal_entry_position_last_seen,
)?;
let generations = PyList::empty(py);
for g in &b.generations {
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)?;
}
e.set_item("generations", generations)?;
e.set_item("compacting", b.compacting)?;
e.set_item(
"memtables",
b.memtables
.as_ref()
.map(|ms| {
let l = PyList::empty(py);
for m in ms {
let d = PyDict::new(py);
d.set_item("generation", m.generation)?;
d.set_item("rows", m.rows)?;
d.set_item("bytes", m.bytes)?;
d.set_item("batches", m.batches)?;
d.set_item("indexes", m.indexes.clone())?;
l.append(d)?;
}
PyResult::Ok(l.unbind())
})
.transpose()?,
)?;
buckets.append(e)?;
}
out.set_item("buckets", buckets)?;
Ok(out.unbind())
}
#[derive(FromPyObject)]
enum PredicateArg {
Expr(PyExpr),
@@ -185,12 +246,22 @@ impl From<lancedb::table::MergeResult> for MergeResult {
}
}
/// Render for `__repr__`, so the default reads as Python's `None` rather than
/// Rust's `Some([..])`.
fn fmt_maintained(maintained: &Option<Vec<String>>) -> String {
match maintained {
Some(names) => format!("{:?}", names),
None => "None".to_string(),
}
}
/// Specification selecting Lance's MemWAL LSM-style write path for
/// `merge_insert`.
///
/// Constructed via the `bucket(...)`, `identity(...)`, or `unsharded()`
/// classmethods, then optionally chain `with_maintained_indexes(...)` and
/// `with_writer_config_defaults(...)`.
/// `with_writer_config_defaults(...)`. A fresh spec maintains every index the
/// MemWAL supports, resolved on install.
#[pyclass(from_py_object)]
#[derive(Clone, Debug)]
pub struct LsmWriteSpec {
@@ -230,11 +301,11 @@ impl LsmWriteSpec {
}
}
/// Replace the list of indexes the MemWAL should keep up to date as
/// rows are appended. Each name must reference an index that
/// already exists on the table at the time `set_lsm_write_spec`
/// is called.
pub fn with_maintained_indexes(&self, indexes: Vec<String>) -> Self {
/// 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.
#[pyo3(signature = (indexes))]
pub fn with_maintained_indexes(&self, indexes: Option<Vec<String>>) -> Self {
Self {
inner: self.inner.clone().with_maintained_indexes(indexes),
}
@@ -256,23 +327,29 @@ impl LsmWriteSpec {
maintained_indexes,
writer_config_defaults,
} => format!(
"LsmWriteSpec.bucket(column={:?}, num_buckets={}, maintained_indexes={:?}, writer_config_defaults={:?})",
column, num_buckets, maintained_indexes, writer_config_defaults,
"LsmWriteSpec.bucket(column={:?}, num_buckets={}, maintained_indexes={}, writer_config_defaults={:?})",
column,
num_buckets,
fmt_maintained(maintained_indexes),
writer_config_defaults,
),
lancedb::table::LsmWriteSpec::Identity {
column,
maintained_indexes,
writer_config_defaults,
} => format!(
"LsmWriteSpec.identity(column={:?}, maintained_indexes={:?}, writer_config_defaults={:?})",
column, maintained_indexes, writer_config_defaults,
"LsmWriteSpec.identity(column={:?}, maintained_indexes={}, writer_config_defaults={:?})",
column,
fmt_maintained(maintained_indexes),
writer_config_defaults,
),
lancedb::table::LsmWriteSpec::Unsharded {
maintained_indexes,
writer_config_defaults,
} => format!(
"LsmWriteSpec.unsharded(maintained_indexes={:?}, writer_config_defaults={:?})",
maintained_indexes, writer_config_defaults,
"LsmWriteSpec.unsharded(maintained_indexes={}, writer_config_defaults={:?})",
fmt_maintained(maintained_indexes),
writer_config_defaults,
),
}
}
@@ -307,10 +384,10 @@ impl LsmWriteSpec {
}
}
/// Names of indexes the MemWAL should keep up to date during writes.
/// Indexes the MemWAL keeps up to date, or `None` for every supported one.
#[getter]
pub fn maintained_indexes(&self) -> Vec<String> {
self.inner.maintained_indexes().to_vec()
pub fn maintained_indexes(&self) -> Option<Vec<String>> {
self.inner.maintained_indexes().map(<[String]>::to_vec)
}
/// Default `ShardWriter` configuration recorded by this spec.
@@ -745,6 +822,9 @@ impl Table {
#[allow(private_interfaces)]
pub fn delete(self_: PyRef<'_, Self>, condition: PredicateArg) -> PyResult<Bound<'_, PyAny>> {
// Do not hold the Python borrow across the await. The cloned Rust table
// handle is thread-safe and allows deletes on the same Python table to
// run concurrently without PyO3 reporting "Already borrowed".
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let result = match &condition {
@@ -1336,6 +1416,51 @@ impl Table {
})
}
/// Converge the table's LSM 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>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
inner.checkpoint_lsm().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>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
inner.compact_lsm().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(
self_: PyRef<'_, Self>,
include_generation_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)
.await
.infer_error()?;
Python::attach(|py| stats.map(|s| lsm_stats_to_py(py, &s)).transpose())
})
}
pub fn close_lsm_writers(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
+1 -1
View File
@@ -1998,7 +1998,7 @@ requires-dist = [
{ name = "pillow", marker = "extra == 'clip'", specifier = ">=12.1.1" },
{ name = "pillow", marker = "extra == 'embeddings'", specifier = ">=12.1.1" },
{ name = "pillow", marker = "extra == 'siglip'", specifier = ">=12.1.1" },
{ name = "polars", marker = "extra == 'tests'", specifier = ">=0.19,<=1.3.0" },
{ name = "polars", marker = "extra == 'tests'", specifier = ">=0.19,<=1.32.3" },
{ name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.5.0" },
{ name = "pyarrow", specifier = ">=16" },
{ name = "pyarrow", marker = "extra == 'tests'", specifier = "<25" },