fix: filter stale scalar index fragments

rust/lance/src/dataset/write/merge_insert.rs

@@ -10429,9 +10429,10 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n
     // in a fragment) followed by partial merge_insert should not produce
     // "fragment id N does not exist" errors.
     //
-    // The bug: stale btree entries reference the deleted fragment. The deletion mask doesn't
-    // block those addresses because the fragment isn't in the index bitmap. TakeExec tries
-    // to read from a non-existent fragment.
+    // The bug: stale btree entries reference the deleted fragment. A deletion in a covered
+    // fragment forces the restricted prefilter to use a block list. That block list must also
+    // block historical fragments outside the index bitmap, or TakeExec tries to read from a
+    // non-existent fragment.
     #[tokio::test]
     async fn test_partial_merge_insert_stale_index_fragment_not_exist() {
         let dataset = create_indexed_3frag_dataset().await;
@@ -10451,11 +10452,15 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n
             .execute()
             .await
             .unwrap();
-        let dataset = update_result.new_dataset;
+        let mut dataset = (*update_result.new_dataset).clone();
+
+        // Step 4: Delete a row from fragment 0, which is still covered by the index.
+        // This forces the restricted deletion mask's block-list branch.
+        dataset.delete("id = 'id-0000'").await.unwrap();
 
-        // Step 4: Partial merge_insert on the same rows.
+        // Step 5: Partial merge_insert on the same rows.
         // This should succeed, not fail with "fragment does not exist".
-        let dataset = partial_merge_insert(dataset, 100..200, 888.0).await;
+        let dataset = partial_merge_insert(Arc::new(dataset), 100..200, 888.0).await;
 
         // Verify correctness
         let batches = dataset
@@ -10473,7 +10478,7 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n
             Field::new("value_b", DataType::Float64, false),
         ]));
         let combined = concat_batches(&all_schema, &batches).unwrap();
-        assert_eq!(combined.num_rows(), 300);
+        assert_eq!(combined.num_rows(), 299);
     }
 
     // Regression test: partial-schema merge_insert followed by update (deleting SOME rows
@@ -10638,21 +10643,24 @@ MergeInsert: on=[id], when_matched=DoNothing, when_not_matched=InsertAll, when_n
             .await
             .unwrap();
 
-        // Check no duplicate ids
+        // Every logical row should be returned exactly once. Checking only for
+        // duplicates is insufficient because stale addresses can be discarded by
+        // TakeExec after they have displaced valid candidates from the top-k.
         let ids = results
             .column_by_name("id")
             .unwrap()
             .as_any()
             .downcast_ref::<StringArray>()
             .unwrap();
-        let unique_ids: std::collections::HashSet<&str> =
+        let actual_ids: std::collections::HashSet<&str> =
             (0..ids.len()).map(|i| ids.value(i)).collect();
+        let expected_ids: std::collections::HashSet<String> =
+            (0..total_rows).map(|i| format!("id-{i:04}")).collect();
+        assert_eq!(ids.len(), total_rows);
+        assert_eq!(actual_ids.len(), total_rows);
         assert_eq!(
-            unique_ids.len(),
-            ids.len(),
-            "Found duplicate ids in KNN results: {} unique out of {} total",
-            unique_ids.len(),
-            ids.len()
+            actual_ids,
+            expected_ids.iter().map(String::as_str).collect(),
         );
     }
 

rust/lance/src/index/prefilter.rs

@@ -52,6 +52,10 @@ pub struct DatasetPreFilter {
     // Fragment IDs whose data is still in the index but has been removed from the dataset.
     // Used by FTS merge-on-read to prune stale fragments at search time.
     pub(super) deleted_fragments: Option<RoaringBitmap>,
+    // Physical fragments outside every index bitmap. Store these as full-fragment
+    // blocks so mask-based consumers enforce index coverage without expanding
+    // partially deleted fragments into full-fragment allow lists.
+    pub(super) excluded_fragments: Option<RoaringBitmap>,
     // When the tasks are finished this is the combined filter
     pub(super) final_mask: Mutex<OnceCell<Arc<RowAddrMask>>>,
 }
