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