From 1647231fe4a477295b42973a0fa54134275f0536 Mon Sep 17 00:00:00 2001 From: JaySon-Huang Date: Sun, 30 Aug 2026 20:20:26 +0800 Subject: [PATCH 1/2] perf: assemble stable-row-id prefilter allow lists from cached per-fragment pieces do_create_deletion_mask_row_id builds a whole-dataset allow list by loading every fragment's row id sequence and deletion vector, so the cost is proportional to the dataset size rather than to what changed. The result is cached under (manifest version, restrict set), which amortizes less than it looks: every commit invalidates the entry, distinct restrict sets each pay their own full rebuild, and the resident mask scales with the dataset. Assemble the allow list from per-fragment pieces instead. Each piece is one fragment's live stable row ids (row id sequence minus deletion vector), cached under its content identity: the row id generation (the same identity RowIdSequenceKey uses for generation-safe sequence caching) plus the deletion file identity. Pieces are OR-ed into the requested mask: - commits that leave a fragment's row ids and deletions untouched keep its piece warm, so a post-commit rebuild only reloads touched fragments - overlapping restrict sets share pieces instead of reloading them - content is unchanged: a piece is exactly what the old fold OR-ed in for that fragment, and missing fragments contribute nothing The whole-mask cache under (version, restrict_hash) is kept, so repeated identical queries still hit it unchanged. --- rust/lance/src/index/prefilter.rs | 193 ++++++++++++++++++++++-------- rust/lance/src/session/caches.rs | 69 ++++++++++- 2 files changed, 209 insertions(+), 53 deletions(-) diff --git a/rust/lance/src/index/prefilter.rs b/rust/lance/src/index/prefilter.rs index 78dcbd3ad43..1b480dfbaab 100644 --- a/rust/lance/src/index/prefilter.rs +++ b/rust/lance/src/index/prefilter.rs @@ -6,7 +6,6 @@ //! Based on the query, we might have information about which fragment ids and //! row ids can be excluded from the search. -use std::borrow::Cow; use std::cell::OnceCell; use std::collections::HashMap; use std::sync::Arc; @@ -18,12 +17,10 @@ use futures::StreamExt; use futures::TryStreamExt; use futures::future::BoxFuture; use futures::stream; -use lance_core::utils::deletion::DeletionVector; use lance_core::utils::tokio::spawn_cpu; use lance_select::{RowAddrMask, RowAddrTreeMap}; use lance_table::format::Fragment; use lance_table::format::IndexMetadata; -use lance_table::rowids::RowIdSequence; use roaring::RoaringBitmap; use tokio::join; use tracing::Instrument; @@ -159,29 +156,28 @@ impl DatasetPreFilter { // twice — once via the BTREE (which holds the row's stable_row_id) and // once via the unindexed scan (which holds the fragment the row now // lives in). See issue #6877. - async fn load_row_ids_and_deletions( - dataset: &Dataset, - restrict_to: Option<&RoaringBitmap>, - ) -> Result, Option>)>> { - let frags: Vec<_> = dataset - .get_fragments() - .into_iter() - .filter(|f| { - restrict_to - .map(|allow| allow.contains(f.id() as u32)) - .unwrap_or(true) - }) - .collect(); - stream::iter(frags) - .map(|frag| async move { - let row_ids = load_row_id_sequence(dataset, frag.metadata()); - let deletion_vector = frag.get_deletion_vector(); - let (row_ids, deletion_vector) = join!(row_ids, deletion_vector); - Ok::<_, crate::Error>((row_ids?, deletion_vector?)) - }) - .buffer_unordered(dataset.object_store.as_ref().io_parallelism()) - .try_collect::>() - .await + // + // The allow list is assembled from per-fragment pieces, each cached + // under its content identity (row id generation + deletion file), so + // repeated builds with overlapping `restrict_to` sets reload only the + // fragments they have not seen yet, and commits that leave a fragment + // untouched keep its piece warm. + async fn build_allow_list_piece( + dataset: Arc, + frag: FileFragment, + ) -> Result { + let row_ids = load_row_id_sequence(&dataset, frag.metadata()); + let deletion_vector = frag.get_deletion_vector(); + let (row_ids, deletion_vector) = join!(row_ids, deletion_vector); + let row_ids = row_ids?; + match deletion_vector? { + Some(deletion_vector) => { + let mut row_ids = row_ids.as_ref().clone(); + row_ids.mask(deletion_vector.to_sorted_iter())?; + Ok(RowAddrTreeMap::from(&row_ids)) + } + None => Ok(RowAddrTreeMap::from(row_ids.as_ref())), + } } let restrict_hash = restrict_to.as_ref().map(|b| { @@ -204,35 +200,59 @@ impl DatasetPreFilter { dataset .metadata_cache .as_ref() - .get_or_insert_with_key(key, move || { - async move { - let row_ids_and_deletions = - load_row_ids_and_deletions(&dataset_clone, restrict_for_load.as_ref()) - .await?; - - // The process of computing the final mask is CPU-bound, so we spawn it - // on a blocking thread. - let allow_list = spawn_cpu(move || { - Result::Ok(row_ids_and_deletions.into_iter().fold( - RowAddrTreeMap::new(), - |mut allow_list, (row_ids, deletion_vector)| { - let seq = if let Some(deletion_vector) = deletion_vector { - let mut row_ids = row_ids.as_ref().clone(); - row_ids.mask(deletion_vector.to_sorted_iter()).unwrap(); - Cow::::Owned(row_ids) - } else { - Cow::::Borrowed(row_ids.as_ref()) - }; - let treemap = RowAddrTreeMap::from(seq.as_ref()); - allow_list |= treemap; - allow_list - }, - )) + .get_or_insert_with_key(key, move || async move { + let fragments: Vec = dataset_clone + .get_fragments() + .into_iter() + .filter(|frag| { + restrict_for_load + .as_ref() + .map(|allow| allow.contains(frag.id() as u32)) + .unwrap_or(true) }) + .collect(); + + let pieces = stream::iter(fragments) + .map(|frag| { + let dataset = dataset_clone.clone(); + async move { + let meta = frag.metadata(); + let row_id_meta = meta.row_id_meta.clone().ok_or_else(|| { + crate::Error::internal(format!( + "fragment {} is missing row id meta", + meta.id + )) + })?; + let piece_key = crate::session::caches::RowIdAllowListPieceKey { + fragment_id: meta.id, + row_id_meta, + deletion_file: meta.deletion_file.clone(), + }; + dataset + .metadata_cache + .as_ref() + .get_or_insert_with_key(piece_key, || { + build_allow_list_piece(dataset.clone(), frag) + }) + .await + } + }) + .buffer_unordered(dataset_clone.object_store.as_ref().io_parallelism()) + .try_collect::>() .await?; - Ok(RowAddrMask::from_allowed(allow_list)) - } + // Merging the pieces is CPU-bound, so we spawn it on a + // blocking thread. + let allow_list = spawn_cpu(move || { + let mut allow_list = RowAddrTreeMap::new(); + for piece in &pieces { + allow_list |= piece.as_ref(); + } + Result::Ok(allow_list) + }) + .await?; + + Ok(RowAddrMask::from_allowed(allow_list)) }) .await } @@ -681,4 +701,73 @@ mod test { .unwrap(); assert_eq!(mask.allow_list().and_then(|x| x.len()), Some(0)); } + + // The allow list is assembled from per-fragment pieces cached by content + // identity (row id generation + deletion file). A commit that deletes rows + // must invalidate exactly the touched fragments' pieces: a mask rebuilt in + // the same session must reflect the new deletions while untouched + // fragments are served from cache. + #[tokio::test] + async fn test_row_id_allow_list_pieces_invalidate_on_new_deletions() { + let test_data = BatchGenerator::new() + .col(Box::new(IncrementingInt32::new().named("x"))) + .batch(9); + let mut dataset = Dataset::write( + test_data, + "memory://test_pieces_invalidate", + Some(WriteParams { + max_rows_per_file: 3, + enable_stable_row_ids: true, + ..Default::default() + }), + ) + .await + .unwrap(); + + // Three fragments of three rows each; stable row ids are 0..8, and + // fragment k holds ids 3k..3k+2. Delete x=8 (fragment 2) and build + // the allow list cold. + dataset.delete("x = 8").await.unwrap(); + let ds = Arc::new(dataset.clone()); + let mask = DatasetPreFilter::create_deletion_mask( + ds.clone(), + RoaringBitmap::from_iter(0..3), + ) + .expect("mask present") + .await + .unwrap(); + assert_eq!(mask.allow_list().and_then(|x| x.len()), Some(8)); + + // Delete x=2 (fragment 0) and rebuild with the same session. The + // rebuilt mask must exclude the newly deleted row and keep everything + // else, including the row ids whose home fragment was untouched by + // this commit. + dataset.delete("x = 2").await.unwrap(); + let ds = Arc::new(dataset); + let mask = DatasetPreFilter::create_deletion_mask( + ds.clone(), + RoaringBitmap::from_iter(0..3), + ) + .expect("mask present") + .await + .unwrap(); + assert_eq!(mask.allow_list().and_then(|x| x.len()), Some(7)); + assert!(!mask.selected(2)); + assert!(mask.selected(0)); + assert!(mask.selected(3)); + assert!(mask.selected(7)); + + // The restricted variant must honor its bitmap restriction against + // the rebuilt pieces: fragments {1, 2} hold ids 3..8 minus the + // deleted 8. + let mask = DatasetPreFilter::create_restricted_deletion_mask( + ds.clone(), + RoaringBitmap::from_iter(1..3), + ) + .expect("restricted mask present") + .await + .unwrap(); + let expected = RowAddrTreeMap::from_iter(3..8); + assert_eq!(mask.allow_list(), Some(&expected)); + } } diff --git a/rust/lance/src/session/caches.rs b/rust/lance/src/session/caches.rs index 330dbdbd0b0..3aba650d841 100644 --- a/rust/lance/src/session/caches.rs +++ b/rust/lance/src/session/caches.rs @@ -17,7 +17,7 @@ use lance_core::{ cache::{CacheKey, CacheKeySchema, KeyBuilder, LanceCache}, utils::deletion::DeletionVector, }; -use lance_select::RowAddrMask; +use lance_select::{RowAddrMask, RowAddrTreeMap}; use lance_table::{ format::{DeletionFile, DeletionFileType, Manifest, RowIdMeta}, rowids::{RowIdIndex, RowIdSequence}, @@ -197,6 +197,73 @@ impl CacheKey for RowAddrMaskKey { } } +/// One fragment's live stable-row-id set (an allow-list "piece") for +/// stable-row-id datasets. +/// +/// [`RowAddrMaskKey`] is keyed by manifest version, so every commit +/// invalidates the whole-dataset mask and every distinct `restrict_to` set +/// pays its own full rebuild (lance-format/lance#8849). This key identifies +/// the piece by the fragment's *content* instead: the row id generation +/// (same identity [`RowIdSequenceKey`] uses, #7645) plus the deletion file +/// identity. Commits that leave a fragment's row ids and deletions untouched +/// keep hitting the same piece, so a whole-dataset allow list is assembled by +/// OR-ing cached pieces and only fragments whose content changed reload. +#[derive(Debug)] +pub struct RowIdAllowListPieceKey { + pub fragment_id: u64, + pub row_id_meta: RowIdMeta, + pub deletion_file: Option, +} + +impl CacheKey for RowIdAllowListPieceKey { + type ValueType = RowAddrTreeMap; + // Only the legacy display form. Identity comes from `write_key` below. + fn key(&self) -> Cow<'_, str> { + Cow::Owned(format!("row_id_allow_list_piece/{}", self.fragment_id)) + } + fn type_name() -> &'static str { + "RowAddrTreeMap" + } + + fn schema() -> CacheKeySchema { + CacheKeySchema::new("lance.dataset.row-id-allow-list-piece-key", 1) + } + + fn write_key(&self, builder: &mut KeyBuilder) { + builder.write_u64(self.fragment_id); + match &self.row_id_meta { + RowIdMeta::Inline(data) => { + builder.write_variant(0); + builder.write_fixed_bytes(data.digest()); + } + RowIdMeta::External(file) => { + builder.write_variant(1); + builder.write_str(&file.path); + builder.write_u64(file.offset); + builder.write_u64(file.size); + } + } + match &self.deletion_file { + None => builder.write_variant(0), + Some(deletion_file) => { + builder.write_variant(1); + builder.write_u64(deletion_file.read_version); + builder.write_u64(deletion_file.id); + builder.write_variant(match deletion_file.file_type { + DeletionFileType::Array => 0, + DeletionFileType::Bitmap => 1, + }); + if let Some(base_id) = deletion_file.base_id { + builder.write_some(); + builder.write_u32(base_id); + } else { + builder.write_none(); + } + } + } + } +} + #[derive(Debug)] pub struct RowIdIndexKey { pub version: u64, From ac85e697c4a79b216268f4c5c36da27ac943d221 Mon Sep 17 00:00:00 2001 From: JaySon-Huang Date: Sun, 30 Aug 2026 20:56:07 +0800 Subject: [PATCH 2/2] fix: run per-piece allow-list construction on the CPU pool Loading a fragment's row id sequence and deletion vector stays async, but the sequence clone, deletion masking and tree construction for each piece are CPU-bound and were running inline on the async runtime. Move them into spawn_cpu, matching the whole-mask fold's original async/CPU separation while keeping the per-piece pipelining: each piece loads on the IO path and builds on the CPU pool independently, so IO of one piece overlaps CPU of another. The final merge stays on the CPU pool. --- rust/lance/src/index/prefilter.rs | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/rust/lance/src/index/prefilter.rs b/rust/lance/src/index/prefilter.rs index 1b480dfbaab..98e7bfa52ce 100644 --- a/rust/lance/src/index/prefilter.rs +++ b/rust/lance/src/index/prefilter.rs @@ -166,18 +166,25 @@ impl DatasetPreFilter { dataset: Arc, frag: FileFragment, ) -> Result { + // Load the row id sequence and deletion vector asynchronously, + // then run the CPU-bound clone, deletion masking and tree + // construction for each piece on the CPU pool, keeping the async + // runtime free (the whole-mask fold previously ran this work in + // one spawn_cpu block after all loads). let row_ids = load_row_id_sequence(&dataset, frag.metadata()); let deletion_vector = frag.get_deletion_vector(); let (row_ids, deletion_vector) = join!(row_ids, deletion_vector); let row_ids = row_ids?; - match deletion_vector? { + let deletion_vector = deletion_vector?; + spawn_cpu(move || match deletion_vector { Some(deletion_vector) => { let mut row_ids = row_ids.as_ref().clone(); row_ids.mask(deletion_vector.to_sorted_iter())?; Ok(RowAddrTreeMap::from(&row_ids)) } None => Ok(RowAddrTreeMap::from(row_ids.as_ref())), - } + }) + .await } let restrict_hash = restrict_to.as_ref().map(|b| {