diff --git a/docs/src/js/classes/MergeInsertBuilder.md b/docs/src/js/classes/MergeInsertBuilder.md index ac0493bad..beb6cdfce 100644 --- a/docs/src/js/classes/MergeInsertBuilder.md +++ b/docs/src/js/classes/MergeInsertBuilder.md @@ -76,24 +76,23 @@ the query optimizer chooses a suboptimal path. *** -### useLsmWrite() +### useLsm() ```ts -useLsmWrite(useLsmWrite): MergeInsertBuilder +useLsm(enable): MergeInsertBuilder ``` -Controls whether the merge uses the MemWAL LSM write path. +Control MemWAL routing for this merge. By default (unset), a `mergeInsert` 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 — `mergeInsert` rejects if none -is installed. +routed through Lance's MemWAL shard writer, and a table without one uses the +standard path. #### Parameters -* **useLsmWrite**: `boolean` - Whether to use the LSM write path. +* **enable**: `boolean` + `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. #### Returns diff --git a/docs/src/js/classes/Query.md b/docs/src/js/classes/Query.md index c4bc3ff23..6ebaebd75 100644 --- a/docs/src/js/classes/Query.md +++ b/docs/src/js/classes/Query.md @@ -497,6 +497,42 @@ ArrowTable. *** +### useLsm() + +```ts +useLsm(enable): this +``` + +Control MemWAL read routing for this query. + +By default (unset), when the table carries a MemWAL write spec (see +[Table#setLsmWriteSpec](Table.md#setlsmwritespec)), reads are routed through the LSM scanner so +they also return data written via the `mergeInsert` 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**: `boolean` + `true` forces the LSM scanner and errors if the table has no + MemWAL write spec. `false` bypasses the MemWAL and reads the base table only, + even when a spec is present. + Note: the LSM scanner does not support every query shape (e.g. reranking, + hybrid search, `orderBy`). On a MemWAL table those shapes error unless + `useLsm(false)` is set, because a base-only read would silently exclude + un-compacted MemWAL data. + +#### Returns + +`this` + +#### Inherited from + +`StandardQueryBase.useLsm` + +*** + ### where() ```ts diff --git a/docs/src/js/classes/TakeQuery.md b/docs/src/js/classes/TakeQuery.md index 0d658b32c..6ae2f9c8c 100644 --- a/docs/src/js/classes/TakeQuery.md +++ b/docs/src/js/classes/TakeQuery.md @@ -273,6 +273,29 @@ ArrowTable. *** +### useLsm() + +```ts +useLsm(enable): this +``` + +Control MemWAL 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 and, on a +MemWAL table, auto-routes to it and errors otherwise. + +#### Parameters + +* **enable**: `boolean` + `false` reads the base table only. + +#### Returns + +`this` + +*** + ### withRowId() ```ts diff --git a/docs/src/js/classes/VectorQuery.md b/docs/src/js/classes/VectorQuery.md index af3317029..f9412c76d 100644 --- a/docs/src/js/classes/VectorQuery.md +++ b/docs/src/js/classes/VectorQuery.md @@ -746,6 +746,42 @@ ArrowTable. *** +### useLsm() + +```ts +useLsm(enable): this +``` + +Control MemWAL read routing for this query. + +By default (unset), when the table carries a MemWAL write spec (see +[Table#setLsmWriteSpec](Table.md#setlsmwritespec)), reads are routed through the LSM scanner so +they also return data written via the `mergeInsert` 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**: `boolean` + `true` forces the LSM scanner and errors if the table has no + MemWAL write spec. `false` bypasses the MemWAL and reads the base table only, + even when a spec is present. + Note: the LSM scanner does not support every query shape (e.g. reranking, + hybrid search, `orderBy`). On a MemWAL table those shapes error unless + `useLsm(false)` is set, because a base-only read would silently exclude + un-compacted MemWAL data. + +#### Returns + +`this` + +#### Inherited from + +`StandardQueryBase.useLsm` + +*** + ### where() ```ts diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 44edaa093..15248f5cf 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -527,6 +527,14 @@ describe.each([arrow15, arrow16, arrow17, arrow18])( ); }); + it("should expose useLsm on takeRowIds as the base-only escape hatch", async () => { + await table.add([{ id: 1 }, { id: 2 }, { id: 3 }]); + // useLsm(false) is reachable on TakeQuery (the escape hatch for MemWAL tables, + // where take-by-row-id auto-routes to the LSM scanner and is rejected). + const res = await table.takeRowIds([0, 2]).useLsm(false).toArray(); + expect(res.map((r) => r.id)).toEqual([1, 3]); + }); + it("should throw for negative number in takeRowIds", () => { expect(() => table.takeRowIds([-1])).toThrow("Row id cannot be negative"); expect(() => table.takeRowIds([0, -5, 2])).toThrow( @@ -3199,14 +3207,14 @@ describe("LSM merge insert", () => { await table.closeLsmWriters(); }); - it("falls back to the standard path with useLsmWrite(false)", async () => { + it("falls back to the standard path with useLsm(false)", async () => { const conn = await connect(tmpDir.name); const table = await bucketTable(conn); const res = await table .mergeInsert("id") .whenNotMatchedInsertAll() - .useLsmWrite(false) + .useLsm(false) .execute([ { id: "b", value: 9 }, { id: "e", value: 5 }, @@ -3240,4 +3248,36 @@ describe("LSM merge insert", () => { .execute([{ id: "g", value: 7 }]), ).rejects.toThrow(); }); + + it("auto-routes reads through the MemWAL scanner", async () => { + const conn = await connect(tmpDir.name); + const table = await bucketTable(conn); // base ids "a", "b" + + await table + .mergeInsert("id") + .whenMatchedUpdateAll() + .whenNotMatchedInsertAll() + .execute([{ id: "c", value: 3 }]); + + // Default read auto-routes and includes the active memtable row. + const lsm = await table.query().toArray(); + expect(lsm.map((r) => r.id).sort()).toEqual(["a", "b", "c"]); + + // useLsm(false) bypasses the MemWAL and reads the base table only. + const baseOnly = await table.query().useLsm(false).toArray(); + expect(baseOnly.map((r) => r.id).sort()).toEqual(["a", "b"]); + }); + + it("reads the base table when no LSM spec is installed", async () => { + const conn = await connect(tmpDir.name); + const table = await conn.createEmptyTable( + "plain", + new arrow.Schema([new arrow.Field("id", new arrow.Utf8(), false)]), + ); + // No spec: default read and useLsm(false) both succeed against the base table. + await expect(table.query().toArray()).resolves.toBeDefined(); + await expect(table.query().useLsm(false).toArray()).resolves.toBeDefined(); + // useLsm(true) demands MemWAL routing; without a spec it errors. + await expect(table.query().useLsm(true).toArray()).rejects.toThrow(); + }); }); diff --git a/nodejs/lancedb/merge.ts b/nodejs/lancedb/merge.ts index 08321427f..30bde7281 100644 --- a/nodejs/lancedb/merge.ts +++ b/nodejs/lancedb/merge.ts @@ -88,21 +88,17 @@ export class MergeInsertBuilder { ); } /** - * Controls whether the merge uses the MemWAL LSM write path. + * Control MemWAL routing for this merge. * * By default (unset), a `mergeInsert` 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 — `mergeInsert` rejects if none - * is installed. + * routed through Lance's MemWAL shard writer, and a table without one uses the + * standard path. * - * @param useLsmWrite - Whether to use the LSM write path. + * @param enable - `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. */ - useLsmWrite(useLsmWrite: boolean): MergeInsertBuilder { - return new MergeInsertBuilder( - this.#native.useLsmWrite(useLsmWrite), - this.#schema, - ); + useLsm(enable: boolean): MergeInsertBuilder { + return new MergeInsertBuilder(this.#native.useLsm(enable), this.#schema); } /** * Controls how an LSM merge checks that its input targets a single shard. diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 7c9d96750..843a1276f 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -460,6 +460,30 @@ export class StandardQueryBase< this.doCall((inner: NativeQueryType) => inner.fastSearch()); return this; } + + /** + * Control MemWAL read routing for this query. + * + * By default (unset), when the table carries a MemWAL write spec (see + * {@link Table#setLsmWriteSpec}), reads are routed through the LSM scanner so + * they also return data written via the `mergeInsert` 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. + * + * @param enable - `true` forces the LSM scanner and errors if the table has no + * MemWAL write spec. `false` bypasses the MemWAL and reads the base table only, + * even when a spec is present. + * + * Note: the LSM scanner does not support every query shape (e.g. reranking, + * hybrid search, `orderBy`). On a MemWAL table those shapes error unless + * `useLsm(false)` is set, because a base-only read would silently exclude + * un-compacted MemWAL data. + */ + useLsm(enable: boolean): this { + this.doCall((inner: NativeQueryType) => inner.useLsm(enable)); + return this; + } } /** @@ -748,6 +772,20 @@ export class TakeQuery extends QueryBase { constructor(inner: NativeTakeQuery) { super(inner); } + + /** + * Control MemWAL 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 and, on a + * MemWAL table, auto-routes to it and errors otherwise. + * + * @param enable - `false` reads the base table only. + */ + useLsm(enable: boolean): this { + this.doCall((inner: NativeTakeQuery) => inner.useLsm(enable)); + return this; + } } /** A builder for LanceDB queries. diff --git a/nodejs/src/merge.rs b/nodejs/src/merge.rs index 1f9609160..ac8649419 100644 --- a/nodejs/src/merge.rs +++ b/nodejs/src/merge.rs @@ -51,9 +51,9 @@ impl NativeMergeInsertBuilder { } #[napi] - pub fn use_lsm_write(&self, use_lsm_write: bool) -> Self { + pub fn use_lsm(&self, enable: bool) -> Self { let mut this = self.clone(); - this.inner.use_lsm_write(use_lsm_write); + this.inner.use_lsm(enable); this } diff --git a/nodejs/src/query.rs b/nodejs/src/query.rs index 7266e05fc..3828023a9 100644 --- a/nodejs/src/query.rs +++ b/nodejs/src/query.rs @@ -168,6 +168,11 @@ impl Query { self.inner = self.inner.clone().with_row_id(); } + #[napi] + pub fn use_lsm(&mut self, enable: bool) { + self.inner = self.inner.clone().use_lsm(enable); + } + #[napi] pub fn order_by(&mut self, ordering: Option>) -> napi::Result<()> { let ordering = ordering.map(|ordering| { @@ -374,6 +379,11 @@ impl VectorQuery { self.inner = self.inner.clone().with_row_id(); } + #[napi] + pub fn use_lsm(&mut self, enable: bool) { + self.inner = self.inner.clone().use_lsm(enable); + } + #[napi] pub fn rerank( &mut self, @@ -479,6 +489,11 @@ impl TakeQuery { self.inner = self.inner.clone().with_row_id(); } + #[napi] + pub fn use_lsm(&mut self, enable: bool) { + self.inner = self.inner.clone().use_lsm(enable); + } + #[napi(catch_unwind)] pub async fn output_schema(&self) -> napi::Result { let schema = self.inner.output_schema().await.default_error()?; diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index c922aa69b..9c000a831 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -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] diff --git a/python/python/lancedb/merge.py b/python/python/lancedb/merge.py index 386c592cd..3debed97b 100644 --- a/python/python/lancedb/merge.py +++ b/python/python/lancedb/merge.py @@ -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( diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 9a730e35f..7063f08cc 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -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. diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 5f89f6308..f61a1f1cf 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -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, ), ) diff --git a/python/python/tests/test_merge_insert_lsm.py b/python/python/tests/test_merge_insert_lsm.py index abdfb306d..5674a05ab 100644 --- a/python/python/tests/test_merge_insert_lsm.py +++ b/python/python/tests/test_merge_insert_lsm.py @@ -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)) diff --git a/python/src/query.rs b/python/src/query.rs index 70ef3b281..affbdf4fd 100644 --- a/python/src/query.rs +++ b/python/src/query.rs @@ -294,6 +294,7 @@ pub struct PyQueryRequest { pub select: PySelect, pub fast_search: Option, pub with_row_id: Option, + pub use_lsm: Option, pub column: Option, pub query_vector: Option, pub minimum_nprobes: Option, @@ -324,6 +325,7 @@ impl From for PyQueryRequest { select: PySelect(query_request.select), fast_search: Some(query_request.fast_search), with_row_id: Some(query_request.with_row_id), + use_lsm: query_request.use_lsm, column: None, query_vector: None, minimum_nprobes: None, @@ -348,6 +350,7 @@ impl From for PyQueryRequest { select: PySelect(vector_query.base.select), fast_search: Some(vector_query.base.fast_search), with_row_id: Some(vector_query.base.with_row_id), + use_lsm: vector_query.base.use_lsm, column: vector_query.column, query_vector: Some(PyQueryVectors(vector_query.query_vector)), minimum_nprobes: Some(vector_query.minimum_nprobes), @@ -474,6 +477,10 @@ impl Query { self.inner = self.inner.clone().fast_search(); } + pub fn use_lsm(&mut self, enable: bool) { + self.inner = self.inner.clone().use_lsm(enable); + } + pub fn with_row_id(&mut self) { self.inner = self.inner.clone().with_row_id(); } @@ -636,6 +643,10 @@ impl TakeQuery { self.inner = self.inner.clone().with_row_id(); } + pub fn use_lsm(&mut self, enable: bool) { + self.inner = self.inner.clone().use_lsm(enable); + } + #[pyo3(signature = ())] pub fn output_schema(self_: PyRef<'_, Self>) -> PyResult> { let inner = self_.inner.clone(); @@ -745,6 +756,10 @@ impl FTSQuery { self.inner = self.inner.clone().fast_search(); } + pub fn use_lsm(&mut self, enable: bool) { + self.inner = self.inner.clone().use_lsm(enable); + } + pub fn with_row_id(&mut self) { self.inner = self.inner.clone().with_row_id(); } @@ -892,6 +907,10 @@ impl VectorQuery { self.inner = self.inner.clone().fast_search(); } + pub fn use_lsm(&mut self, enable: bool) { + self.inner = self.inner.clone().use_lsm(enable); + } + pub fn with_row_id(&mut self) { self.inner = self.inner.clone().with_row_id(); } @@ -1086,6 +1105,11 @@ impl HybridQuery { self.inner_fts.postfilter(); } + pub fn use_lsm(&mut self, enable: bool) { + self.inner_vec.use_lsm(enable); + self.inner_fts.use_lsm(enable); + } + pub fn add_query_vector(&mut self, vector: Bound<'_, PyAny>) -> PyResult<()> { self.inner_vec.add_query_vector(vector) } diff --git a/python/src/table.rs b/python/src/table.rs index 5d6910685..b62aa7377 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1215,8 +1215,8 @@ impl Table { if let Some(use_index) = parameters.use_index { builder.use_index(use_index); } - if let Some(use_lsm_write) = parameters.use_lsm_write { - builder.use_lsm_write(use_lsm_write); + if let Some(use_lsm) = parameters.use_lsm { + builder.use_lsm(use_lsm); } if let Some(validate_single_shard) = parameters.validate_single_shard { builder.validate_single_shard(validate_single_shard); @@ -1454,7 +1454,7 @@ pub struct MergeInsertParams { when_not_matched_by_source_condition_expr: Option, timeout: Option, use_index: Option, - use_lsm_write: Option, + use_lsm: Option, validate_single_shard: Option, } diff --git a/rust/lancedb/src/lib.rs b/rust/lancedb/src/lib.rs index 0b2962682..35a564d41 100644 --- a/rust/lancedb/src/lib.rs +++ b/rust/lancedb/src/lib.rs @@ -167,6 +167,11 @@ //! # } //! ``` +// The MemWAL LSM read path (`table::query::lsm`) deepens the `create_plan` future's +// type graph enough to overflow the default trait-recursion limit while evaluating +// auto-traits (`Send`) through the Linux io_uring build's moka cache. Raise it. +#![recursion_limit = "256"] + pub mod arrow; pub mod blob; pub mod connection; diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index 727d87c38..b76865043 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -523,6 +523,26 @@ pub trait QueryBase { /// /// This allows ordering query results by one or more columns in either ascending or descending order. fn order_by(self, ordering: Option>) -> Self; + + /// Control MemWAL read routing for this query. + /// + /// By default (unset), when the table carries a MemWAL write spec (see + /// [`crate::Table::set_lsm_write_spec`]), reads are routed through the LSM + /// scanner so they also return data written via the `merge_insert` LSM path + /// that has not yet been compacted into the base table (active/frozen + /// memtables and flushed generations); a table without a spec reads the base + /// table. + /// + /// - `use_lsm(true)` forces LSM routing and errors if the table has no + /// MemWAL write spec. + /// - `use_lsm(false)` bypasses the MemWAL and reads the base table only, + /// even when a spec is present. + /// + /// Note: the LSM scanner does not support every query shape (e.g. reranking, + /// hybrid search, `order_by`). On a MemWAL table those shapes error unless + /// `use_lsm(false)` is set, because a base-only read would silently + /// exclude un-compacted MemWAL data. + fn use_lsm(self, enable: bool) -> Self; } pub trait HasQuery { @@ -593,6 +613,11 @@ impl QueryBase for T { self.mut_query().order_by = ordering; self } + + fn use_lsm(mut self, enable: bool) -> Self { + self.mut_query().use_lsm = Some(enable); + self + } } /// Options for controlling the execution of a query @@ -844,6 +869,20 @@ pub struct QueryRequest { /// /// This allows ordering query results by one or more columns in either ascending or descending order. pub order_by: Option>, + + /// Controls MemWAL read routing. When unset (the default), a query against a + /// table that carries a MemWAL write spec (see + /// [`crate::Table::set_lsm_write_spec`]) is routed through the LSM scanner so + /// it also sees data written via the `merge_insert` LSM path that has not yet + /// been compacted into the base table — the active and frozen in-memory + /// memtables and the flushed (L0) generations, deduplicated by primary key + /// against the base table (newest generation wins); a table without a spec + /// reads the base table. + /// + /// - `Some(true)` forces LSM routing and errors if the table has no MemWAL + /// write spec. + /// - `Some(false)` reads only the base table, bypassing the MemWAL. + pub use_lsm: Option, } impl Default for QueryRequest { @@ -862,6 +901,7 @@ impl Default for QueryRequest { norm: None, disable_scoring_autoprojection: false, order_by: None, + use_lsm: None, } } } diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 92c4ad296..b17e6c2bb 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -617,6 +617,11 @@ impl RemoteTable { ) -> Result<()> { params.check_filter()?; body["prefilter"] = params.prefilter.into(); + // Only forward use_lsm when explicitly set; a server that predates it + // ignores the field and routes as it would by default. + if let Some(use_lsm) = params.use_lsm { + body["use_lsm"] = serde_json::Value::Bool(use_lsm); + } if let Some(offset) = params.offset { body["offset"] = serde_json::Value::Number(serde_json::Number::from(offset)); } @@ -2843,6 +2848,10 @@ struct MergeInsertRequest { // (the default is true) #[serde(skip_serializing_if = "is_true")] use_index: bool, + // Only serialize use_lsm when explicitly set (Some); a server that predates + // it ignores the field and routes as it would by default. + #[serde(skip_serializing_if = "Option::is_none")] + use_lsm: Option, } fn is_true(b: &bool) -> bool { @@ -2894,6 +2903,7 @@ impl TryFrom for MergeInsertRequest { when_not_matched_by_source_delete_filt, // Only serialize use_index when it's false for backwards compatibility use_index: value.use_index, + use_lsm: value.use_lsm, }) } } diff --git a/rust/lancedb/src/table/merge.rs b/rust/lancedb/src/table/merge.rs index a122ba2f2..13a633c67 100644 --- a/rust/lancedb/src/table/merge.rs +++ b/rust/lancedb/src/table/merge.rs @@ -73,7 +73,7 @@ pub struct MergeInsertBuilder { pub(crate) when_not_matched_by_source_delete_filt: Option, pub(crate) timeout: Option, pub(crate) use_index: bool, - pub(crate) use_lsm_write: Option, + pub(crate) use_lsm: Option, pub(crate) validate_single_shard: bool, } @@ -89,7 +89,7 @@ impl MergeInsertBuilder { when_not_matched_by_source_delete_filt: None, timeout: None, use_index: true, - use_lsm_write: None, + use_lsm: None, validate_single_shard: true, } } @@ -187,16 +187,17 @@ impl MergeInsertBuilder { self } - /// Controls whether `merge_insert` uses the MemWAL LSM write path. + /// Control MemWAL routing for this `merge_insert`. /// /// By default (unset), a `merge_insert` on a table with an - /// [`LsmWriteSpec`](super::LsmWriteSpec) installed is routed through - /// Lance's MemWAL shard writer, and a table without one uses the standard - /// path. Calling this with `false` forces the standard path even when a - /// spec is set. Calling it with `true` requires a spec — `merge_insert` - /// errors if none is installed. - pub fn use_lsm_write(&mut self, use_lsm_write: bool) -> &mut Self { - self.use_lsm_write = Some(use_lsm_write); + /// [`LsmWriteSpec`](super::LsmWriteSpec) installed is routed through Lance's + /// MemWAL shard writer; a table without one uses the standard path. + /// + /// - `use_lsm(true)` forces MemWAL routing and errors if the table has no + /// LSM write spec. + /// - `use_lsm(false)` forces the standard write path even when a spec is set. + pub fn use_lsm(&mut self, enable: bool) -> &mut Self { + self.use_lsm = Some(enable); self } @@ -626,7 +627,7 @@ mod lsm_tests { } #[tokio::test] - async fn lsm_merge_insert_use_lsm_write_false_falls_back() { + async fn lsm_merge_insert_use_lsm_false_falls_back() { let dir = tempdir().unwrap(); let table = id_value_table(&dir).await; table @@ -634,9 +635,10 @@ mod lsm_tests { .await .unwrap(); - // use_lsm_write(false) opts out: the standard path runs and commits. + // use_lsm(false) opts out: the standard path runs and commits even though + // a spec is installed. let mut builder = table.merge_insert(&["id"]); - builder.when_not_matched_insert_all().use_lsm_write(false); + builder.when_not_matched_insert_all().use_lsm(false); let result = builder .execute(id_value_reader(vec![3, 4, 5])) .await @@ -646,6 +648,25 @@ mod lsm_tests { assert_eq!(table.count_rows(None).await.unwrap(), 5); } + #[tokio::test] + async fn lsm_merge_insert_use_lsm_true_without_spec_errors() { + let dir = tempdir().unwrap(); + let table = id_value_table(&dir).await; + + // use_lsm(true) demands MemWAL routing; without a write spec it errors + // rather than silently falling back to the standard path. + let mut builder = table.merge_insert(&["id"]); + builder + .when_matched_update_all(None) + .when_not_matched_insert_all() + .use_lsm(true); + let err = builder + .execute(id_value_reader(vec![3, 4, 5])) + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}"); + } + #[tokio::test] async fn lsm_merge_insert_rejects_on_not_primary_key() { let dir = tempdir().unwrap(); @@ -754,18 +775,23 @@ mod lsm_tests { } #[tokio::test] - async fn lsm_merge_insert_use_lsm_write_true_requires_spec() { + async fn lsm_merge_insert_no_spec_uses_standard_path() { let dir = tempdir().unwrap(); // id_value_table sets a primary key but no LSM write spec. let table = id_value_table(&dir).await; + // Without a spec, a default merge_insert (use_lsm unset) simply uses + // the standard path and commits — no opt-out required, no error. let mut builder = table.merge_insert(&["id"]); builder .when_matched_update_all(None) - .when_not_matched_insert_all() - .use_lsm_write(true); - let err = builder.execute(id_value_reader(vec![4])).await.unwrap_err(); - assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}"); + .when_not_matched_insert_all(); + let result = builder + .execute(id_value_reader(vec![3, 4, 5])) + .await + .unwrap(); + assert_eq!(result.num_inserted_rows, 2); + assert_eq!(table.count_rows(None).await.unwrap(), 5); } #[tokio::test] @@ -818,4 +844,387 @@ mod lsm_tests { .await .unwrap(); } + + // --------------------------------------------------------------------- + // LSM read path + // --------------------------------------------------------------------- + + use crate::arrow::SendableRecordBatchStream; + use crate::query::{ExecutableQuery, QueryBase}; + use arrow::array::AsArray; + use arrow::datatypes::Int64Type; + use futures::TryStreamExt; + + /// Collect `(id, value)` pairs from a result stream, sorted by id. + async fn collect_id_value(stream: SendableRecordBatchStream) -> Vec<(i64, i64)> { + let batches: Vec<_> = stream.try_collect().await.unwrap(); + let mut rows = Vec::new(); + for batch in &batches { + let ids = batch + .column_by_name("id") + .unwrap() + .as_primitive::(); + let values = batch + .column_by_name("value") + .unwrap() + .as_primitive::(); + for i in 0..batch.num_rows() { + rows.push((ids.value(i), values.value(i))); + } + } + rows.sort(); + rows + } + + /// Upsert `ids` (value = 0..n) through the LSM `merge_insert` path. + async fn lsm_upsert(table: &Table, ids: Vec) { + let mut builder = table.merge_insert(&[]); + builder + .when_matched_update_all(None) + .when_not_matched_insert_all(); + builder.execute(id_value_reader(ids)).await.unwrap(); + } + + #[tokio::test] + async fn lsm_read_sees_active_memtable() { + let dir = tempdir().unwrap(); + let table = id_value_table(&dir).await; // base: ids 1,2,3 (value 0,1,2) + table + .set_lsm_write_spec(LsmWriteSpec::unsharded()) + .await + .unwrap(); + + // Insert ids 4,5 into the active memtable (not committed to base). + lsm_upsert(&table, vec![4, 5]).await; + + // Default read auto-routes through the LSM scanner: base ∪ active memtable. + let lsm = table.query().execute().await.unwrap(); + let rows = collect_id_value(lsm).await; + assert_eq!( + rows.iter().map(|(id, _)| *id).collect::>(), + vec![1, 2, 3, 4, 5] + ); + + // use_lsm(false) bypasses the MemWAL and reads the base table only. + let base_only = table.query().use_lsm(false).execute().await.unwrap(); + let rows = collect_id_value(base_only).await; + assert_eq!( + rows.iter().map(|(id, _)| *id).collect::>(), + vec![1, 2, 3] + ); + } + + #[tokio::test] + async fn lsm_read_dedup_newest_wins() { + let dir = tempdir().unwrap(); + let table = id_value_table(&dir).await; // base: id 2 -> value 1 + table + .set_lsm_write_spec(LsmWriteSpec::unsharded()) + .await + .unwrap(); + + // Upsert ids 2,3,4 with values 0,1,2. id 2 and 3 shadow the base rows. + lsm_upsert(&table, vec![2, 3, 4]).await; + + let lsm = table.query().execute().await.unwrap(); + let rows = collect_id_value(lsm).await; + // id 1 from base (value 0); ids 2,3,4 from memtable (values 0,1,2). + assert_eq!(rows, vec![(1, 0), (2, 0), (3, 1), (4, 2)]); + } + + #[tokio::test] + async fn lsm_read_point_lookup_filter() { + let dir = tempdir().unwrap(); + let table = id_value_table(&dir).await; + table + .set_lsm_write_spec(LsmWriteSpec::unsharded()) + .await + .unwrap(); + lsm_upsert(&table, vec![2, 3, 4]).await; // id 2 -> value 0 (shadows base) + + let lsm = table.query().only_if("id = 2").execute().await.unwrap(); + let rows = collect_id_value(lsm).await; + assert_eq!(rows, vec![(2, 0)]); + } + + #[tokio::test] + async fn lsm_read_multi_shard() { + let dir = tempdir().unwrap(); + let table = id_value_table(&dir).await; + table + .set_lsm_write_spec(LsmWriteSpec::bucket("id", 8)) + .await + .unwrap(); + + // Two single-row upserts that route to (likely) different buckets; each + // closes the writer so the next opens a fresh shard. + lsm_upsert(&table, vec![10]).await; + table.close_lsm_writers().await.unwrap(); + lsm_upsert(&table, vec![11]).await; + + let lsm = table.query().execute().await.unwrap(); + let rows = collect_id_value(lsm).await; + let ids: Vec = rows.iter().map(|(id, _)| *id).collect(); + // Base 1,2,3 + flushed/active shards for 10 and 11. + assert_eq!(ids, vec![1, 2, 3, 10, 11]); + } + + #[tokio::test] + async fn lsm_read_after_close_sees_flushed() { + let dir = tempdir().unwrap(); + let table = id_value_table(&dir).await; + table + .set_lsm_write_spec(LsmWriteSpec::unsharded()) + .await + .unwrap(); + lsm_upsert(&table, vec![4, 5]).await; + // close flushes the active memtable to an on-disk generation and drops + // the cached writer; the read must still see those rows via the shard + // manifest snapshot. + table.close_lsm_writers().await.unwrap(); + + let lsm = table.query().execute().await.unwrap(); + let ids: Vec = collect_id_value(lsm) + .await + .iter() + .map(|(id, _)| *id) + .collect(); + assert_eq!(ids, vec![1, 2, 3, 4, 5]); + } + + #[tokio::test] + async fn lsm_read_without_spec_reads_base() { + let dir = tempdir().unwrap(); + let table = id_value_table(&dir).await; // no LSM write spec + + // With no spec installed there is nothing to route: the default read and + // an explicit use_lsm(false) both read the base table without error. + for query in [table.query(), table.query().use_lsm(false)] { + let rows = collect_id_value(query.execute().await.unwrap()).await; + assert_eq!( + rows.iter().map(|(id, _)| *id).collect::>(), + vec![1, 2, 3] + ); + } + } + + #[tokio::test] + async fn lsm_read_unsupported_shape_errors_without_use_lsm_false() { + let dir = tempdir().unwrap(); + let table = id_value_table(&dir).await; + table + .set_lsm_write_spec(LsmWriteSpec::unsharded()) + .await + .unwrap(); + lsm_upsert(&table, vec![4]).await; + + // `with_row_id` is a shape the LSM scanner cannot honor. On a MemWAL + // table the default (auto-routed) read hard-errors rather than silently + // reading a stale base-only result that would exclude un-compacted row 4. + let err = table + .query() + .with_row_id() + .execute() + .await + .err() + .expect("unsupported shape on a MemWAL table must error"); + assert!(matches!(err, Error::NotSupported { .. }), "got {err:?}"); + + // use_lsm(false) is the escape hatch: it reads the base table only. + let rows = collect_id_value( + table + .query() + .with_row_id() + .use_lsm(false) + .execute() + .await + .unwrap(), + ) + .await; + assert_eq!( + rows.iter().map(|(id, _)| *id).collect::>(), + vec![1, 2, 3] + ); + } + + /// A reader of `[id: Int64, text: Utf8]` rows. + fn id_text_reader(rows: Vec<(i64, &str)>) -> Box { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("text", DataType::Utf8, false), + ])); + let ids: Vec = rows.iter().map(|(id, _)| *id).collect(); + let texts: Vec<&str> = rows.iter().map(|(_, t)| *t).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(ids)), + Arc::new(StringArray::from(texts)), + ], + ) + .unwrap(); + Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema)) + } + + #[tokio::test] + async fn lsm_read_full_text_search() { + use crate::index::Index; + use lance_index::scalar::FullTextSearchQuery; + + let dir = tempdir().unwrap(); + let conn = connect(dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let table = conn + .create_table( + "t", + id_text_reader(vec![(1, "alpha"), (2, "beta"), (3, "gamma")]), + ) + .execute() + .await + .unwrap(); + table.set_unenforced_primary_key(["id"]).await.unwrap(); + table + .create_index(&["text"], Index::FTS(Default::default())) + .execute() + .await + .unwrap(); + let fts_index = table.list_indices().await.unwrap()[0].name.clone(); + table + .set_lsm_write_spec(LsmWriteSpec::unsharded().with_maintained_indexes([fts_index])) + .await + .unwrap(); + + // Insert a row whose term ("zebra") exists in no base row. + let mut builder = table.merge_insert(&[]); + builder + .when_matched_update_all(None) + .when_not_matched_insert_all(); + builder + .execute(id_text_reader(vec![(99, "zebra")])) + .await + .unwrap(); + + let search = |term: &str| { + let q = FullTextSearchQuery::new(term.to_string()) + .with_column("text".to_string()) + .unwrap(); + table.query().full_text_search(q) + }; + + // "zebra" lives only in the active memtable; LSM read finds it. + let stream = search("zebra").execute().await.unwrap(); + let batches: Vec<_> = stream.try_collect().await.unwrap(); + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 1, "LSM FTS must surface the memtable row"); + + // A base-only term still matches the base table through the LSM scan. + let stream = search("alpha").execute().await.unwrap(); + let batches: Vec<_> = stream.try_collect().await.unwrap(); + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 1, "LSM FTS must still see base rows"); + } + + #[tokio::test] + async fn lsm_read_vector_search() { + use crate::index::Index; + use crate::index::vector::IvfPqIndexBuilder; + use arrow::array::{FixedSizeListBuilder, Float32Builder}; + use arrow::datatypes::Int64Type; + + const DIM: i32 = 8; + const N: i64 = 256; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new( + "vec", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), DIM), + false, + ), + ])); + let make_batch = |rows: Vec<(i64, f32)>| -> RecordBatch { + let ids: Vec = rows.iter().map(|(id, _)| *id).collect(); + let mut vb = FixedSizeListBuilder::new(Float32Builder::new(), DIM); + for (_, fill) in &rows { + for _ in 0..DIM { + vb.values().append_value(*fill); + } + vb.append(true); + } + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from(ids)), Arc::new(vb.finish())], + ) + .unwrap() + }; + + let dir = tempdir().unwrap(); + let conn = connect(dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + // Base rows fill each vector with its own id (0..256); all far from 1000. + let base = make_batch((0..N).map(|i| (i, i as f32)).collect()); + let base_reader: Box = + Box::new(RecordBatchIterator::new(vec![Ok(base)], schema.clone())); + let table = conn.create_table("t", base_reader).execute().await.unwrap(); + table.set_unenforced_primary_key(["id"]).await.unwrap(); + table + .create_index( + &["vec"], + Index::IvfPq( + IvfPqIndexBuilder::default() + .num_partitions(1) + .num_sub_vectors(2), + ), + ) + .execute() + .await + .unwrap(); + let vec_index = table.list_indices().await.unwrap()[0].name.clone(); + table + .set_lsm_write_spec(LsmWriteSpec::unsharded().with_maintained_indexes([vec_index])) + .await + .unwrap(); + + // Insert a vector (filled with 1000) that is nearest to the query. + let mut builder = table.merge_insert(&[]); + builder + .when_matched_update_all(None) + .when_not_matched_insert_all(); + let insert_reader: Box = Box::new(RecordBatchIterator::new( + vec![Ok(make_batch(vec![(9999, 1000.0)]))], + schema.clone(), + )); + builder.execute(insert_reader).await.unwrap(); + + // KNN near [1000; DIM]: the default (auto-routed) read surfaces the + // memtable row. + let stream = table + .query() + .nearest_to(&[1000.0_f32; 8]) + .unwrap() + .limit(1) + .execute() + .await + .unwrap(); + let batches: Vec<_> = stream.try_collect().await.unwrap(); + let ids: Vec = batches + .iter() + .flat_map(|b| { + b.column_by_name("id") + .unwrap() + .as_primitive::() + .values() + .to_vec() + }) + .collect(); + assert_eq!( + ids, + vec![9999], + "LSM vector search must rank the memtable row first" + ); + } } diff --git a/rust/lancedb/src/table/merge/lsm.rs b/rust/lancedb/src/table/merge/lsm.rs index 9bc6ad088..0eb7c0231 100644 --- a/rust/lancedb/src/table/merge/lsm.rs +++ b/rust/lancedb/src/table/merge/lsm.rs @@ -306,6 +306,39 @@ impl ShardWriterEntry { } Ok(()) } + + /// The cached writer's latest in-memory manifest (current generation + + /// flushed generations). `Ok(None)` if the writer was already closed. + /// Used by the LSM read path to snapshot this shard authoritatively + /// without re-reading the on-disk manifest. + async fn manifest(&self) -> Result> { + let guard = self.inner.read().await; + let Some(writer) = guard.as_ref() else { + return Ok(None); + }; + writer.manifest().await.map_err(|e| Error::Runtime { + message: format!("read: shard writer manifest read failed: {}", e), + }) + } + + /// Atomically capture the cached writer's active + frozen-awaiting-flush + /// memtables for unified LSM scanning. `Ok(None)` if the writer was + /// already closed. + async fn in_memory_memtable_refs( + &self, + ) -> Result> { + let guard = self.inner.read().await; + let Some(writer) = guard.as_ref() else { + return Ok(None); + }; + writer + .in_memory_memtable_refs() + .await + .map(Some) + .map_err(|e| Error::Runtime { + message: format!("read: shard writer memtable capture failed: {}", e), + }) + } } impl ShardWriterCache { @@ -345,6 +378,36 @@ impl ShardWriterCache { Ok(entry) } + /// Snapshot the cached writer's shard for the LSM read path: its shard id, + /// authoritative in-memory manifest, and active + frozen memtable refs. + /// Returns `None` when no writer is currently cached (e.g. nothing has been + /// written this session, or the writer was closed). + #[allow(clippy::redundant_pub_crate)] + pub(crate) async fn read_snapshot( + &self, + ) -> Result< + Option<( + Uuid, + Option, + Option, + )>, + > { + let cached = { + let guard = self.slot.read().await; + guard.as_ref().map(|(id, entry)| (*id, entry.clone())) + }; + let Some((shard_id, entry)) = cached else { + return Ok(None); + }; + // Capture memtables before the manifest. If a flush interleaves, dedup + // tolerates the same rows appearing in both a memtable and a freshly + // flushed generation, but would drop rows present in neither. Manifest + // last guarantees any generation flushed mid-capture is still covered. + let memtables = entry.in_memory_memtable_refs().await?; + let manifest = entry.manifest().await?; + Ok(Some((shard_id, manifest, memtables))) + } + /// Close the cached writer, if any, and clear the slot. #[allow(clippy::redundant_pub_crate)] pub(crate) async fn drain_and_close(&self) -> Result<()> { @@ -408,19 +471,20 @@ pub(crate) async fn lsm_dispatch_decision( table: &NativeTable, params: &MergeInsertBuilder, ) -> Result { - // `Some(false)` is an explicit opt-out: use the standard path. - if params.use_lsm_write == Some(false) { + // Explicit opt-out: use the standard path regardless of any installed spec. + if params.use_lsm == Some(false) { return Ok(LsmDispatch::Standard); } let dataset = table.dataset.get().await?; let Some(details) = dataset.mem_wal_index_details().await? else { - // No LSM write spec installed. `Some(true)` explicitly asked for the - // LSM path, which is meaningless without a spec; `None` (the default) - // just falls back to the standard path. - if params.use_lsm_write == Some(true) { + // No write spec installed. `use_lsm(true)` demanded MemWAL routing, so + // that is an error; otherwise fall back to the standard path. + if params.use_lsm == Some(true) { return Err(Error::InvalidInput { - message: "merge_insert: use_lsm_write(true) requires an LSM write spec on the table; call set_lsm_write_spec first".to_string(), + message: "use_lsm(true) was set but the table has no MemWAL write spec; \ + install one with set_lsm_write_spec or leave use_lsm unset" + .to_string(), }); } return Ok(LsmDispatch::Standard); @@ -449,7 +513,7 @@ pub(crate) async fn lsm_dispatch_decision( if !is_upsert_only(params) { return Err(Error::InvalidInput { - message: "merge_insert: when an LSM write spec is set, only the upsert form (when_matched_update_all without a filter + when_not_matched_insert_all, no by-source delete) is supported; call use_lsm_write(false) to use the standard merge_insert path".to_string(), + message: "merge_insert: when an LSM write spec is set, only the upsert form (when_matched_update_all without a filter + when_not_matched_insert_all, no by-source delete) is supported; call use_lsm(false) to use the standard merge_insert path".to_string(), }); } diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index 28845e96f..9feb9d5ab 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -3,6 +3,8 @@ use std::sync::Arc; +mod lsm; + use super::NativeTable; use crate::connection::NamespaceClientPushdownOperation; use crate::error::{Error, Result}; @@ -20,6 +22,7 @@ use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::union::UnionExec; use futures::future::try_join_all; +use lance::dataset::mem_wal::DatasetMemWalExt; use lance::dataset::scanner::DatasetRecordBatchStream; use lance::dataset::scanner::Scanner; use lance_datafusion::exec::{analyze_plan as lance_analyze_plan, execute_plan}; @@ -53,7 +56,7 @@ pub async fn execute_query( // QueryTable pushdown runs the query server-side, but only on the main // branch: the namespace request carries no branch yet, so a branch handle // must fall through to local execution. - if can_execute_namespace_query(table, query) + if can_execute_namespace_query(table, query).await? && let Some(ref namespace_client) = table.namespace_client { return execute_namespace_query(table, namespace_client.clone(), query, options).await; @@ -61,18 +64,35 @@ pub async fn execute_query( execute_generic_query(table, query, options).await } -fn can_execute_namespace_query(table: &NativeTable, query: &AnyQuery) -> bool { - table +async fn can_execute_namespace_query(table: &NativeTable, query: &AnyQuery) -> Result { + if !(table .pushdown_operations .contains(&NamespaceClientPushdownOperation::QueryTable) && table.namespace_client.is_some() && table.dataset.current_branch().is_none() - && !requires_local_namespace_execution(query) + && !requires_local_namespace_execution(query)) + { + return Ok(false); + } + // A MemWAL write spec means reads auto-route through the LSM scanner in + // `create_plan` even when `use_lsm` is unset. The namespace request has no + // use_lsm field, so pushing the default query down would silently omit + // un-compacted rows — force local execution whenever a spec is installed. + let dataset = table.dataset.get().await?; + if dataset.mem_wal_index_details().await?.is_some() { + return Ok(false); + } + Ok(true) } fn requires_local_namespace_execution(query: &AnyQuery) -> bool { - // The namespace QueryTable request has no approx_mode field yet, so - // pushing this query down would silently ignore the user's setting. + // The namespace QueryTable request has no approx_mode or use_lsm field yet, so + // pushing these down would silently ignore the user's setting. For use_lsm that + // is worse than a tuning miss: MemWAL read routing lives only in `create_plan`, + // so a pushed-down query would return stale base-only data with no error. + if query.base().use_lsm.is_some() { + return true; + } matches!( query, AnyQuery::VectorQuery(VectorQueryRequest { @@ -120,6 +140,29 @@ pub async fn create_plan( query.base.check_filter()?; let ds_ref = table.dataset.get().await?; + + // MemWAL read routing driven by `use_lsm`: + // * unset — route through the LSM scanner iff the table carries a write spec + // * Some(true) — force LSM routing; error if the table has no write spec + // * Some(false) — read the base table only, bypassing the MemWAL + // The LSM scanner surfaces in-flight `merge_insert` data (active/frozen + // memtables + flushed generations); validation and dispatch live in `lsm`. + let has_spec = ds_ref.mem_wal_index_details().await?.is_some(); + let use_lsm = match query.base.use_lsm { + Some(true) if !has_spec => { + return Err(Error::InvalidInput { + message: "use_lsm(true) was set but the table has no MemWAL write spec; \ + install one with set_lsm_write_spec or leave use_lsm unset" + .to_string(), + }); + } + Some(enable) => enable, + None => has_spec, + }; + if use_lsm { + return lsm::create_lsm_plan(table, ds_ref, query).await; + } + let schema = ds_ref.schema(); let mut column = query.column.clone(); @@ -904,6 +947,65 @@ mod tests { assert_eq!(namespace_client.query_table_calls.load(Ordering::SeqCst), 0); } + #[tokio::test] + async fn test_execute_query_use_lsm_with_namespace_pushdown_runs_locally() { + use crate::connect; + use crate::table::query::execute_query; + use arrow_array::{Int32Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + + let conn = connect("memory://").execute().await.unwrap(); + + let vectors = Arc::new(fixed_size_list_array( + vec![0.0, 0.0, 10.0, 10.0, 20.0, 20.0], + 2, + )); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(Int32Array::from(vec![1, 2, 3])), vectors], + ) + .unwrap(); + + let table = conn + .create_table("test_use_lsm_namespace_fallback", batch) + .execute() + .await + .unwrap(); + let namespace_client = Arc::new(CountingNamespaceClient::default()); + let mut native_table = table.as_native().unwrap().clone(); + native_table.namespace_client = Some(namespace_client.clone()); + native_table + .pushdown_operations + .insert(NamespaceClientPushdownOperation::QueryTable); + + // `use_lsm` set (even to false) must force local execution — the namespace + // request has no use_lsm field, so a pushdown would silently ignore it. + let query_vector = Arc::new(Float32Array::from(vec![0.0, 0.0])); + let query = AnyQuery::VectorQuery(VectorQueryRequest { + base: QueryRequest { + limit: Some(1), + use_lsm: Some(false), + ..Default::default() + }, + column: Some("vector".to_string()), + query_vector: vec![query_vector as ArrayRef], + ..Default::default() + }); + + let stream = execute_query(&native_table, &query, QueryExecutionOptions::default()) + .await + .unwrap(); + let batches = stream.try_collect::>().await.unwrap(); + let count: usize = batches.iter().map(|b| b.num_rows()).sum(); + + assert_eq!(count, 1); + assert_eq!(namespace_client.query_table_calls.load(Ordering::SeqCst), 0); + } + #[tokio::test] async fn test_create_plan_multivector_structure() { use arrow_array::{Float32Array, RecordBatch}; diff --git a/rust/lancedb/src/table/query/lsm.rs b/rust/lancedb/src/table/query/lsm.rs new file mode 100644 index 000000000..074d13476 --- /dev/null +++ b/rust/lancedb/src/table/query/lsm.rs @@ -0,0 +1,786 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! MemWAL LSM read path. +//! +//! When a table has an LSM write spec installed (see [`set_lsm_write_spec`]), +//! reads are routed through Lance's [`LsmScanner`] instead of the plain +//! base-table scan unless the query sets +//! [`use_lsm(false)`](crate::query::QueryBase::use_lsm). This makes data +//! written via the LSM `merge_insert` path — which lives in the active/frozen +//! in-memory memtables and the flushed SSTable generations until an external +//! compaction merges it into the base table — visible to queries, deduplicated by +//! primary key (newest generation wins). +//! +//! Three query shapes are supported, mirroring the standard scan: a plain scan +//! (filter / projection / limit), full-text search, and vector (ANN) search. All +//! three run through a single [`LsmScanner`], so a `where` filter is honored as a +//! prefilter uniformly — including for vector search, where `LsmScanner` threads +//! it into the vector planner's prefilter. Shapes the LSM path cannot honor are +//! rejected with [`Error::NotSupported`]; the caller must set `use_lsm(false)` to +//! run those against the base table. +//! +//! [`set_lsm_write_spec`]: crate::Table::set_lsm_write_spec + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow_array::Array; +use arrow_schema::{DataType, Schema as ArrowSchema}; +use datafusion_physical_plan::expressions::Column; +use datafusion_physical_plan::projection::ProjectionExec; +use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr}; +use lance::Dataset; +use lance::dataset::mem_wal::scanner::InMemoryMemTables; +use lance::dataset::mem_wal::{ + DatasetMemWalExt, LsmScanner, ShardManifestStore, ShardSnapshot, ShardWriterConfig, +}; +use lance_index::mem_wal::{MemWalIndexDetails, ShardManifest}; +use uuid::Uuid; + +use super::NativeTable; +use crate::DistanceType; +use crate::error::{Error, Result}; +use crate::query::{DEFAULT_TOP_K, QueryFilter, Select, VectorQueryRequest}; +use crate::utils::default_vector_column; + +/// Over-fetch factor for the LSM vector/FTS arms. With the default of `1.0` a +/// source blocked by cross-generation PK dedup fetches exactly `k` and can return +/// fewer than `k` live rows; upserts routinely create such blocked candidates, so +/// widen the per-source fetch to keep result pages filled. +const LSM_OVERFETCH_FACTOR: f64 = 2.0; + +/// Build the LSM read plan for a MemWAL-routed query. +/// +/// The caller guarantees `ds_ref` carries a MemWAL write spec (routing is decided +/// in [`create_plan`](super::create_plan)). Errors with [`Error::NotSupported`] +/// for query shapes the LSM scanner cannot honor — the caller must set +/// `use_lsm(false)` to run those against the base table. +pub(super) async fn create_lsm_plan( + table: &NativeTable, + ds_ref: Arc, + query: VectorQueryRequest, +) -> Result> { + reject_unsupported(&query)?; + + // A time-traveled (checked-out) handle pins an older dataset version, but the + // WAL manifests and cached writer expose current live state — mixing them would + // surface WAL rows written after the requested version. `use_lsm(false)` reads + // the base table at the pinned version. + if table.dataset.time_travel_version().is_some() { + return Err(Error::NotSupported { + message: "the MemWAL LSM scanner cannot read from a time-traveled dataset version; set use_lsm(false) to read the base table at this version".to_string(), + }); + } + + // Routing guarantees a write spec is installed (see `create_plan`). + let details = ds_ref + .mem_wal_index_details() + .await? + .ok_or_else(|| Error::Runtime { + message: "the MemWAL LSM write spec disappeared during read planning".to_string(), + })?; + + let pk_columns = pk_columns(&ds_ref)?; + // The base index an indexed arm relies on may lag compaction; resolve it so the + // snapshot retains SSTables the index has not yet caught up to. + let arm_index = arm_maintained_index_name(&ds_ref, &query, &details).await?; + let (snapshots, in_memory) = + build_read_context(table, &ds_ref, &details, arm_index.as_deref()).await?; + + let limit = query.base.limit; + let offset = query.base.offset; + + let plan = if !query.query_vector.is_empty() { + vector_plan( + &ds_ref, + &query, + &details, + pk_columns.clone(), + snapshots, + in_memory, + limit, + offset, + ) + .await? + } else if let Some(fts) = &query.base.full_text_search { + fts_plan( + &ds_ref, + fts.clone(), + &query, + &details, + pk_columns.clone(), + snapshots, + in_memory, + limit, + offset, + ) + .await? + } else { + plain_plan( + &ds_ref, + &query, + pk_columns.clone(), + snapshots, + in_memory, + limit, + offset, + ) + .await? + }; + + // Lance appends the primary-key columns internally for dedup and keeps them in + // the output; drop the ones the user did not request so the projection matches. + restore_projection(plan, &query, &pk_columns) +} + +/// Reject query shapes the LSM read path does not implement. On a MemWAL table +/// reads route through the LSM scanner by default, so an unsupported shape is a +/// hard error rather than a silent fallback to the base-only scan — which would +/// exclude un-compacted MemWAL data. The caller must set `use_lsm(false)` to +/// run these against the base table, accepting that the results omit un-compacted +/// MemWAL data. +/// +/// A `where` filter is intentionally *not* rejected: every arm routes through +/// [`LsmScanner`], which applies it as a prefilter (see [`base_scanner`]). +fn reject_unsupported(query: &VectorQueryRequest) -> Result<()> { + let unsupported = |what: &str| { + Err(Error::NotSupported { + message: format!( + "the MemWAL LSM scanner does not support {what}; set use_lsm(false) to read the base table only (results will exclude un-compacted MemWAL data)" + ), + }) + }; + if query.query_vector.len() > 1 { + return unsupported("multiple query vectors"); + } + if !query.query_vector.is_empty() && query.base.full_text_search.is_some() { + return unsupported("hybrid (vector + full-text) search"); + } + if query.base.with_row_id { + return unsupported("with_row_id (the LSM scanner exposes _rowaddr, not a stable _rowid)"); + } + if query.base.reranker.is_some() { + return unsupported("reranking / hybrid search"); + } + if query.base.order_by.is_some() { + return unsupported("order_by"); + } + // Vector-only knobs the LSM scanner cannot honor. Both change results rather + // than just recall, so error instead of silently ignoring them: distance_range + // would return rows outside the bound, and use_index(false) asks for a + // brute-force search the index-only base arm can't do. (ef / approx_mode / + // maximum_nprobes are recall/speed knobs and are left to no-op — and + // maximum_nprobes defaults to Some, so it cannot be rejected on presence.) + if !query.query_vector.is_empty() { + if query.lower_bound.is_some() || query.upper_bound.is_some() { + return unsupported("distance_range on vector search"); + } + if !query.use_index { + return unsupported( + "use_index(false) / brute-force vector search (the LSM base arm is index-only)", + ); + } + } + // Postfilter changes result semantics for both vector and full-text search, and + // the LSM scanner always prefilters — reject a requested postfilter for either. + if (!query.query_vector.is_empty() || query.base.full_text_search.is_some()) + && !query.base.prefilter + { + return unsupported( + "postfilter on vector or full-text search (the LSM scanner always prefilters)", + ); + } + match &query.base.select { + Select::All | Select::Columns(_) => {} + Select::Dynamic(_) | Select::Expr(_) => return unsupported("dynamic column projection"), + } + if let Some(QueryFilter::Substrait(_)) = &query.base.filter { + return unsupported("Substrait filters"); + } + // Take-by-row-id / row-offset queries carry a `_rowid` / `_rowoffset` filter, + // columns the LSM scanner never exposes (only `_rowaddr`); reject with guidance + // rather than failing deep in datafusion with a column-not-found error. + if let Some(QueryFilter::Datafusion(expr)) = &query.base.filter + && expr + .column_refs() + .iter() + .any(|c| c.name == "_rowid" || c.name == "_rowoffset") + { + return unsupported( + "take by row id or row offset (the LSM scanner has no stable _rowid / _rowoffset)", + ); + } + Ok(()) +} + +/// Primary-key column names from the dataset's unenforced primary key. +fn pk_columns(dataset: &Dataset) -> Result> { + let pk: Vec = dataset + .schema() + .unenforced_primary_key() + .iter() + .map(|f| f.name.clone()) + .collect(); + if pk.is_empty() { + return Err(Error::InvalidInput { + message: + "the MemWAL LSM scanner requires an unenforced primary key, but the table has none" + .to_string(), + }); + } + Ok(pk) +} + +/// Per-shard SSTable exclusion watermark: the generation at or below which SSTables +/// are safe to drop for this arm. A generation is droppable only once it is +/// compacted into the base table AND covered by `index_name`'s catch-up (for an +/// indexed arm); a plain scan (`index_name == None`) uses the compaction watermark +/// alone. Capping at the index catch-up keeps rows the base index has not yet +/// indexed visible through their SSTable. First occurrence per shard mirrors Lance's +/// `compacted_generation_for_shard`. +fn exclusion_watermarks( + details: &MemWalIndexDetails, + index_name: Option<&str>, +) -> HashMap { + let mut exclude: HashMap = HashMap::new(); + for entry in &details.compacted_sstables { + let mut watermark = entry.generation; + if let Some(name) = index_name + && let Some(caught_up) = details + .index_catchup + .iter() + .find(|icp| icp.index_name == name) + .and_then(|icp| icp.caught_up_generation_for_shard(&entry.shard_id)) + { + watermark = watermark.min(caught_up); + } + exclude.entry(entry.shard_id).or_insert(watermark); + } + exclude +} + +/// Assemble the per-shard snapshots (flushed SSTable generations) and the +/// in-memory memtables (active + frozen) for the table. +/// +/// Snapshots for all shards come from their on-disk manifests; for the shard +/// with a live cached `ShardWriter` (this session's in-flight writes) the +/// writer's authoritative in-memory manifest and memtables override the +/// on-disk view so a read sees data not yet flushed. +async fn build_read_context( + table: &NativeTable, + dataset: &Dataset, + details: &MemWalIndexDetails, + index_name: Option<&str>, +) -> Result<(Vec, HashMap)> { + let exclude = exclusion_watermarks(details, index_name); + + let shard_ids = dataset.list_mem_wal_latest_shard_ids().await?; + // Use the dataset's own object store (not `ObjectStore::from_uri`, which + // builds a fresh registry and would miss `memory://` and custom-registered + // stores). The base path matches `list_mem_wal_latest_shard_ids`. + let store = dataset.object_store(None).await?; + let base_path = dataset.branch_location().path; + let scan_batch_size = ShardWriterConfig::default().manifest_scan_batch_size; + + let mut snapshots: Vec = Vec::new(); + for shard_id in shard_ids { + let manifest_store = + ShardManifestStore::new(store.clone(), &base_path, shard_id, scan_batch_size); + if let Some(manifest) = manifest_store.read_latest().await? { + snapshots.push(snapshot_from_manifest(shard_id, &manifest, &exclude)); + } + } + + // WAL-only writers (enable_memtable=false) keep no in-memory memtable, and + // `in_memory_memtable_refs` errors in that mode; the on-disk manifests above + // already cover their flushed SSTables, so skip the live-writer snapshot. (Lance + // forbids maintained indexes in WAL-only mode, so only plain scans reach here.) + let wal_only = details + .writer_config_defaults + .get("enable_memtable") + .map(|v| v == "false") + .unwrap_or(false); + + // Override the active shard with the cached writer's in-memory view. + let mut in_memory: HashMap = HashMap::new(); + if !wal_only + && let Some((shard_id, manifest, memtables)) = + table.dataset.shard_writer().read_snapshot().await? + { + if let Some(manifest) = manifest { + let snapshot = snapshot_from_manifest(shard_id, &manifest, &exclude); + match snapshots.iter_mut().find(|s| s.shard_id == shard_id) { + Some(existing) => *existing = snapshot, + None => snapshots.push(snapshot), + } + } + if let Some(memtables) = memtables { + in_memory.insert(shard_id, memtables); + } + } + + Ok((snapshots, in_memory)) +} + +/// Convert a shard manifest into a read snapshot (current + not-yet-compacted +/// flushed SSTables). SSTable generations at or below the shard's compaction +/// watermark are already in the base table and are skipped. +fn snapshot_from_manifest( + shard_id: Uuid, + manifest: &ShardManifest, + compacted: &HashMap, +) -> ShardSnapshot { + let mut snapshot = ShardSnapshot::new(shard_id) + .with_spec_id(manifest.shard_spec_id) + .with_current_generation(manifest.current_generation); + let watermark = compacted.get(&shard_id).copied(); + for sstable in &manifest.sstables { + if watermark.is_some_and(|w| sstable.generation <= w) { + continue; + } + snapshot = snapshot.with_sstable(sstable.generation, sstable.path.clone()); + } + snapshot +} + +/// Columns selected by the query, if an explicit projection was requested. +fn selected_columns(query: &VectorQueryRequest) -> Option> { + match &query.base.select { + Select::Columns(columns) => Some(columns.clone()), + _ => None, + } +} + +/// Non-negative `Option` limit/offset as the `Option` the scanner +/// expects. +fn as_i64(value: Option) -> Option { + value.map(|v| v as i64) +} + +/// Build a base `LsmScanner` configured with sources, filter, and projection. +/// +/// The filter set here is applied as a prefilter across every arm — plain scan, +/// full-text search, and vector search — since all three terminate on this +/// scanner's `create_plan`. +fn base_scanner( + dataset: &Dataset, + query: &VectorQueryRequest, + pk_columns: Vec, + snapshots: Vec, + in_memory: HashMap, +) -> Result { + let mut scanner = LsmScanner::new(Arc::new(dataset.clone()), snapshots, pk_columns); + for (shard_id, memtables) in in_memory { + scanner = scanner.with_in_memory_memtables(shard_id, memtables); + } + if let Some(columns) = selected_columns(query) { + let refs: Vec<&str> = columns.iter().map(String::as_str).collect(); + scanner = scanner.project(&refs)?; + } + if let Some(filter) = &query.base.filter { + scanner = match filter { + QueryFilter::Sql(sql) => scanner.filter(sql)?, + QueryFilter::Datafusion(expr) => scanner.filter_expr(expr.clone()), + QueryFilter::Substrait(_) => { + return Err(Error::NotSupported { + message: "the MemWAL LSM scanner does not support Substrait filters; set use_lsm(false) to read the base table only".to_string(), + }); + } + }; + } + Ok(scanner) +} + +/// Plain scan: filter / projection / limit over base ∪ SSTables ∪ in-memory. +/// The plain scan applies limit and offset inside the planner. +async fn plain_plan( + dataset: &Dataset, + query: &VectorQueryRequest, + pk_columns: Vec, + snapshots: Vec, + in_memory: HashMap, + limit: Option, + offset: Option, +) -> Result> { + let scanner = base_scanner(dataset, query, pk_columns, snapshots, in_memory)? + .limit(as_i64(limit), as_i64(offset))?; + Ok(scanner.create_plan().await?) +} + +/// Full-text search over base ∪ SSTables ∪ in-memory, merged by local BM25 score. +/// The scanner threads the query filter in as a prefilter and pages via limit/offset. +#[allow(clippy::too_many_arguments)] +async fn fts_plan( + dataset: &Dataset, + fts: lance_index::scalar::FullTextSearchQuery, + query: &VectorQueryRequest, + details: &MemWalIndexDetails, + pk_columns: Vec, + snapshots: Vec, + in_memory: HashMap, + limit: Option, + offset: Option, +) -> Result> { + // Pre-check for a lancedb-flavored error; `LsmScanner` also validates the + // single-column requirement, but without the `use_lsm(false)` guidance. + let columns: Vec = fts.columns().into_iter().collect(); + if columns.len() > 1 { + return Err(Error::NotSupported { + message: "the MemWAL LSM scanner full-text search supports a single column; set use_lsm(false) to read the base table only".to_string(), + }); + } + let column = columns.first().ok_or_else(|| Error::NotSupported { + message: "the MemWAL LSM scanner full-text search requires an explicit FTS column" + .to_string(), + })?; + + // Without a maintained in-memory FTS index for this column, the active memtable + // arm produces an empty plan (`active_source_can_execute_fts` returns false), so + // the search silently omits un-compacted documents. Reject rather than mislead. + if !index_maintained( + dataset, + column, + &details.maintained_indexes, + "InvertedIndexDetails", + ) + .await? + { + return Err(Error::NotSupported { + message: format!( + "the MemWAL LSM scanner full-text search requires the FTS index on '{column}' to be maintained by the write spec (LsmWriteSpec::with_maintained_indexes); otherwise un-compacted documents are omitted. set use_lsm(false) to read the base table only" + ), + }); + } + + let scanner = base_scanner(dataset, query, pk_columns, snapshots, in_memory)? + .with_overfetch_factor(LSM_OVERFETCH_FACTOR) + .full_text_search(fts)? + .limit(as_i64(limit), as_i64(offset))?; + Ok(scanner.create_plan().await?) +} + +/// Whether an index of `type_url_suffix` covering `column` is in the MemWAL spec's +/// maintained set. Only a maintained index has its catch-up tracked (so exclusion is +/// gated correctly) and its in-memory arm kept current; an unmaintained base index +/// falls back to the compaction watermark and can drop rows it has not re-indexed. +/// The type must match specifically — a maintained BTree on the same column is not +/// the FTS/vector index the arm relies on. +async fn index_maintained( + dataset: &Dataset, + column: &str, + maintained: &[String], + type_url_suffix: &str, +) -> Result { + use lance::index::DatasetIndexExt; + let Some(field) = dataset.schema().field(column) else { + return Ok(false); + }; + let indices = dataset.load_indices().await?; + Ok(indices.iter().any(|idx| { + idx.fields.contains(&field.id) + && maintained.iter().any(|m| m == &idx.name) + && idx + .index_details + .as_ref() + .is_some_and(|d| d.type_url.ends_with(type_url_suffix)) + })) +} + +/// The maintained base index the query's arm relies on (vector index for ANN, FTS +/// index for full-text), used to gate SSTable compaction exclusion by index catch-up. +/// `None` for a plain scan or when no maintained index covers the searched column. +async fn arm_maintained_index_name( + dataset: &Dataset, + query: &VectorQueryRequest, + details: &MemWalIndexDetails, +) -> Result> { + use lance::index::DatasetIndexExt; + // Resolve the arm's searched column, the index-detail type it relies on, and a + // label for diagnostics — catch-up is taken from the vector/FTS index + // specifically, not a BTree on the same column. + let (column, type_url_suffix, arm) = if !query.query_vector.is_empty() { + let arrow_schema = ArrowSchema::from(dataset.schema()); + let column = match &query.column { + Some(column) => column.clone(), + None => { + let dim = query.query_vector.first().map(|v| v.len() as i32); + default_vector_column(&arrow_schema, dim)? + } + }; + (column, "VectorIndexDetails", "vector") + } else if let Some(fts) = &query.base.full_text_search { + match fts.columns().into_iter().next() { + Some(column) => (column, "InvertedIndexDetails", "full-text"), + None => return Ok(None), + } + } else { + return Ok(None); + }; + let Some(field) = dataset.schema().field(&column) else { + return Ok(None); + }; + let indices = dataset.load_indices().await?; + let segment_names: Vec = indices + .iter() + .filter(|idx| { + idx.fields.contains(&field.id) + && idx + .index_details + .as_ref() + .is_some_and(|d| d.type_url.ends_with(type_url_suffix)) + }) + .map(|idx| idx.name.clone()) + .collect(); + resolve_single_index(segment_names, &details.maintained_indexes, arm, &column) +} + +/// Resolve the single logical index from the names of its matching physical +/// segments. `load_indices` returns one entry per segment, so one logical index can +/// appear multiple times (same name); dedupe by name before counting. Errors when +/// more than one *distinct* index covers the field — the base planner's choice is +/// ambiguous and their catch-up watermarks can diverge, so gating exclusion on the +/// wrong one could drop SSTables the used index has not caught up to. Otherwise +/// returns the name only when it is maintained (else the caller falls back to the +/// compaction watermark). +fn resolve_single_index( + mut names: Vec, + maintained: &[String], + arm: &str, + column: &str, +) -> Result> { + names.sort(); + names.dedup(); + if names.len() > 1 { + return Err(Error::NotSupported { + message: format!( + "the MemWAL LSM scanner cannot resolve the {arm} index catch-up watermark for '{column}': it has multiple {arm} indexes; set use_lsm(false) to read the base table only" + ), + }); + } + Ok(names + .into_iter() + .next() + .filter(|name| maintained.contains(name))) +} + +/// Drop the primary-key columns Lance appends internally for dedup when the user's +/// explicit projection did not request them, restoring the requested output schema. +/// `Select::All` legitimately includes the pk columns and is left untouched. +fn restore_projection( + plan: Arc, + query: &VectorQueryRequest, + pk_columns: &[String], +) -> Result> { + let Select::Columns(selected) = &query.base.select else { + return Ok(plan); + }; + let schema = plan.schema(); + // Keep a column unless it is a pk column the user did not select (this preserves + // user columns and score columns like `_distance`, dropping only leaked pk). + let keep: Vec<(Arc, String)> = schema + .fields() + .iter() + .enumerate() + .filter(|(_, f)| { + selected.iter().any(|c| c == f.name()) || !pk_columns.iter().any(|pk| pk == f.name()) + }) + .map(|(i, f)| { + ( + Arc::new(Column::new(f.name(), i)) as Arc, + f.name().clone(), + ) + }) + .collect(); + if keep.len() == schema.fields().len() { + return Ok(plan); + } + Ok(Arc::new(ProjectionExec::try_new(keep, plan)?)) +} + +/// Vector (ANN) search over base ∪ SSTables ∪ in-memory, routed through the same +/// [`LsmScanner`] as the other arms so the query filter is applied as a prefilter. +/// +/// Note: the base and SSTable arms use `fast_search` (indexed data only), so a +/// base-table row not covered by a vector index is invisible here — it surfaces +/// only via the memtable or `use_lsm(false)`. +#[allow(clippy::too_many_arguments)] +async fn vector_plan( + dataset: &Dataset, + query: &VectorQueryRequest, + details: &MemWalIndexDetails, + pk_columns: Vec, + snapshots: Vec, + in_memory: HashMap, + limit: Option, + offset: Option, +) -> Result> { + let query_vector = query + .query_vector + .first() + .cloned() + .ok_or_else(|| Error::InvalidInput { + message: "vector search requires a query vector".to_string(), + })?; + + let arrow_schema = ArrowSchema::from(dataset.schema()); + let column = match &query.column { + Some(column) => column.clone(), + None => default_vector_column(&arrow_schema, Some(query_vector.len() as i32))?, + }; + + // The base arm relies on the column's vector index (`fast_search`). Unless it is + // maintained, its catch-up is untracked and exclusion falls back to the + // compaction watermark — dropping compacted SSTables the (lagging) base index has + // not re-indexed. Reject rather than silently omit rows, mirroring the FTS arm. + if !index_maintained( + dataset, + &column, + &details.maintained_indexes, + "VectorIndexDetails", + ) + .await? + { + return Err(Error::NotSupported { + message: format!( + "the MemWAL LSM scanner requires the vector index on '{column}' to be maintained by the write spec (LsmWriteSpec::with_maintained_indexes); otherwise compacted rows not yet re-indexed are omitted. set use_lsm(false) to read the base table only" + ), + }); + } + + // The LSM vector planner is Float32-only; reject binary (uint8) vectors with a + // clear error rather than failing deep in the planner. + if is_binary_vector_column(&arrow_schema, &column) { + return Err(Error::NotSupported { + message: "the MemWAL LSM scanner does not support binary (uint8) vector search; set use_lsm(false) to read the base table only".to_string(), + }); + } + + let distance_type = resolve_distance_type(dataset, query, &column).await?; + + // `nearest` takes a flat query vector and builds the fixed-size list itself. + // Guard `k` to at least 1 so a degenerate `limit(0)` is trimmed by `limit` + // below rather than rejected by `nearest`. + let k = limit.unwrap_or(DEFAULT_TOP_K).max(1); + let mut scanner = base_scanner(dataset, query, pk_columns, snapshots, in_memory)? + .with_overfetch_factor(LSM_OVERFETCH_FACTOR) + .nearest(&column, query_vector.as_ref(), k)? + .nprobes(query.minimum_nprobes) + .distance_metric(distance_type.into()); + if let Some(refine_factor) = query.refine_factor { + scanner = scanner.refine(refine_factor); + } + scanner = scanner.limit(as_i64(limit), as_i64(offset))?; + Ok(scanner.create_plan().await?) +} + +/// Whether `column` stores binary (uint8) vectors, which the LSM vector planner +/// does not support. +fn is_binary_vector_column(schema: &ArrowSchema, column: &str) -> bool { + matches!( + schema.field_with_name(column).map(|f| f.data_type()), + Ok(DataType::FixedSizeList(field, _)) if matches!(field.data_type(), DataType::UInt8) + ) +} + +/// Resolve the distance metric for the vector arm: the explicit query metric if +/// set, else the metric of the column's vector index, else L2. +async fn resolve_distance_type( + dataset: &Dataset, + query: &VectorQueryRequest, + column: &str, +) -> Result { + if let Some(dt) = query.distance_type { + return Ok(dt); + } + // Inherit the column's vector-index metric so cross-source distances match + // the metric the maintained memtable index was built with. + use lance::index::{DatasetIndexExt, DatasetIndexInternalExt}; + use lance_index::metrics::NoOpMetricsCollector; + let field = dataset.schema().field(column); + if let Some(field) = field { + let indices = dataset.load_indices().await?; + for index in indices.iter() { + if index.fields.contains(&field.id) + && let Ok(vector_index) = dataset + .open_vector_index(column, &index.uuid, &NoOpMetricsCollector) + .await + { + return Ok(vector_index.metric_type().into()); + } + } + } + Ok(DistanceType::L2) +} + +#[cfg(test)] +mod tests { + use super::*; + use lance_index::mem_wal::{CompactedSsTable, IndexCatchupProgress}; + + #[test] + fn exclusion_watermark_gates_on_lagging_index_catchup() { + let shard = Uuid::from_u128(1); + let details = MemWalIndexDetails { + // Compaction has drained generations through 5 into the base table... + compacted_sstables: vec![CompactedSsTable::new(shard, 5)], + // ...but the FTS index has only caught up through generation 2. + index_catchup: vec![IndexCatchupProgress::new( + "fts_idx".to_string(), + vec![CompactedSsTable::new(shard, 2)], + )], + maintained_indexes: vec!["fts_idx".to_string()], + ..Default::default() + }; + + // Plain scan: drop every compacted generation (through 5). + assert_eq!(exclusion_watermarks(&details, None).get(&shard), Some(&5)); + + // FTS arm with a lagging index: exclusion is capped at the index catch-up + // (2), so SSTable generations 3..=5 are retained until the index covers + // them — otherwise those documents would silently vanish from FTS results. + assert_eq!( + exclusion_watermarks(&details, Some("fts_idx")).get(&shard), + Some(&2) + ); + + // A caught-up index — or one untracked in index_catchup — falls back to the + // compaction watermark. + assert_eq!( + exclusion_watermarks(&details, Some("caught_up_idx")).get(&shard), + Some(&5) + ); + } + + #[test] + fn resolve_single_index_dedupes_segments() { + let maintained = vec!["fts_idx".to_string()]; + // Two physical segments of ONE logical index must not count as "multiple". + assert_eq!( + resolve_single_index( + vec!["fts_idx".to_string(), "fts_idx".to_string()], + &maintained, + "full-text", + "text" + ) + .unwrap(), + Some("fts_idx".to_string()) + ); + // Two distinct indexes on the field are ambiguous → error. + assert!( + resolve_single_index( + vec!["fts_a".to_string(), "fts_b".to_string()], + &maintained, + "full-text", + "text" + ) + .is_err() + ); + // A single unmaintained index resolves to None (compaction-watermark fallback). + assert_eq!( + resolve_single_index(vec!["other".to_string()], &maintained, "full-text", "text") + .unwrap(), + None + ); + } +}