mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-18 20:18:37 +00:00
feat(query): add use_lsm to read MemWAL LSM data (#3489)
## What
MemWAL LSM **read** support. When a table has an LSM write spec
(`set_lsm_write_spec`), `merge_insert` upserts live in the MemWAL
active/frozen memtables and flushed SSTables until an external
compaction merges them into the base table, so a normal scan returns
**stale** data. This routes reads through Lance's `LsmScanner` so
queries also surface that in-flight data, deduplicated by primary key
(newest generation wins).
## How
- Adds a **`use_lsm: Option<bool>`** query flag, symmetric with the
`merge_insert` flag:
- **unset** — auto-route through the LSM scanner when the table carries
a write spec
- **`use_lsm(true)`** — force the LSM path; error if there is no spec
- **`use_lsm(false)`** — read the base table only (the escape hatch)
- Plain scan, single-column full-text search, and single-vector ANN all
run through one `LsmScanner` (assembled from on-disk shard manifests
plus the cached writer's in-memory memtables), so a `where` predicate is
honored as a **prefilter** uniformly — including for vector search.
- **Compaction-aware snapshots:** an SSTable generation is dropped only
once it is both compacted into the base table and covered by the arm's
base-index catch-up (`index_catchup`); plain scans use the compaction
watermark alone.
- Query shapes the scanner cannot honor hard-error with guidance to set
`use_lsm(false)`: hybrid, multi/binary vectors, `with_row_id`,
reranking, `order_by`, dynamic/Substrait projection or filters,
`distance_range`, `use_index(false)`, postfilter, take-by-row-id/offset,
reads from a time-traveled version, and an unmaintained or ambiguous
(multiple) FTS/vector index. Namespace-pushdown queries fall back to
local execution when a spec is present; WAL-only writers are handled.
- Exposed across the Rust core and the Python (`use_lsm`) and TypeScript
(`useLsm`) bindings, including `TakeQuery`.
Rebased from Lance `7.2.0-beta.3` to `10.0.0-beta.3`.
This commit is contained in:
@@ -391,6 +391,7 @@ class Query:
|
||||
def fast_search(self): ...
|
||||
def with_row_id(self): ...
|
||||
def postfilter(self): ...
|
||||
def use_lsm(self, enable: bool): ...
|
||||
def nearest_to(self, query_vec: pa.Array) -> VectorQuery: ...
|
||||
def nearest_to_text(self, query: dict) -> FTSQuery: ...
|
||||
def order_by(self, ordering: Optional[List[ColumnOrdering]]): ...
|
||||
@@ -407,6 +408,7 @@ class Query:
|
||||
class TakeQuery:
|
||||
def select(self, columns: List[str]): ...
|
||||
def with_row_id(self): ...
|
||||
def use_lsm(self, enable: bool): ...
|
||||
async def output_schema(self) -> pa.Schema: ...
|
||||
async def execute(self) -> RecordBatchStream: ...
|
||||
async def explain_plan(self, verbose: Optional[bool]) -> str: ...
|
||||
@@ -425,6 +427,7 @@ class FTSQuery:
|
||||
def fast_search(self): ...
|
||||
def with_row_id(self): ...
|
||||
def postfilter(self): ...
|
||||
def use_lsm(self, enable: bool): ...
|
||||
def get_query(self) -> str: ...
|
||||
def add_query_vector(self, query_vec: pa.Array) -> None: ...
|
||||
def nearest_to(self, query_vec: pa.Array) -> HybridQuery: ...
|
||||
@@ -452,6 +455,7 @@ class VectorQuery:
|
||||
def column(self, column: str): ...
|
||||
def distance_type(self, distance_type: str): ...
|
||||
def postfilter(self): ...
|
||||
def use_lsm(self, enable: bool): ...
|
||||
def refine_factor(self, refine_factor: int): ...
|
||||
def nprobes(self, nprobes: int): ...
|
||||
def minimum_nprobes(self, minimum_nprobes: int): ...
|
||||
@@ -475,6 +479,7 @@ class HybridQuery:
|
||||
def fast_search(self): ...
|
||||
def with_row_id(self): ...
|
||||
def postfilter(self): ...
|
||||
def use_lsm(self, enable: bool): ...
|
||||
def distance_type(self, distance_type: str): ...
|
||||
def refine_factor(self, refine_factor: int): ...
|
||||
def nprobes(self, nprobes: int): ...
|
||||
@@ -499,6 +504,7 @@ class PyQueryRequest:
|
||||
select: Optional[Union[str, List[str]]]
|
||||
fast_search: Optional[bool]
|
||||
with_row_id: Optional[bool]
|
||||
use_lsm: Optional[bool]
|
||||
column: Optional[str]
|
||||
query_vector: Optional[List[pa.Array]]
|
||||
minimum_nprobes: Optional[int]
|
||||
|
||||
@@ -37,7 +37,7 @@ class LanceMergeInsertBuilder(object):
|
||||
self._when_not_matched_by_source_condition_expr = None
|
||||
self._timeout = None
|
||||
self._use_index = True
|
||||
self._use_lsm_write = None
|
||||
self._use_lsm = None
|
||||
self._validate_single_shard = None
|
||||
|
||||
def when_matched_update_all(
|
||||
@@ -113,22 +113,22 @@ class LanceMergeInsertBuilder(object):
|
||||
self._use_index = use_index
|
||||
return self
|
||||
|
||||
def use_lsm_write(self, use_lsm_write: bool) -> LanceMergeInsertBuilder:
|
||||
def use_lsm(self, enable: bool) -> LanceMergeInsertBuilder:
|
||||
"""
|
||||
Controls whether the merge uses the MemWAL LSM write path.
|
||||
Control MemWAL routing for this merge.
|
||||
|
||||
By default (unset), a `merge_insert` on a table with an LSM write spec
|
||||
is routed through Lance's MemWAL shard writer, and a table without one
|
||||
uses the standard path. Pass `False` to force the standard path even
|
||||
when a spec is set. Pass `True` to require a spec — `merge_insert`
|
||||
raises an error if none is installed.
|
||||
By default (unset), a `merge_insert` on a table with an LSM write spec is
|
||||
routed through Lance's MemWAL shard writer, and a table without one uses
|
||||
the standard path.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
use_lsm_write: bool
|
||||
Whether to use the LSM write path.
|
||||
enable: bool
|
||||
``True`` forces MemWAL routing and errors if the table has no LSM
|
||||
write spec. ``False`` forces the standard write path even when a spec
|
||||
is set.
|
||||
"""
|
||||
self._use_lsm_write = use_lsm_write
|
||||
self._use_lsm = enable
|
||||
return self
|
||||
|
||||
def validate_single_shard(
|
||||
|
||||
@@ -778,6 +778,11 @@ class Query(pydantic.BaseModel):
|
||||
# if true, will only search the indexed data
|
||||
fast_search: Optional[bool] = None
|
||||
|
||||
# MemWAL LSM read routing: None auto-routes when the table carries a write
|
||||
# spec, True forces the LSM scanner (errors without a spec), False reads the
|
||||
# base table only
|
||||
use_lsm: Optional[bool] = None
|
||||
|
||||
# size of the nearest neighbor list maintained during HNSW search
|
||||
ef: Optional[int] = None
|
||||
|
||||
@@ -795,6 +800,9 @@ class Query(pydantic.BaseModel):
|
||||
query.full_text_query = req.full_text_search
|
||||
query.columns = req.select
|
||||
query.with_row_id = req.with_row_id
|
||||
# use_lsm is a genuine tri-state (None / True / False); preserve it as-is
|
||||
# so a round-tripped query keeps an explicit False.
|
||||
query.use_lsm = req.use_lsm
|
||||
query.vector_column = req.column
|
||||
query.vector = req.query_vector
|
||||
query.distance_type = req.distance_type
|
||||
@@ -967,6 +975,7 @@ class LanceQueryBuilder(ABC):
|
||||
self._with_row_address = None
|
||||
self._fragments = None
|
||||
self._fragment_ids = None
|
||||
self._use_lsm = None
|
||||
self._vector = None
|
||||
self._text = None
|
||||
self._ef = None
|
||||
@@ -1326,6 +1335,30 @@ class LanceQueryBuilder(ABC):
|
||||
self._fragment_ids = fragment_ids
|
||||
return self
|
||||
|
||||
def use_lsm(self, enable: bool) -> Self:
|
||||
"""Control MemWAL LSM read routing for this query.
|
||||
|
||||
By default (unset), a query against a table with an LSM write spec is
|
||||
routed through the LSM scanner so it also returns data written via the
|
||||
``merge_insert`` LSM path that has not yet been compacted into the base
|
||||
table (active/frozen memtables + flushed generations); a table without a
|
||||
spec reads the base table.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
enable : bool
|
||||
``True`` forces the LSM scanner and errors if the table has no LSM
|
||||
write spec. ``False`` bypasses the MemWAL and reads the base table
|
||||
only, even when a spec is present.
|
||||
|
||||
Returns
|
||||
-------
|
||||
LanceQueryBuilder
|
||||
The LanceQueryBuilder object.
|
||||
"""
|
||||
self._use_lsm = enable
|
||||
return self
|
||||
|
||||
def explain_plan(self, verbose: Optional[bool] = False) -> str:
|
||||
"""Return the execution plan for this query.
|
||||
|
||||
@@ -1788,6 +1821,7 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
|
||||
with_row_address=self._with_row_address,
|
||||
fragments=self._fragments,
|
||||
fragment_ids=self._fragment_ids,
|
||||
use_lsm=self._use_lsm,
|
||||
offset=self._offset,
|
||||
fast_search=self._fast_search,
|
||||
ef=self._ef,
|
||||
@@ -2012,6 +2046,7 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
||||
with_row_address=self._with_row_address,
|
||||
fragments=self._fragments,
|
||||
fragment_ids=self._fragment_ids,
|
||||
use_lsm=self._use_lsm,
|
||||
full_text_query=FullTextSearchQuery(
|
||||
query=self._query_with_phrase_semantics(), columns=self._fts_columns
|
||||
),
|
||||
@@ -2078,6 +2113,7 @@ class LanceEmptyQueryBuilder(LanceQueryBuilder):
|
||||
with_row_address=self._with_row_address,
|
||||
fragments=self._fragments,
|
||||
fragment_ids=self._fragment_ids,
|
||||
use_lsm=self._use_lsm,
|
||||
offset=self._offset,
|
||||
order_by=self._order_by,
|
||||
)
|
||||
@@ -2655,6 +2691,9 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
if self._with_row_id:
|
||||
self._vector_query.with_row_id(True)
|
||||
self._fts_query.with_row_id(True)
|
||||
if self._use_lsm is not None:
|
||||
self._vector_query.use_lsm(self._use_lsm)
|
||||
self._fts_query.use_lsm(self._use_lsm)
|
||||
if self._phrase_query:
|
||||
self._fts_query.phrase_query(True)
|
||||
if self._distance_type:
|
||||
@@ -3231,6 +3270,27 @@ class AsyncStandardQuery(AsyncQueryBase):
|
||||
self._inner.fast_search()
|
||||
return self
|
||||
|
||||
def use_lsm(self, enable: bool) -> Self:
|
||||
"""
|
||||
Control MemWAL LSM read routing for this query.
|
||||
|
||||
By default (unset), a query against a table with an LSM write spec (see
|
||||
[AsyncTable.set_lsm_write_spec][lancedb.table.AsyncTable.set_lsm_write_spec])
|
||||
is routed through the LSM scanner so it also returns data written via the
|
||||
``merge_insert`` LSM path that has not yet been compacted into the base
|
||||
table (the active/frozen in-memory memtables and the flushed generations),
|
||||
deduplicated by primary key; a table without a spec reads the base table.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
enable : bool
|
||||
``True`` forces the LSM scanner and errors if the table has no LSM
|
||||
write spec. ``False`` bypasses the MemWAL and reads the base table
|
||||
only, even when a spec is present.
|
||||
"""
|
||||
self._inner.use_lsm(enable)
|
||||
return self
|
||||
|
||||
def postfilter(self) -> Self:
|
||||
"""
|
||||
If this is called then filtering will happen after the search instead of
|
||||
@@ -3944,6 +4004,15 @@ class AsyncTakeQuery(AsyncQueryBase):
|
||||
def __init__(self, inner: LanceTakeQuery, table: Optional["AsyncTable"] = None):
|
||||
super().__init__(inner, table)
|
||||
|
||||
def use_lsm(self, enable: bool) -> "AsyncTakeQuery":
|
||||
"""Control MemWAL LSM read routing for this take query.
|
||||
|
||||
``False`` bypasses the MemWAL and reads the base table only — the escape
|
||||
hatch, since take-by-row-id/offset is not supported on the LSM scanner.
|
||||
"""
|
||||
self._inner.use_lsm(enable)
|
||||
return self
|
||||
|
||||
async def _plain_scan_to_pandas(
|
||||
self,
|
||||
blob_mode: BlobMode,
|
||||
@@ -4002,6 +4071,16 @@ class BaseQueryBuilder(object):
|
||||
self._inner.with_row_id()
|
||||
return self
|
||||
|
||||
def use_lsm(self, enable: bool) -> Self:
|
||||
"""
|
||||
Control MemWAL LSM read routing for this query.
|
||||
|
||||
``False`` bypasses the MemWAL and reads the base table only, the escape
|
||||
hatch for shapes the LSM scanner cannot honor (e.g. take-by-row-id).
|
||||
"""
|
||||
self._inner.use_lsm(enable)
|
||||
return self
|
||||
|
||||
def with_row_address(self, with_row_address: bool = True) -> Self:
|
||||
"""
|
||||
Include the _rowaddr column in scanner-backed plain query results.
|
||||
|
||||
@@ -5363,6 +5363,8 @@ class AsyncTable:
|
||||
async_query = async_query.where(query.filter)
|
||||
if query.fast_search:
|
||||
async_query = async_query.fast_search()
|
||||
if query.use_lsm is not None:
|
||||
async_query = async_query.use_lsm(query.use_lsm)
|
||||
if query.with_row_id:
|
||||
async_query = async_query.with_row_id()
|
||||
if query.order_by:
|
||||
@@ -5483,7 +5485,7 @@ class AsyncTable:
|
||||
when_not_matched_by_source_condition_expr=merge._when_not_matched_by_source_condition_expr,
|
||||
timeout=merge._timeout,
|
||||
use_index=merge._use_index,
|
||||
use_lsm_write=merge._use_lsm_write,
|
||||
use_lsm=merge._use_lsm,
|
||||
validate_single_shard=merge._validate_single_shard,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -9,6 +9,7 @@ import lancedb
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from lancedb._lancedb import LsmWriteSpec
|
||||
from lancedb.index import FTS, IvfPq
|
||||
|
||||
SCHEMA = pa.schema(
|
||||
[
|
||||
@@ -102,19 +103,35 @@ def test_lsm_merge_insert_identity(tmp_path):
|
||||
assert result.num_rows == 2
|
||||
|
||||
|
||||
def test_lsm_merge_insert_use_lsm_write_false(tmp_path):
|
||||
def test_lsm_merge_insert_use_lsm_false(tmp_path):
|
||||
table = _bucket_table(tmp_path) # rows id = 1, 2, 3
|
||||
# use_lsm_write(False) opts out: the standard path runs and commits.
|
||||
# use_lsm(False) opts out: the standard path runs and commits even with a spec.
|
||||
result = (
|
||||
table.merge_insert("id")
|
||||
.when_not_matched_insert_all()
|
||||
.use_lsm_write(False)
|
||||
.use_lsm(False)
|
||||
.execute(_reader([3, 4, 5]))
|
||||
)
|
||||
assert result.num_inserted_rows == 2
|
||||
assert table.count_rows() == 5
|
||||
|
||||
|
||||
def test_lsm_merge_insert_use_lsm_true_without_spec_errors(tmp_path):
|
||||
# A table with a primary key but no LSM write spec installed.
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _reader([1, 2, 3]))
|
||||
table.set_unenforced_primary_key("id")
|
||||
# use_lsm(True) demands MemWAL routing; without a spec it errors.
|
||||
with pytest.raises(Exception, match="use_lsm"):
|
||||
(
|
||||
table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.use_lsm(True)
|
||||
.execute(_reader([3, 4, 5]))
|
||||
)
|
||||
|
||||
|
||||
def test_lsm_merge_insert_validate_single_shard_off(tmp_path):
|
||||
table = _bucket_table(tmp_path)
|
||||
result = (
|
||||
@@ -127,19 +144,20 @@ def test_lsm_merge_insert_validate_single_shard_off(tmp_path):
|
||||
assert result.num_rows == 3
|
||||
|
||||
|
||||
def test_lsm_merge_insert_use_lsm_write_true_requires_spec(tmp_path):
|
||||
def test_lsm_merge_insert_no_spec_uses_standard_path(tmp_path):
|
||||
# A table with a primary key but no LSM write spec installed.
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _reader([1, 2, 3]))
|
||||
table.set_unenforced_primary_key("id")
|
||||
with pytest.raises(Exception, match="use_lsm_write"):
|
||||
(
|
||||
table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.use_lsm_write(True)
|
||||
.execute(_reader([4]))
|
||||
)
|
||||
# With no spec, a default merge_insert uses the standard path and commits.
|
||||
result = (
|
||||
table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.execute(_reader([3, 4, 5]))
|
||||
)
|
||||
assert result.num_inserted_rows == 2
|
||||
assert table.count_rows() == 5
|
||||
|
||||
|
||||
def test_lsm_merge_insert_rejects_on_not_primary_key(tmp_path):
|
||||
@@ -194,3 +212,445 @@ async def test_async_lsm_merge_insert(tmp_path):
|
||||
result = await builder.execute(_reader([3, 4, 5]))
|
||||
assert result.num_rows == 3
|
||||
await table.close_lsm_writers()
|
||||
|
||||
|
||||
def _lsm_upsert(table, ids):
|
||||
"""Upsert ``ids`` (value = 0..n) through the LSM merge_insert path."""
|
||||
(
|
||||
table.merge_insert([])
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.execute(_reader(ids))
|
||||
)
|
||||
|
||||
|
||||
def test_lsm_read_sees_active_memtable(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _reader([1, 2, 3])) # base ids 1,2,3
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
|
||||
|
||||
_lsm_upsert(table, [4, 5]) # active memtable only, not committed to base
|
||||
|
||||
# Default read auto-routes through the LSM scanner: base ∪ active memtable.
|
||||
lsm = table.search().to_arrow()
|
||||
assert sorted(lsm["id"].to_pylist()) == [1, 2, 3, 4, 5]
|
||||
|
||||
# use_lsm(False) bypasses the MemWAL and reads the base table only.
|
||||
base_only = table.search().use_lsm(False).to_arrow()
|
||||
assert sorted(base_only["id"].to_pylist()) == [1, 2, 3]
|
||||
|
||||
|
||||
def test_lsm_read_dedup_newest_wins(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _reader([1, 2, 3])) # id 2 -> value 1
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
|
||||
|
||||
_lsm_upsert(table, [2, 3, 4]) # ids 2,3,4 -> values 0,1,2
|
||||
|
||||
lsm = table.search().to_arrow().sort_by("id")
|
||||
assert lsm["id"].to_pylist() == [1, 2, 3, 4]
|
||||
# id 1 from base (value 0); 2,3,4 from memtable (values 0,1,2).
|
||||
assert lsm["value"].to_pylist() == [0, 0, 1, 2]
|
||||
|
||||
|
||||
def test_lsm_read_without_spec_reads_base(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _reader([1, 2, 3]))
|
||||
table.set_unenforced_primary_key("id") # no LSM write spec
|
||||
|
||||
# No spec: default read and use_lsm(False) both read the base table, no error.
|
||||
assert sorted(table.search().to_arrow()["id"].to_pylist()) == [1, 2, 3]
|
||||
assert sorted(table.search().use_lsm(False).to_arrow()["id"].to_pylist()) == [
|
||||
1,
|
||||
2,
|
||||
3,
|
||||
]
|
||||
|
||||
|
||||
def test_lsm_read_unsupported_shape_errors_without_use_lsm_false(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _reader([1, 2, 3]))
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
|
||||
_lsm_upsert(table, [4])
|
||||
|
||||
# with_row_id is unsupported by the LSM scanner; on a MemWAL table the default
|
||||
# (auto-routed) read hard-errors instead of silently reading a stale base.
|
||||
with pytest.raises(Exception):
|
||||
table.search().with_row_id(True).to_arrow()
|
||||
|
||||
# use_lsm(False) is the escape hatch: it reads the base table only.
|
||||
base = table.search().with_row_id(True).use_lsm(False).to_arrow()
|
||||
assert sorted(base["id"].to_pylist()) == [1, 2, 3]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_lsm_read(tmp_path):
|
||||
db = await lancedb.connect_async(
|
||||
tmp_path, read_consistency_interval=timedelta(seconds=0)
|
||||
)
|
||||
table = await db.create_table("t", _reader([1, 2, 3]))
|
||||
await table.set_unenforced_primary_key("id")
|
||||
await table.set_lsm_write_spec(LsmWriteSpec.unsharded())
|
||||
|
||||
builder = (
|
||||
table.merge_insert([]).when_matched_update_all().when_not_matched_insert_all()
|
||||
)
|
||||
await builder.execute(_reader([4, 5]))
|
||||
|
||||
arrow = await table.query().to_arrow()
|
||||
assert sorted(arrow["id"].to_pylist()) == [1, 2, 3, 4, 5]
|
||||
|
||||
|
||||
VECTOR_DIM = 8
|
||||
|
||||
VECTOR_SCHEMA = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64(), nullable=False),
|
||||
pa.field("category", pa.utf8(), nullable=False),
|
||||
pa.field("vector", pa.list_(pa.float32(), VECTOR_DIM), nullable=False),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _vector_reader(rows):
|
||||
"""Rows are ``(id, category, [f32; VECTOR_DIM])`` tuples."""
|
||||
batch = pa.RecordBatch.from_arrays(
|
||||
[
|
||||
pa.array([row[0] for row in rows], type=pa.int64()),
|
||||
pa.array([row[1] for row in rows], type=pa.utf8()),
|
||||
pa.array([row[2] for row in rows], type=pa.list_(pa.float32(), VECTOR_DIM)),
|
||||
],
|
||||
schema=VECTOR_SCHEMA,
|
||||
)
|
||||
return pa.RecordBatchReader.from_batches(VECTOR_SCHEMA, [batch])
|
||||
|
||||
|
||||
def _vector_table(tmp_path):
|
||||
"""Base table whose vector column is indexed so its rows are visible to the LSM
|
||||
vector scanner (the base arm uses ``fast_search`` — indexed data only), plus an
|
||||
unsharded LSM spec that maintains that index for the memtable.
|
||||
|
||||
Rows 1,2 are category ``a``, row 3 is ``b``, and 4..60 are filler ``c`` that
|
||||
give the tiny IVF index enough data to train.
|
||||
"""
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
rows = [
|
||||
(
|
||||
i,
|
||||
"a" if i in (1, 2) else "b" if i == 3 else "c",
|
||||
[float((i * 7 + j) % 13) for j in range(VECTOR_DIM)],
|
||||
)
|
||||
for i in range(1, 61)
|
||||
]
|
||||
table = db.create_table("t", _vector_reader(rows))
|
||||
table.set_unenforced_primary_key("id")
|
||||
# num_partitions=1 makes the search exhaustive within the single partition
|
||||
# (deterministic); num_bits=4 keeps PQ training viable on a tiny dataset.
|
||||
table.create_index(
|
||||
"vector", config=IvfPq(num_partitions=1, num_sub_vectors=2, num_bits=4)
|
||||
)
|
||||
index_name = table.list_indices()[0].name
|
||||
table.set_lsm_write_spec(
|
||||
LsmWriteSpec.unsharded().with_maintained_indexes([index_name])
|
||||
)
|
||||
return table
|
||||
|
||||
|
||||
def _vector_upsert(table, rows):
|
||||
(
|
||||
table.merge_insert([])
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.execute(_vector_reader(rows))
|
||||
)
|
||||
|
||||
|
||||
def test_lsm_read_vector_sees_memtable(tmp_path):
|
||||
table = _vector_table(tmp_path)
|
||||
# id 1000 lands in the active memtable, not committed to the base table.
|
||||
_vector_upsert(table, [(1000, "a", [1.0] * VECTOR_DIM)])
|
||||
|
||||
query = [1.0] * VECTOR_DIM
|
||||
# Vector search auto-routes through the LSM scanner: indexed base ∪ memtable.
|
||||
ids = set(table.search(query).limit(100).to_arrow()["id"].to_pylist())
|
||||
assert {1, 2, 3} <= ids # indexed base rows
|
||||
assert 1000 in ids # in-flight memtable row
|
||||
|
||||
# use_lsm(False) bypasses the MemWAL, so the in-flight row is not visible.
|
||||
base_ids = set(
|
||||
table.search(query).use_lsm(False).limit(100).to_arrow()["id"].to_pylist()
|
||||
)
|
||||
assert {1, 2, 3} <= base_ids
|
||||
assert 1000 not in base_ids
|
||||
|
||||
|
||||
def test_lsm_read_vector_prefilter(tmp_path):
|
||||
table = _vector_table(tmp_path)
|
||||
# in-flight rows in both categories.
|
||||
_vector_upsert(
|
||||
table, [(1000, "a", [1.0] * VECTOR_DIM), (1001, "b", [1.0] * VECTOR_DIM)]
|
||||
)
|
||||
|
||||
query = [1.0] * VECTOR_DIM
|
||||
# The `where` predicate must apply as a prefilter across base ∪ memtable —
|
||||
# regression test for the vector arm silently dropping the filter.
|
||||
rows = table.search(query).where("category = 'a'").limit(100).to_arrow()
|
||||
assert set(rows["id"].to_pylist()) == {1, 2, 1000}
|
||||
assert set(rows["category"].to_pylist()) == {"a"}
|
||||
|
||||
# Sanity: without the filter, other categories are returned too.
|
||||
unfiltered = set(table.search(query).limit(100).to_arrow()["category"].to_pylist())
|
||||
assert unfiltered != {"a"}
|
||||
|
||||
|
||||
def test_lsm_read_plain_prefilter(tmp_path):
|
||||
table = _vector_table(tmp_path)
|
||||
_vector_upsert(
|
||||
table, [(1000, "a", [1.0] * VECTOR_DIM), (1001, "b", [1.0] * VECTOR_DIM)]
|
||||
)
|
||||
|
||||
# Plain scan + filter over base ∪ memtable: base 'a' rows 1,2 and memtable 1000.
|
||||
rows = table.search().where("category = 'a'").to_arrow()
|
||||
assert set(rows["id"].to_pylist()) == {1, 2, 1000}
|
||||
|
||||
|
||||
FTS_SCHEMA = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64(), nullable=False),
|
||||
pa.field("text", pa.utf8(), nullable=False),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _fts_reader(rows):
|
||||
"""Rows are ``(id, text)`` tuples."""
|
||||
batch = pa.RecordBatch.from_arrays(
|
||||
[
|
||||
pa.array([row[0] for row in rows], type=pa.int64()),
|
||||
pa.array([row[1] for row in rows], type=pa.utf8()),
|
||||
],
|
||||
schema=FTS_SCHEMA,
|
||||
)
|
||||
return pa.RecordBatchReader.from_batches(FTS_SCHEMA, [batch])
|
||||
|
||||
|
||||
def test_lsm_read_fts_sees_memtable(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table(
|
||||
"t",
|
||||
_fts_reader(
|
||||
[
|
||||
(1, "the quick brown fox"),
|
||||
(2, "lazy dog sleeps"),
|
||||
(3, "quick red fox"),
|
||||
]
|
||||
),
|
||||
)
|
||||
table.set_unenforced_primary_key("id")
|
||||
# Native FTS index (tantivy is not compatible with the LSM memtable index).
|
||||
table.create_index("text", config=FTS())
|
||||
index_name = table.list_indices()[0].name
|
||||
table.set_lsm_write_spec(
|
||||
LsmWriteSpec.unsharded().with_maintained_indexes([index_name])
|
||||
)
|
||||
|
||||
# in-flight doc 4 lands in the memtable's maintained FTS index.
|
||||
(
|
||||
table.merge_insert([])
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.execute(_fts_reader([(4, "brown fox jumps")]))
|
||||
)
|
||||
|
||||
# Full-text search auto-routes through the LSM scanner: base ∪ memtable.
|
||||
ids = set(
|
||||
table.search("fox", query_type="fts", fts_columns="text")
|
||||
.limit(10)
|
||||
.to_arrow()["id"]
|
||||
.to_pylist()
|
||||
)
|
||||
assert ids == {1, 3, 4}
|
||||
|
||||
# Prefilter restricts the FTS results across both tiers.
|
||||
filtered = set(
|
||||
table.search("fox", query_type="fts", fts_columns="text")
|
||||
.where("id > 1")
|
||||
.limit(10)
|
||||
.to_arrow()["id"]
|
||||
.to_pylist()
|
||||
)
|
||||
assert filtered == {3, 4}
|
||||
|
||||
|
||||
def test_lsm_read_vector_unsupported_knobs_error(tmp_path):
|
||||
table = _vector_table(tmp_path)
|
||||
_vector_upsert(table, [(1000, "a", [1.0] * VECTOR_DIM)])
|
||||
query = [1.0] * VECTOR_DIM
|
||||
|
||||
# distance_range and use_index(False) change the vector result set/mode, which
|
||||
# the LSM scanner can't honor, so it hard-errors instead of silently returning
|
||||
# wrong results (matching the prefilter / unsupported-shape contract).
|
||||
with pytest.raises(Exception, match="distance_range"):
|
||||
table.search(query).distance_range(0.0, 0.5).to_arrow()
|
||||
with pytest.raises(Exception, match="use_index"):
|
||||
table.search(query).bypass_vector_index().to_arrow()
|
||||
|
||||
# use_lsm(False) is the escape hatch: the base-only standard path honors them.
|
||||
base = table.search(query).distance_range(0.0, 100.0).use_lsm(False).to_arrow()
|
||||
assert 1000 not in set(base["id"].to_pylist())
|
||||
|
||||
|
||||
def test_lsm_read_vector_limit_offset(tmp_path):
|
||||
table = _vector_table(tmp_path)
|
||||
_vector_upsert(table, [(1000, "a", [1.0] * VECTOR_DIM)])
|
||||
query = [1.0] * VECTOR_DIM
|
||||
# Lance's plan_vector over-fetches k + offset internally, so paging is correct:
|
||||
# the second page is a full page (not truncated) and disjoint from the first.
|
||||
page1 = table.search(query).limit(3).offset(0).to_arrow()["id"].to_pylist()
|
||||
page2 = table.search(query).limit(3).offset(3).to_arrow()["id"].to_pylist()
|
||||
assert len(page1) == 3
|
||||
# If k ignored offset, page2 would be empty (limit - offset = 0); a full second
|
||||
# page that differs from the first proves offset widens the candidate pool.
|
||||
assert len(page2) == 3
|
||||
assert set(page1) != set(page2)
|
||||
|
||||
|
||||
def test_lsm_read_vector_postfilter_errors(tmp_path):
|
||||
table = _vector_table(tmp_path)
|
||||
_vector_upsert(table, [(1000, "a", [1.0] * VECTOR_DIM)])
|
||||
query = [1.0] * VECTOR_DIM
|
||||
# The LSM scanner always prefilters; a requested postfilter changes results, so
|
||||
# it hard-errors rather than silently prefiltering.
|
||||
with pytest.raises(Exception, match="postfilter"):
|
||||
table.search(query).where("category = 'a'").postfilter().to_arrow()
|
||||
|
||||
|
||||
def test_lsm_read_projection_excludes_pk(tmp_path):
|
||||
table = _vector_table(tmp_path)
|
||||
_vector_upsert(table, [(1000, "a", [1.0] * VECTOR_DIM)])
|
||||
# Selecting only 'category' must not leak the 'id' primary key Lance appends
|
||||
# internally for dedup.
|
||||
rows = table.search().select(["category"]).where("category = 'a'").to_arrow()
|
||||
assert rows.column_names == ["category"]
|
||||
|
||||
|
||||
def test_lsm_read_fts_unmaintained_index_errors(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _fts_reader([(1, "quick fox"), (2, "lazy dog")]))
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.create_index("text", config=FTS())
|
||||
# No maintained indexes: the active memtable FTS arm cannot serve un-compacted
|
||||
# docs, so the search would silently omit them — reject instead.
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
|
||||
with pytest.raises(Exception, match="maintained"):
|
||||
table.search("fox", query_type="fts", fts_columns="text").to_arrow()
|
||||
|
||||
|
||||
def test_lsm_read_time_travel_errors(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _reader([1, 2, 3]))
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
|
||||
pinned = table.version
|
||||
table.add(_reader([4, 5])) # standard add commits a newer version
|
||||
table.checkout(pinned) # detached head at the historical version
|
||||
|
||||
# The WAL/manifest expose current live state, so an LSM read at a pinned
|
||||
# historical version is rejected.
|
||||
with pytest.raises(Exception, match="time-travel"):
|
||||
table.search().to_arrow()
|
||||
# use_lsm(False) reads the base table at the pinned version.
|
||||
base = table.search().use_lsm(False).to_arrow()
|
||||
assert sorted(base["id"].to_pylist()) == [1, 2, 3]
|
||||
|
||||
|
||||
def test_lsm_read_take_row_ids_errors(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _reader([1, 2, 3]))
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
|
||||
_lsm_upsert(table, [4])
|
||||
# take-by-row-id auto-routes through the LSM scanner, which has no stable _rowid,
|
||||
# so it hard-errors instead of failing with an opaque column-not-found error.
|
||||
with pytest.raises(Exception, match="row id"):
|
||||
table.take_row_ids([0, 1]).to_arrow()
|
||||
# use_lsm(False) is the escape hatch: it reads the base table.
|
||||
base = table.take_row_ids([0, 1]).use_lsm(False).to_arrow()
|
||||
assert base.num_rows == 2
|
||||
|
||||
|
||||
def test_lsm_read_fts_postfilter_errors(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _fts_reader([(1, "quick fox"), (2, "lazy dog")]))
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.create_index("text", config=FTS())
|
||||
index_name = table.list_indices()[0].name
|
||||
table.set_lsm_write_spec(
|
||||
LsmWriteSpec.unsharded().with_maintained_indexes([index_name])
|
||||
)
|
||||
# The LSM scanner always prefilters; postfilter on FTS changes result semantics,
|
||||
# so it hard-errors (previously only the vector arm rejected it).
|
||||
with pytest.raises(Exception, match="postfilter"):
|
||||
(
|
||||
table.search("fox", query_type="fts", fts_columns="text")
|
||||
.where("id > 0")
|
||||
.postfilter()
|
||||
.to_arrow()
|
||||
)
|
||||
|
||||
|
||||
def test_lsm_read_fts_multiple_same_type_indexes_errors(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _fts_reader([(1, "quick fox"), (2, "lazy dog")]))
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.create_index("text", config=FTS(), name="fts_a")
|
||||
table.create_index("text", config=FTS(), name="fts_b", replace=False)
|
||||
table.set_lsm_write_spec(
|
||||
LsmWriteSpec.unsharded().with_maintained_indexes(["fts_a"])
|
||||
)
|
||||
# Two FTS indexes on the column: the base planner's chosen index is ambiguous, so
|
||||
# the scanner can't pick a catch-up watermark and rejects rather than risk
|
||||
# dropping rows the actually-used index has not caught up to.
|
||||
with pytest.raises(Exception, match="multiple"):
|
||||
table.search("fox", query_type="fts", fts_columns="text").to_arrow()
|
||||
|
||||
|
||||
def test_lsm_read_vector_unmaintained_index_errors(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
rows = [
|
||||
(i, "a", [float((i * 7 + j) % 13) for j in range(VECTOR_DIM)])
|
||||
for i in range(1, 61)
|
||||
]
|
||||
table = db.create_table("t", _vector_reader(rows))
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.create_index(
|
||||
"vector", config=IvfPq(num_partitions=1, num_sub_vectors=2, num_bits=4)
|
||||
)
|
||||
# Spec with NO maintained indexes: the base vector index's catch-up is untracked,
|
||||
# so the scanner rejects rather than risk dropping compacted-but-unindexed rows.
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
|
||||
with pytest.raises(Exception, match="maintained"):
|
||||
table.search([1.0] * VECTOR_DIM).to_arrow()
|
||||
|
||||
|
||||
def test_lsm_read_fts_optimized_index_not_rejected(tmp_path):
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
||||
table = db.create_table("t", _fts_reader([(i, "quick fox") for i in range(1, 6)]))
|
||||
table.set_unenforced_primary_key("id")
|
||||
table.create_index("text", config=FTS())
|
||||
table.add(_fts_reader([(i, "lazy fox") for i in range(6, 11)]))
|
||||
table.optimize() # may split the FTS index into multiple physical segments
|
||||
name = table.list_indices()[0].name
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([name]))
|
||||
# Multiple physical segments of one logical index must not be miscounted as
|
||||
# multiple indexes and rejected.
|
||||
ids = set(
|
||||
table.search("fox", query_type="fts", fts_columns="text")
|
||||
.limit(20)
|
||||
.to_arrow()["id"]
|
||||
.to_pylist()
|
||||
)
|
||||
assert ids == set(range(1, 11))
|
||||
|
||||
Reference in New Issue
Block a user