From dd29553669a3a2fa7fc31c247edea969a2778499 Mon Sep 17 00:00:00 2001 From: Daniel Rammer Date: Thu, 8 Oct 2026 17:36:42 -0500 Subject: [PATCH 1/3] perf(mem_wal): rank only prefiltered, newest memtable rows in a filtered vector search A filtered vector search over the memtable runs `MemTableBruteForceVectorExec`, which on every query hashed every visible key to find each one's newest version, computed a distance for every visible row, applied the filter only afterwards, and did all of it synchronously inside `execute`. On 100,000 rows that is the whole cost of the search, however selective the filter. Narrow the rows before any distance is computed: - Ask the filter indexes for candidates, through the same split of the filter into index searches and a leftover that a filtered scan uses, and apply the filter only to candidates the indexes did not settle. - Keep each key's newest version without the pass over every key when no key was ever rewritten (the condition that already lets an unfiltered search use HNSW), and with one key-index seek per row when few rows pass the filter. - Compute distances only for the rows left, or for the whole batch when most of it is left. - Rank on the CPU pool from inside the stream rather than in `execute`. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../mem_wal/memtable/scanner/builder.rs | 34 ++- .../scanner/exec/brute_force_vector.rs | 272 ++++++++++++++---- 2 files changed, 231 insertions(+), 75 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs index 514d5fadfe1..b2a3e9ee22c 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs @@ -1318,18 +1318,28 @@ impl MemTableScanner { self.with_row_id, )?) } else { - Arc::new( - MemTableBruteForceVectorExec::new( - self.batch_store.clone(), - query.clone(), - max_readable, - projection_indices, - base_schema, - self.with_row_id, - )? - .with_filter(filter_predicate) - .with_pk_columns(self.pk_columns.clone()), - ) + let mut exec = MemTableBruteForceVectorExec::new( + self.batch_store.clone(), + query.clone(), + max_readable, + projection_indices, + base_schema, + self.with_row_id, + )? + .with_filter(filter_predicate) + .with_pk_columns(self.pk_columns.clone()) + .with_indexes(self.indexes.clone()); + // Rank only the rows a filtered scan would read, via the same + // index/leftover split of the filter. + if self.use_index + && let Some(filter) = &self.filter + { + let optimized = Planner::new(self.schema.clone()).optimize_expr(filter.clone())?; + if let Some(indexed) = plan_filter(&optimized, self.indexes.filter_catalog())? { + exec = exec.with_index_filter(indexed.searches, indexed.is_whole_filter); + } + } + Arc::new(exec) }; self.apply_post_index_ops(exec).await } diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/brute_force_vector.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/brute_force_vector.rs index b8da5e93a2b..9a55320debf 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/brute_force_vector.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/exec/brute_force_vector.rs @@ -11,11 +11,13 @@ //! or new rows in the window between commit and next memtable rotation), this //! exec keeps KNN correct by computing exact distances row-by-row. +use std::collections::HashSet; use std::fmt::{Debug, Formatter}; use std::sync::Arc; use arrow_array::{Array, BooleanArray, Float32Array, RecordBatch, UInt64Array, cast::AsArray}; use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use datafusion::common::ScalarValue; use datafusion::common::stats::Precision; use datafusion::error::Result as DataFusionResult; use datafusion::execution::TaskContext; @@ -27,15 +29,19 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, Statistics, }; use datafusion_physical_expr::{EquivalenceProperties, PhysicalExprRef}; -use futures::stream::{self, StreamExt}; +use futures::stream::{self, StreamExt, TryStreamExt}; +use lance_core::utils::tokio::spawn_cpu; use lance_core::{Error, Result}; +use lance_index::scalar::expression::ScalarIndexExpr; use lance_linalg::distance::DistanceType; use super::super::builder::VectorQuery; use super::newest_pk_positions; use super::vector::DISTANCE_COLUMN; +use crate::dataset::mem_wal::index::{SearchContext, evaluate_index_filter}; use crate::dataset::mem_wal::memtable::scanner::exec::{scan_record_batch, take_projected_columns}; -use crate::dataset::mem_wal::write::BatchStore; +use crate::dataset::mem_wal::scanner::exec::resolve_pk_indices; +use crate::dataset::mem_wal::write::{BatchStore, IndexStore}; /// Distance metric used when [`VectorQuery::distance_type`] is `None`. The /// indexed path defers to the index's own metric, but with no index there is @@ -43,8 +49,33 @@ use crate::dataset::mem_wal::write::BatchStore; /// SSTable/base arms use when re-ranking unindexed candidates. const DEFAULT_DISTANCE_TYPE: DistanceType = DistanceType::L2; +/// Once over `1 / NEWEST_SEEK_SHARE` of visible rows pass, one pass hashing +/// every visible key beats a key-index seek per row. +const NEWEST_SEEK_SHARE: u64 = 8; + +/// The part of the prefilter the memtable's filter indexes answer. +#[derive(Debug, Clone)] +struct IndexFilter { + /// The index searches the filter was split into. + expr: ScalarIndexExpr, + /// Whether the searches are the whole filter; if not, even an exact answer + /// is only a superset. + covers_filter: bool, +} + +/// How a search keeps only each primary key's newest visible version. +enum Newest { + /// Every row is its key's newest version, or there is no key. + Any, + /// The newest versions' positions, from a pass over every visible key. + Among(HashSet), + /// A seek in the key index per row. + Seek, +} + /// Brute-force KNN over an active memtable without an HNSW. Produces the same /// output schema as [`super::VectorIndexExec`]. +#[derive(Clone)] pub struct MemTableBruteForceVectorExec { batch_store: Arc, query: VectorQuery, @@ -62,6 +93,11 @@ pub struct MemTableBruteForceVectorExec { /// eligible for top-k. With a filter, this evaluates the predicate against /// the current PK version instead of falling back to a stale older version. pk_columns: Option>, + /// Key index for newest-version checks; filter indexes to narrow the + /// prefilter's rows. + indexes: Option>, + /// The prefilter as searches in `indexes`, when they answer some of it. + index_filter: Option, } impl Debug for MemTableBruteForceVectorExec { @@ -116,6 +152,8 @@ impl MemTableBruteForceVectorExec { with_row_id, filter: None, pk_columns: None, + indexes: None, + index_filter: None, }) } @@ -133,6 +171,23 @@ impl MemTableBruteForceVectorExec { self } + /// Use the memtable's indexes to skip or narrow newest-version checks. + /// Required by [`Self::with_index_filter`]. + pub fn with_indexes(mut self, indexes: Arc) -> Self { + self.indexes = Some(indexes); + self + } + + /// Narrow the prefilter's rows with `expr`, the indexed part of the filter; + /// `covers_filter` if it is the whole filter. Inert without [`Self::with_indexes`]. + pub fn with_index_filter(mut self, expr: ScalarIndexExpr, covers_filter: bool) -> Self { + self.index_filter = Some(IndexFilter { + expr, + covers_filter, + }); + self + } + /// Evaluate the prefilter predicate against a memtable batch, returning a /// keep-mask (`true` = retain). `Ok(None)` when no filter is configured. fn filter_mask(&self, batch: &RecordBatch) -> Result> { @@ -197,10 +252,110 @@ impl MemTableBruteForceVectorExec { Ok(self.query.query_vector.clone()) } - /// Compute `(distance, row_position)` for every visible row, then top-k by - /// distance ascending. Rows where the vector column is null or where the - /// computed distance is non-finite are skipped — same convention as the - /// HNSW search (which filters on `result.distance.is_finite()`). + /// Ascending candidate positions from the filter indexes (`None`: every + /// visible row), and whether they are exactly the rows the filter keeps. + fn index_candidates(&self, max_readable_row: u64) -> (Option>, bool) { + let (Some(indexes), Some(index_filter)) = (&self.indexes, &self.index_filter) else { + return (None, false); + }; + let ctx = SearchContext::new(max_readable_row); + match evaluate_index_filter(&index_filter.expr, indexes, &ctx) { + Ok(result) => { + let settled = result.is_exact() && index_filter.covers_filter; + if result.at_most.len() == max_readable_row + 1 { + (None, settled) + } else { + (Some(result.at_most.into()), settled) + } + } + // A failed search falls back to filtering every row, never "no rows". + Err(error) => { + log::warn!( + "a memtable index failed to search {}; filtering every row instead: {error}", + index_filter.expr + ); + (None, false) + } + } + } + + /// How to keep only each key's newest visible version, given `passed` + /// prefiltered rows: skip with no rewrites, seek when few, else hash every key. + fn newest_rule(&self, max_readable_row: u64, passed: u64) -> Result { + let Some(pk_columns) = &self.pk_columns else { + return Ok(Newest::Any); + }; + if let Some(indexes) = &self.indexes + && indexes.has_pk_index() + { + // Read at execution: the flag is set before a rewrite is visible + // and never cleared. + if !indexes.pk_has_overrides() { + return Ok(Newest::Any); + } + if passed <= (max_readable_row + 1) / NEWEST_SEEK_SHARE { + return Ok(Newest::Seek); + } + } + newest_pk_positions( + &self.batch_store, + pk_columns, + self.readable_count, + max_readable_row, + ) + .map(Newest::Among) + .map_err(|e| Error::invalid_input(e.to_string())) + } + + /// Visible rows passing the prefilter, per batch: `(first row position, + /// batch, kept row offsets)`. Rows are tracked by offset, never copied. + fn prefiltered_rows(&self, max_readable_row: u64) -> Result)>> { + let (index_candidates, settled) = self.index_candidates(max_readable_row); + let mut current_row: u64 = 0; + let mut next_candidate = 0usize; + let mut passed = Vec::new(); + for (batch_position, stored_batch) in self.batch_store.iter().enumerate() { + let n = stored_batch.num_rows; + let start = current_row; + current_row += n as u64; + if n == 0 || batch_position >= self.readable_count || start > max_readable_row { + continue; + } + let end = current_row.min(max_readable_row + 1); + + let mut rows: Vec = match &index_candidates { + Some(positions) => { + let first = next_candidate; + while next_candidate < positions.len() && positions[next_candidate] < end { + next_candidate += 1; + } + positions[first..next_candidate] + .iter() + .map(|&pos| (pos - start) as u32) + .collect() + } + None => (0..(end - start) as u32).collect(), + }; + if rows.is_empty() { + continue; + } + let scan_batch = scan_record_batch(&stored_batch.data)?; + + // A NULL predicate result excludes the row, as in SQL. + if !settled && let Some(mask) = self.filter_mask(&scan_batch)? { + rows.retain(|&row| mask.is_valid(row as usize) && mask.value(row as usize)); + } + if !rows.is_empty() { + passed.push((start, scan_batch, rows)); + } + } + Ok(passed) + } + + /// Top-k `(distance, row_position)`, nearest first, over prefiltered + /// newest-version rows; null vectors and non-finite distances are skipped, + /// as in the HNSW search. Distances cover whole batches: copying a vector + /// out costs far more than its distance. fn compute_topk(&self) -> Result> { if self.query.k == 0 { return Ok(Vec::new()); @@ -213,41 +368,48 @@ impl MemTableBruteForceVectorExec { let distance_type = self.query.distance_type.unwrap_or(DEFAULT_DISTANCE_TYPE); let batch_func = distance_type.arrow_batch_func(); + let passed = self.prefiltered_rows(max_readable_row)?; // When PK columns are configured, only the newest version of each PK is // eligible. This keeps top-k slots from being consumed by superseded // rows and makes filtered search evaluate the predicate against the // current version of the PK. - let newest_positions = if let Some(pk_columns) = &self.pk_columns { - Some( - newest_pk_positions( - &self.batch_store, - pk_columns, - self.readable_count, - max_readable_row, - ) - .map_err(|e| Error::invalid_input(e.to_string()))?, - ) - } else { - None - }; + let passed_rows = passed.iter().map(|(_, _, rows)| rows.len() as u64).sum(); + let newest = self.newest_rule(max_readable_row, passed_rows)?; - // Walk batches in append order. `current_row` is the global row offset - // of the *next* row about to be visited; rows past `max_readable_row` - // are dropped before they reach the heap. - let mut current_row: u64 = 0; let mut candidates: Vec<(f32, u64)> = Vec::new(); - - for (batch_position, stored_batch) in self.batch_store.iter().enumerate() { - let n = stored_batch.num_rows; - if n == 0 { - continue; + for (start, scan_batch, mut rows) in passed { + // Only a key's newest version is eligible, so a newer non-matching + // version excludes the key. + match &newest { + Newest::Any => {} + Newest::Among(positions) => { + rows.retain(|&row| positions.contains(&(start + row as u64))); + } + Newest::Seek => { + let (Some(indexes), Some(pk_columns)) = (&self.indexes, &self.pk_columns) + else { + unreachable!("a seek is chosen only with a key index and key columns"); + }; + let pk_indices = resolve_pk_indices(&scan_batch, pk_columns)?; + let mut kept = Vec::with_capacity(rows.len()); + for row in rows { + let values = pk_indices + .iter() + .map(|&column| { + ScalarValue::try_from_array(scan_batch.column(column), row as usize) + }) + .collect::>>()?; + if indexes.pk_is_newest(&values, start + row as u64, max_readable_row) { + kept.push(row); + } + } + rows = kept; + } } - if batch_position >= self.readable_count { - current_row += n as u64; + if rows.is_empty() { continue; } - let scan_batch = scan_record_batch(&stored_batch.data)?; let column = scan_batch.column_by_name(column_name).ok_or_else(|| { Error::invalid_input(format!( "Vector column '{}' not found in memtable schema", @@ -269,38 +431,16 @@ impl MemTableBruteForceVectorExec { )) })?; - // Prefilter: drop rows that fail the predicate before they reach the - // top-k heap (a NULL predicate result excludes the row, matching SQL). - let filter_mask = self.filter_mask(&scan_batch)?; - - for row in 0..n { - let pos = current_row + row as u64; - if pos > max_readable_row { - break; - } - // Skip superseded versions: only the newest version of each PK is - // eligible, so a newer non-matching version excludes the PK. - if let Some(ref newest) = newest_positions - && !newest.contains(&pos) - { - continue; - } - if let Some(ref mask) = filter_mask - && (!mask.is_valid(row) || !mask.value(row)) - { - continue; - } - if distances.is_null(row) { + for row in rows { + if distances.is_null(row as usize) { continue; } - let dist = distances.value(row); + let dist = distances.value(row as usize); if !dist.is_finite() { continue; } - candidates.push((dist, pos)); + candidates.push((dist, start + row as u64)); } - - current_row += n as u64; } // `partial_cmp` defaults Equal on NaN; we filtered non-finite above so @@ -452,14 +592,20 @@ impl ExecutionPlan for MemTableBruteForceVectorExec { _partition: usize, _context: Arc, ) -> DataFusionResult { - let results = self - .compute_topk() - .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; - let batches = self.materialize_rows(&results)?; - let stream = stream::iter(batches.into_iter().map(Ok)).boxed(); + // Ranking is CPU-bound: run it on the CPU pool when polled, not in `execute`. + let exec = self.clone(); + let batches = stream::once(spawn_cpu(move || { + let results = exec + .compute_topk() + .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; + exec.materialize_rows(&results) + })) + .map_ok(|batches| stream::iter(batches.into_iter().map(Ok))) + .try_flatten() + .boxed(); Ok(Box::pin(RecordBatchStreamAdapter::new( self.output_schema.clone(), - stream, + batches, ))) } From 09d6061980099a5499ea213dfaad5bd85c6cec92 Mon Sep 17 00:00:00 2001 From: Daniel Rammer Date: Thu, 8 Oct 2026 17:36:49 -0500 Subject: [PATCH 2/3] test(mem_wal): check a prefiltered memtable vector search against expected answers Ranks a prefiltered search over a keyed memtable and compares ids and distances with the nearest keys computed directly from the rows written: filters an index answers exactly, partly, or not at all, selective and broad, across rewrites, deletes, rows indexed but not yet visible, and distance bounds. An ignored test times a prefiltered search over 100,000 rows. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../mem_wal/memtable/scanner/builder.rs | 264 ++++++++++++++++++ 1 file changed, 264 insertions(+) diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs index b2a3e9ee22c..ef94cf111d8 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs @@ -3714,4 +3714,268 @@ mod tests { } } } + + /// One row of [`prefilter_memtable`], kept for computing expected answers. + struct PrefilterRow { + id: i64, + category: Option, + /// `None` for a delete: a newer version of the key with no vector. + vector: Option>, + batch: usize, + } + + /// A memtable keyed on `id`, with a B-tree on `category` (16 values plus + /// nulls) and an unindexed `vector`. Batches after the first rewrite + /// `rewrites` earlier keys and delete `deletes` more. All batches are + /// indexed; only the first `visible_batches` are visible. + fn prefilter_memtable( + batches: usize, + batch_rows: usize, + dim: i32, + rewrites: usize, + deletes: usize, + visible_batches: usize, + ) -> ( + Arc, + Arc, + SchemaRef, + Vec, + ) { + use crate::dataset::mem_wal::wal::WriterCursors; + use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; + + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("category", DataType::Utf8, true), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), dim), + true, + ), + ])); + let lance_schema = LanceSchema::try_from(schema.as_ref()).unwrap(); + let mut indexes = IndexStore::from_specs( + &[MemIndexSpec::btree("category_idx", 1, "category")], + &lance_schema, + batches * batch_rows, + batches, + ) + .unwrap(); + indexes.enable_pk_index(&[("id".to_string(), 0)]); + let cursors = Arc::new(WriterCursors::new(true)); + indexes.set_durability(cursors.clone(), 0); + let batch_store = Arc::new(BatchStore::with_capacity(batches)); + + let mut seed: u64 = 0x9e3779b97f4a7c15; + let mut next = move || { + seed = seed + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + seed >> 33 + }; + let mut rows = Vec::new(); + let mut next_id = 0i64; + for batch_index in 0..batches { + let (mut ids, mut categories) = (vec![], vec![]); + let mut vectors = FixedSizeListBuilder::new(Float32Builder::new(), dim); + for row in 0..batch_rows { + let id = if batch_index > 0 && row < rewrites + deletes { + (next() % next_id as u64) as i64 + } else { + next_id += 1; + next_id - 1 + }; + let deleted = batch_index > 0 && (rewrites..rewrites + deletes).contains(&row); + let category = match next() % 17 { + _ if deleted => None, + 16 => None, + c => Some(format!("cat{c:02}")), + }; + let vector: Option> = + (!deleted).then(|| (0..dim).map(|_| (next() % 1000) as f32 / 1000.0).collect()); + match &vector { + Some(vector) => { + vectors.values().append_slice(vector); + vectors.append(true); + } + None => { + vectors.values().append_nulls(dim as usize); + vectors.append(false); + } + } + ids.push(id); + categories.push(category.clone()); + rows.push(PrefilterRow { + id, + category, + vector, + batch: batch_index, + }); + } + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(ids)), + Arc::new(StringArray::from(categories)), + Arc::new(vectors.finish()), + ], + ) + .unwrap(); + let (position, offset, _) = batch_store.append(batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&batch, offset, Some(position)) + .unwrap(); + } + cursors.advance_durable(visible_batches); + (batch_store, Arc::new(indexes), schema, rows) + } + + /// `(id, _distance)` of a prefiltered search, nearest first. + async fn prefiltered_nearest( + batch_store: &Arc, + indexes: &Arc, + schema: &SchemaRef, + filter: &str, + query: &[f32], + (lower, upper): (Option, Option), + k: usize, + ) -> Vec<(i64, f32)> { + use arrow_array::cast::AsArray; + + let mut scanner = + MemTableScanner::new(batch_store.clone(), indexes.clone(), schema.clone()); + scanner.with_pk_columns(vec!["id".to_string()]); + scanner.filter(filter).unwrap(); + let query: Arc = + Arc::new(arrow_array::Float32Array::from(query.to_vec())); + scanner.nearest("vector", query.as_ref(), k).unwrap(); + scanner.distance_range(lower, upper); + let batch = scanner.try_into_batch().await.unwrap(); + let ids = batch["id"].as_primitive::(); + let distances = batch["_distance"].as_primitive::(); + let mut out: Vec<(i64, f32)> = ids + .values() + .iter() + .zip(distances.values().iter()) + .map(|(&id, &distance)| (id, distance)) + .collect(); + out.sort_by(|a, b| a.1.total_cmp(&b.1).then(a.0.cmp(&b.0))); + out + } + + /// The nearest `k` keys whose newest version passes `keep`, computed + /// directly from the rows written. + fn expected_nearest( + rows: &[PrefilterRow], + visible_batches: usize, + keep: impl Fn(&PrefilterRow) -> bool, + query: &[f32], + (lower, upper): (Option, Option), + k: usize, + ) -> Vec<(i64, f32)> { + let mut newest: std::collections::HashMap = + std::collections::HashMap::new(); + for row in rows.iter().filter(|row| row.batch < visible_batches) { + newest.insert(row.id, row); + } + let l2 = DistanceType::L2.func::(); + let mut out: Vec<(i64, f32)> = newest + .values() + .filter(|row| keep(row)) + .filter_map(|row| Some((row.id, l2(query, row.vector.as_ref()?)))) + .filter(|&(_, d)| lower.is_none_or(|lb| d >= lb) && upper.is_none_or(|ub| d < ub)) + .collect(); + out.sort_by(|a, b| a.1.total_cmp(&b.1).then(a.0.cmp(&b.0))); + out.truncate(k); + out + } + + /// A prefiltered search ranks exactly the keys whose newest visible version + /// passes the filter, with exact, partial, or no index answers, across + /// rewrites, deletes, indexed-but-invisible rows, and distance bounds. + #[tokio::test] + async fn a_prefiltered_vector_search_ranks_only_the_newest_matching_versions() { + type Keep = fn(&PrefilterRow) -> bool; + let cases: [(&str, Keep); 5] = [ + ("category = 'cat03'", |row| { + row.category.as_deref() == Some("cat03") + }), + ("category IN ('cat01', 'cat05', 'cat09')", |row| { + matches!(row.category.as_deref(), Some("cat01" | "cat05" | "cat09")) + }), + // Matches about half the rows. + ("category < 'cat08'", |row| { + row.category.as_deref().is_some_and(|c| c < "cat08") + }), + // Only half of it is indexed. + ("category = 'cat03' AND id % 2 = 0", |row| { + row.category.as_deref() == Some("cat03") && row.id % 2 == 0 + }), + // Nothing indexed. + ("id % 7 = 3", |row| row.id % 7 == 3), + ]; + // Six batches, the last two indexed but not yet visible. + let (batch_store, indexes, schema, rows) = prefilter_memtable(6, 500, 4, 120, 30, 4); + assert!(indexes.pk_has_overrides(), "the fixture must rewrite keys"); + for (filter, keep) in cases { + for query in [[0.1, 0.2, 0.3, 0.4], [0.9, 0.1, 0.5, 0.7]] { + for bounds in [(None, None), (Some(0.05), Some(0.6))] { + let got = prefiltered_nearest( + &batch_store, + &indexes, + &schema, + filter, + &query, + bounds, + 10, + ) + .await; + let expected = expected_nearest(&rows, 4, keep, &query, bounds, 10); + assert_eq!( + got.iter().map(|(id, _)| *id).collect::>(), + expected.iter().map(|(id, _)| *id).collect::>(), + "filter: {filter}, bounds: {bounds:?}" + ); + for ((_, got), (_, expected)) in got.iter().zip(&expected) { + assert!((got - expected).abs() < 1e-5, "filter: {filter}"); + } + } + } + } + } + + /// Timing for prefiltered searches over 100,000 rows of 128-dim vectors, + /// with and without key rewrites. Run optimized with `--ignored --nocapture`. + #[tokio::test] + #[ignore] + #[allow(clippy::print_stdout)] + async fn time_a_prefiltered_vector_search_over_a_large_memtable() { + let query: Vec = (0..128).map(|i| (i % 10) as f32 / 10.0).collect(); + for rewrites in [0, 1_000] { + let (batch_store, indexes, schema, _) = + prefilter_memtable(10, 10_000, 128, rewrites, 0, 10); + for filter in ["category = 'cat03'", "id % 16 = 3", "category < 'cat08'"] { + let mut times = Vec::new(); + for _ in 0..11 { + let start = std::time::Instant::now(); + prefiltered_nearest( + &batch_store, + &indexes, + &schema, + filter, + &query, + (None, None), + 10, + ) + .await; + times.push(start.elapsed().as_secs_f64() * 1000.0); + } + times.sort_by(f64::total_cmp); + println!( + "rewrites={rewrites:5} filter={filter:20} median={:7.1} ms", + times[times.len() / 2] + ); + } + } + } } From 977ba1e8955985da7d77850e6a08fdfc3d6d2b05 Mon Sep 17 00:00:00 2001 From: Daniel Rammer Date: Mon, 12 Oct 2026 20:18:04 -0500 Subject: [PATCH 3/3] fix(mem_wal): order memtable key columns as the key index expects A scanner given its primary-key columns in another order than the key index built lookup tuples the index never stored, so a selective prefilter that seeks the index dropped valid newest rows. Put the columns in the index's order once per plan, for vector search, FTS and the filtered dedup scan, and reject columns naming another key. Co-Authored-By: Claude Opus 5.5 (1M context) --- rust/lance/src/dataset/mem_wal/index/pk.rs | 12 +- .../mem_wal/memtable/scanner/builder.rs | 159 +++++++++++++++++- 2 files changed, 164 insertions(+), 7 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/index/pk.rs b/rust/lance/src/dataset/mem_wal/index/pk.rs index aec1c8995a3..8903809dbca 100644 --- a/rust/lance/src/dataset/mem_wal/index/pk.rs +++ b/rust/lance/src/dataset/mem_wal/index/pk.rs @@ -13,7 +13,7 @@ use datafusion::common::ScalarValue; use lance_core::{Error, Result}; use super::pk_key::{encode_pk_batch, encode_pk_tuple}; -use super::{BTreeMemIndex, IndexStore, PrimaryKeyIndex, RowPosition}; +use super::{BTreeMemIndex, IndexStore, MemIndex, PrimaryKeyIndex, RowPosition}; /// The column a composite key index is keyed on: the order-preserving encoded /// tuple, stored as `Binary` so a [`BTreeMemIndex`] indexes it directly. @@ -182,6 +182,16 @@ impl IndexStore { self.pk_index.is_some() } + /// The primary-key index's key columns, in the order its lookups take + /// values. + pub fn pk_columns(&self) -> Option<&[String]> { + match self.pk_index.as_ref()? { + PkIndex::Shared { entry, .. } => self.indexes.get(entry).map(|index| index.columns()), + PkIndex::Owned(OwnedPk::Single(index)) => Some(index.columns()), + PkIndex::Owned(OwnedPk::Composite { columns, .. }) => Some(columns), + } + } + /// Sorted `(key, row id)` batches for training the generation's on-disk /// key index: the typed key for one column, the encoded tuple for several. /// Positions are the flushed file's row ids. diff --git a/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs b/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs index ef94cf111d8..e8994e0ec83 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs +++ b/rust/lance/src/dataset/mem_wal/memtable/scanner/builder.rs @@ -1165,6 +1165,7 @@ impl MemTableScanner { /// only if it is its key's newest visible version. It never pushes a limit (the LSM caps /// results above the cross-source merge). pub async fn create_dedup_plan(&self, pk_columns: &[String]) -> Result> { + let pk_columns = &self.key_columns(pk_columns)?; validate_pk_types(&self.schema, pk_columns)?; let pk_indices = pk_columns @@ -1279,7 +1280,12 @@ impl MemTableScanner { let projection_indices = self.compute_projection_indices()?; let base_schema = self.base_output_schema()?; let filter_predicate = self.filter_predicate()?; - if let Some(pk_columns) = &self.pk_columns { + let pk_columns = self + .pk_columns + .as_deref() + .map(|columns| self.key_columns(columns)) + .transpose()?; + if let Some(pk_columns) = &pk_columns { validate_pk_types(&self.schema, pk_columns)?; } @@ -1290,8 +1296,7 @@ impl MemTableScanner { // append-only PK data can still use HNSW safely. This relies on // `IndexStore` marking PK overrides before advancing the visible batch // watermark, so any snapshot that sees a rewrite also sees the flag. - let hnsw_safe_with_pk = self - .pk_columns + let hnsw_safe_with_pk = pk_columns .as_ref() .map(|_| self.indexes.has_pk_index() && !self.indexes.pk_has_overrides()) .unwrap_or(true); @@ -1327,7 +1332,7 @@ impl MemTableScanner { self.with_row_id, )? .with_filter(filter_predicate) - .with_pk_columns(self.pk_columns.clone()) + .with_pk_columns(pk_columns) .with_indexes(self.indexes.clone()); // Rank only the rows a filtered scan would read, via the same // index/leftover split of the filter. @@ -1357,7 +1362,12 @@ impl MemTableScanner { let max_readable = self.readable_count; let projection_indices = self.compute_projection_indices()?; let filter_predicate = self.filter_predicate()?; - if let Some(pk_columns) = &self.pk_columns { + let pk_columns = self + .pk_columns + .as_deref() + .map(|columns| self.key_columns(columns)) + .transpose()?; + if let Some(pk_columns) = &pk_columns { validate_pk_types(&self.schema, pk_columns)?; } @@ -1371,10 +1381,26 @@ impl MemTableScanner { self.with_row_id, )? .with_filter(filter_predicate) - .with_pk_columns(self.pk_columns.clone()); + .with_pk_columns(pk_columns); self.apply_post_index_ops(Arc::new(index_exec)).await } + /// `pk_columns` in the key index's order, which its lookups expect. Errors + /// if they name another key. + fn key_columns(&self, pk_columns: &[String]) -> Result> { + let Some(index_key) = self.indexes.pk_columns() else { + return Ok(pk_columns.to_vec()); + }; + if pk_columns.len() != index_key.len() + || !index_key.iter().all(|column| pk_columns.contains(column)) + { + return Err(Error::invalid_input(format!( + "primary-key columns {pk_columns:?} differ from the memtable's key index columns {index_key:?}" + ))); + } + Ok(index_key.to_vec()) + } + fn empty_fts_plan( &self, document_granularity: DocumentGranularity, @@ -3944,6 +3970,127 @@ mod tests { } } + /// A memtable keyed on `(a, b)` with a B-tree on `category`: 64 keys + /// `(i, 0)` in category `old`, then key `(3, 0)` rewritten into `new`. + fn composite_key_memtable() -> (Arc, Arc, SchemaRef) { + use arrow_array::FixedSizeListArray; + use arrow_array::types::Float32Type; + + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("category", DataType::Utf8, true), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2), + true, + ), + ])); + let lance_schema = LanceSchema::try_from(schema.as_ref()).unwrap(); + let mut indexes = IndexStore::from_specs( + &[MemIndexSpec::btree("category_idx", 2, "category")], + &lance_schema, + 128, + 4, + ) + .unwrap(); + indexes.enable_pk_index(&[("a".to_string(), 0), ("b".to_string(), 1)]); + let batch_store = Arc::new(BatchStore::with_capacity(4)); + let batch = |a: Vec, category: &str| { + let rows = a.len(); + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(a.clone())), + Arc::new(Int32Array::from(vec![0; rows])), + Arc::new(StringArray::from(vec![category; rows])), + Arc::new( + FixedSizeListArray::from_iter_primitive::( + a.iter().map(|&a| Some([Some(a as f32), Some(0.0)])), + 2, + ), + ), + ], + ) + .unwrap() + }; + for batch in [batch((0..64).collect(), "old"), batch(vec![3], "new")] { + let (position, offset, _) = batch_store.append(batch.clone()).unwrap(); + indexes + .insert_with_batch_position(&batch, offset, Some(position)) + .unwrap(); + } + assert!(indexes.pk_has_overrides()); + (batch_store, Arc::new(indexes), schema) + } + + /// Key columns given in another order than the key index's still keep + /// exactly each key's newest version, whether the filter is selective + /// enough to seek the key index or broad enough to hash every key. + #[rstest::rstest] + #[case::seek("category = 'new'", 1)] + #[case::hash("category IN ('old', 'new')", 64)] + #[tokio::test] + async fn reordered_composite_key_columns_keep_the_newest_version( + #[case] filter: &str, + #[case] expected_rows: usize, + ) { + use arrow_array::cast::AsArray; + use arrow_array::types::Int32Type; + + let (batch_store, indexes, schema) = composite_key_memtable(); + let newest_of_3 = |batch: &RecordBatch| { + let a = batch["a"].as_primitive::(); + let category = batch["category"].as_string::(); + let versions: Vec<&str> = (0..batch.num_rows()) + .filter(|&row| a.value(row) == 3) + .map(|row| category.value(row)) + .collect(); + assert_eq!(versions, ["new"], "filter: {filter}"); + assert_eq!(batch.num_rows(), expected_rows, "filter: {filter}"); + }; + let reordered = vec!["b".to_string(), "a".to_string()]; + + let mut scanner = + MemTableScanner::new(batch_store.clone(), indexes.clone(), schema.clone()); + scanner.with_pk_columns(reordered.clone()); + scanner.filter(filter).unwrap(); + let query = arrow_array::Float32Array::from(vec![3.0, 0.0]); + scanner.nearest("vector", &query, 100).unwrap(); + newest_of_3(&scanner.try_into_batch().await.unwrap()); + + let mut scanner = MemTableScanner::new(batch_store, indexes, schema); + scanner.filter(filter).unwrap(); + scanner.with_memtable_filter_indexes(true); + let plan = scanner.create_dedup_plan(&reordered).await.unwrap(); + let ctx = SessionContext::new(); + let batches = datafusion::physical_plan::collect(plan, ctx.task_ctx()) + .await + .unwrap(); + newest_of_3(&arrow_select::concat::concat_batches(&batches[0].schema(), &batches).unwrap()); + } + + /// Key columns naming another key than the key index's are rejected. + #[tokio::test] + async fn key_columns_naming_another_key_are_rejected() { + let (batch_store, indexes, schema) = composite_key_memtable(); + let mut scanner = + MemTableScanner::new(batch_store.clone(), indexes.clone(), schema.clone()); + scanner.with_pk_columns(vec!["a".to_string()]); + let query = arrow_array::Float32Array::from(vec![3.0, 0.0]); + scanner.nearest("vector", &query, 10).unwrap(); + let error = scanner.try_into_batch().await.unwrap_err().to_string(); + assert!(error.contains("key index"), "{error}"); + + let scanner = MemTableScanner::new(batch_store, indexes, schema); + let error = scanner + .create_dedup_plan(&["a".to_string()]) + .await + .unwrap_err() + .to_string(); + assert!(error.contains("key index"), "{error}"); + } + /// Timing for prefiltered searches over 100,000 rows of 128-dim vectors, /// with and without key rewrites. Run optimized with `--ignored --nocapture`. #[tokio::test]