mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-18 12:08:35 +00:00
Merge remote-tracking branch 'origin/main' into gatekeeper/fix-2923-1
This commit is contained in:
+3
-3
@@ -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]
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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]: ...
|
||||
|
||||
|
||||
@@ -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]])
|
||||
|
||||
@@ -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)))
|
||||
|
||||
@@ -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())
|
||||
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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",
|
||||
[
|
||||
|
||||
@@ -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
@@ -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 {
|
||||
|
||||
Generated
+1
-1
@@ -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" },
|
||||
|
||||
Reference in New Issue
Block a user