diff --git a/rust/lance/src/index/prefilter.rs b/rust/lance/src/index/prefilter.rs index 78dcbd3ad43..98e7bfa52ce 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,35 @@ 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 { + // 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?; + 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| { @@ -204,35 +207,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 +708,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,