@@ -71,18 +75,30 @@ impl DatasetPreFilter {
                 fragments |= idx.fragment_bitmap.as_ref().unwrap();
             });
         }
-        let deleted_ids = if all_have_bitmaps {
-            Self::create_restricted_deletion_mask(dataset, fragments)
+        let uses_stable_row_ids = dataset.manifest.uses_stable_row_ids();
+        let deleted_ids = if all_have_bitmaps && uses_stable_row_ids {
+            Self::create_restricted_deletion_mask(dataset.clone(), fragments.clone())
         } else {
-            Self::create_deletion_mask(dataset, fragments)
+            Self::create_deletion_mask(dataset.clone(), fragments.clone())
         }
         .map(SharedPrerequisite::spawn);
+        let excluded_fragments = if all_have_bitmaps && !uses_stable_row_ids {
+            let mut historical_fragments = RoaringBitmap::new();
+            if let Some(max_fragment_id) = dataset.manifest.max_fragment_id() {
+                historical_fragments.insert_range(0..=max_fragment_id as u32);
+            }
+            historical_fragments -= &fragments;
+            (!historical_fragments.is_empty()).then_some(historical_fragments)
+        } else {
+            None
+        };
         let filtered_ids = filter
             .map(|filtered_ids| SharedPrerequisite::spawn(filtered_ids.load().in_current_span()));
         Self {
             deleted_ids,
             filtered_ids,
             deleted_fragments: None,
+            excluded_fragments,
             final_mask: Mutex::new(OnceCell::new()),
         }
     }
