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 514d5fadfe1..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); @@ -1318,18 +1323,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(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. + 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 } @@ -1347,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)?; } @@ -1361,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, @@ -3704,4 +3740,389 @@ 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}"); + } + } + } + } + } + + /// 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] + #[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] + ); + } + } + } } 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, ))) }