@@ -276,20 +292,12 @@ impl DatasetPreFilter {
         // outside the index bitmap. This can happen when a fragment's data was
         // modified but the index was not rewritten (e.g. after DataReplacement
         // or partial merge_insert).
-        //
-        // We materialize the set of non-index fragments here so the slow path
-        // below can fold them into the deletion mask as Full blocks instead of
-        // computing AllowList(Full) - BlockList(Partial), which forces
-        // RoaringBitmap::full() per fragment in RowAddrTreeMap::sub_assign and
-        // is the dominant cost on every merge_insert call.
-        let non_index_frags: Option<RoaringBitmap> = if restrict_to_fragments {
+        let needs_allow_list = if restrict_to_fragments {
             let dataset_frag_ids: RoaringBitmap = frag_map.keys().copied().collect();
-            let outside = &dataset_frag_ids - &fragments;
-            (!outside.is_empty()).then_some(outside)
+            !dataset_frag_ids.is_subset(&fragments)
         } else {
-            None
+            false
         };
-        let needs_allow_list = non_index_frags.is_some();
         for frag_id in fragments.iter() {
             let frag = frag_map.get(&frag_id);
             if let Some(frag) = frag {
@@ -318,29 +326,23 @@ impl DatasetPreFilter {
             }
             Some(async move { Ok(Arc::new(RowAddrMask::from_allowed(allow_list))) }.boxed())
         } else {
-            // There are deletions/missing frags. Build the deletion mask and,
-            // if needed, fold the non-index fragments into it as Full blocks.
-            // Equivalent to BlockList(deletions) | BlockList(non_index) — but
-            // expressed via insert_fragment so we never materialize a
-            // RoaringBitmap::full() per fragment.
+            // There are deletions/missing frags. Build the deletion mask and
+            // optionally combine it with the fragment allow-list. Hot query
+            // paths keep these filters separate to avoid materializing full
+            // roaring bitmaps for partially deleted fragments.
             let fut =
                 Self::do_create_deletion_mask(dataset, missing_frags, frags_with_deletion_files);
-            if let Some(non_index_frags) = non_index_frags {
+            if needs_allow_list {
                 Some(
                     async move {
                         let deletion_mask = fut.await?;
-                        let mut combined = match &*deletion_mask {
-                            RowAddrMask::BlockList(b) => b.clone(),
-                            RowAddrMask::AllowList(_) => {
-                                // do_create_deletion_mask only returns BlockList; this is
-                                // defensive and should be unreachable.
-                                return Ok(deletion_mask);
-                            }
-                        };
-                        for frag_id in non_index_frags.iter() {
-                            combined.insert_fragment(frag_id);
+                        let mut allow_list = RowAddrTreeMap::new();
+                        for frag_id in fragments.iter() {
+                            allow_list.insert_fragment(frag_id);
                         }
-                        Ok(Arc::new(RowAddrMask::from_block(combined)))
+                        Ok(Arc::new(
+                            (*deletion_mask).clone() & RowAddrMask::from_allowed(allow_list),
+                        ))
                     }
                     .boxed(),
                 )
@@ -376,13 +378,16 @@ impl PreFilter for DatasetPreFilter {
             if let Some(deleted_ids) = &self.deleted_ids {
                 combined = combined & (*deleted_ids.get_ready()).clone();
             }
-            if let Some(deleted) = &self.deleted_fragments {
-                let mut block_list = RowAddrTreeMap::new();
-                for frag_id in deleted.iter() {
-                    block_list.insert_fragment(frag_id);
-                }
-                combined = combined & RowAddrMask::from_block(block_list);
+            let blocked_fragments = self
+                .deleted_fragments
+                .iter()
+                .chain(&self.excluded_fragments)
+                .flat_map(|fragments| fragments.iter());
+            let mut fragment_block_list = RowAddrTreeMap::new();
+            for fragment_id in blocked_fragments {
+                fragment_block_list.insert_fragment(fragment_id);
             }
+            combined = combined & RowAddrMask::from_block(fragment_block_list);
             Arc::new(combined)
         });
 
@@ -393,6 +398,7 @@ impl PreFilter for DatasetPreFilter {
         self.deleted_ids.is_none()
             && self.filtered_ids.is_none()
             && self.deleted_fragments.is_none()
+            && self.excluded_fragments.is_none()
     }
 
     /// Get the row id mask for this prefilter
@@ -421,6 +427,7 @@ impl PreFilter for DatasetPreFilter {
 mod test {
     use lance_select::RowSetOps;
     use lance_testing::datagen::{BatchGenerator, IncrementingInt32};
+    use uuid::Uuid;
 
     use crate::dataset::WriteParams;
 
@@ -516,6 +523,20 @@ mod test {
         let mask = mask.unwrap().await.unwrap();
         assert_eq!(mask.block_list().and_then(|x| x.len()), Some(1));
 
+        // A restricted mask must block every historical fragment outside the
+        // bitmap, including fragment 1, which has been removed from the dataset.
+        let mask = DatasetPreFilter::create_restricted_deletion_mask(
+            datasets.deletions_missing_frags.clone(),
+            RoaringBitmap::from_iter(2..3),
+        )
+        .unwrap()
+        .await
+        .unwrap();
+        assert!(mask.selected(2 << 32));
+        assert!(!mask.selected((2 << 32) + 2));
+        assert!(!mask.selected(0));
+        assert!(!mask.selected(1 << 32));
+
         // If there are only missing fragments, we should still get a mask
         let mask = DatasetPreFilter::create_deletion_mask(
             datasets.only_missing_frags.clone(),
@@ -529,6 +550,33 @@ mod test {
         assert_eq!(mask.block_list(), Some(&expected));
     }
 
+    #[tokio::test]
+    async fn test_prefilter_mask_honors_legacy_index_fragment_bitmap() {
+        let datasets = test_datasets(false).await;
+        let mut dataset = (*datasets.no_deletions).clone();
+        // Older manifests derive the high-water mark from the fragment list.
+        Arc::make_mut(&mut dataset.manifest).max_fragment_id = None;
+        let index = IndexMetadata {
+            uuid: Uuid::new_v4(),
+            fields: vec![0],
+            name: "legacy_vector_index".to_string(),
+            dataset_version: dataset.manifest.version,
+            fragment_bitmap: Some(RoaringBitmap::from_iter([1])),
+            index_details: None,
+            index_version: 0,
+            created_at: None,
+            base_id: None,
+            files: None,
+        };
+        let prefilter = DatasetPreFilter::new(Arc::new(dataset), &[index], None);
+        prefilter.wait_for_ready().await.unwrap();
+
+        let mask = prefilter.mask();
+        assert!(!mask.selected(0));
+        assert!(mask.selected(1 << 32));
+        assert!(!mask.selected(2 << 32));
+    }
+
     #[tokio::test]
     async fn test_deletion_mask_stable_row_id() {
         // Here, behavior is different.

rust/lance/src/index/vector/fixture_test.rs

@@ -278,6 +278,7 @@ mod test {
                     deleted_ids: None,
                     filtered_ids: None,
                     deleted_fragments: None,
+                    excluded_fragments: None,
                     final_mask: Mutex::new(OnceCell::new()),
                 }),
                 &NoOpMetricsCollector,

rust/lance/src/io/exec/scalar_index.rs

@@ -490,20 +490,29 @@ impl MapIndexExec {
             });
         }
         let fragment_bitmap = fragment_bitmap.expect("MapIndexExec built with no lookups");
-        let deletion_mask_fut =
-            DatasetPreFilter::create_restricted_deletion_mask(dataset.clone(), fragment_bitmap);
+        let uses_stable_row_ids = dataset.manifest.uses_stable_row_ids();
+        let deletion_mask_fut = if uses_stable_row_ids {
+            DatasetPreFilter::create_restricted_deletion_mask(
+                dataset.clone(),
+                fragment_bitmap.clone(),
+            )
+        } else {
+            DatasetPreFilter::create_deletion_mask(dataset.clone(), fragment_bitmap.clone())
+        };
         let deletion_mask = if let Some(fut) = deletion_mask_fut {
             Some(fut.await?)
         } else {
             None
         };
+        let allowed_fragments = (!uses_stable_row_ids).then_some(Arc::new(fragment_bitmap));
 
         let baseline = BaselineMetrics::new(&metrics_set, partition);
         let elapsed_compute = baseline.elapsed_compute().clone();
         let stream = input.then(move |batch_result| {
             let lookups = lookups.clone();
             let dataset = dataset.clone();
             let deletion_mask = deletion_mask.clone();
+            let allowed_fragments = allowed_fragments.clone();
             let metrics = index_metrics.clone();
             let elapsed_compute = elapsed_compute.clone();
             async move {
@@ -512,7 +521,15 @@ impl MapIndexExec {
                 // the per-batch sargable index evaluation, which is the work
                 // we want attributed here.
                 let _t = elapsed_compute.timer();
-                Self::map_batch(lookups, dataset, deletion_mask, batch, metrics).await
+                Self::map_batch(
+                    lookups,
+                    dataset,
+                    deletion_mask,
+                    allowed_fragments,
+                    batch,
+                    metrics,
+                )
+                .await
             }
         });
         let stream = stream.map(move |batch| {
@@ -560,6 +577,7 @@ impl MapIndexExec {
         lookups: Vec<IndexLookup>,
         dataset: Arc<Dataset>,
         deletion_mask: Option<Arc<RowAddrMask>>,
+        allowed_fragments: Option<Arc<RoaringBitmap>>,
         batch: RecordBatch,
         metrics: Arc<IndexMetrics>,
     ) -> datafusion::error::Result<RecordBatch> {
@@ -570,6 +588,11 @@ impl MapIndexExec {
         }
         let mut row_addr_mask = query_result.upper;
 
+        if let (Some(allowed_fragments), RowAddrMask::AllowList(candidates)) =
+            (allowed_fragments.as_ref(), &mut row_addr_mask)
+        {
+            candidates.retain_fragments(allowed_fragments.iter());
+        }
         if let Some(deletion_mask) = deletion_mask.as_ref() {
             row_addr_mask = row_addr_mask & deletion_mask.as_ref().clone();
         }