fix: filter stale physical row ids in take

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

@@ -38,6 +38,7 @@ use crate::dataset::Dataset;
 use crate::dataset::fragment::{FragReadConfig, FragmentReader};
 use crate::dataset::rowids::get_row_id_index;
 use crate::datatypes::Schema;
+use crate::index::prefilter::DatasetPreFilter;
 
 use super::utils::IoMetrics;
 
@@ -173,11 +174,12 @@ impl TakeStream {
     }
 
     /// Returns the row addresses for the given batch, plus an optional validity
-    /// mask. When stable row IDs are used, some row IDs from stale index results
-    /// (e.g. FTS matches for deleted rows) may no longer exist in the row ID
-    /// index. These are excluded from the returned addresses, and the mask
-    /// indicates which input rows are still valid so the caller can filter the
-    /// batch to match.
+    /// mask. Some row IDs from stale index results (e.g. FTS matches for deleted
+    /// rows) may no longer be valid. For stable row IDs, the row ID index detects
+    /// these entries. For physical row IDs, the dataset deletion mask detects
+    /// deleted rows and removed fragments. Invalid entries are excluded from the
+    /// returned addresses, and the mask indicates which input rows are still
+    /// valid so the caller can filter the batch to match.
     async fn get_row_addrs(
         &self,
         batch: &RecordBatch,
@@ -189,28 +191,52 @@ impl TakeStream {
 
             if let Some(row_id_index) = get_row_id_index(&self.dataset).await? {
                 let row_id_array = row_id_array.as_primitive::<UInt64Type>();
-                let mut addresses = Vec::with_capacity(row_id_array.len());
-                let mut valid = Vec::with_capacity(row_id_array.len());
-
-                for id in row_id_array.values().iter() {
-                    if let Some(address) = row_id_index.get(*id) {
-                        addresses.push(u64::from(address));
-                        valid.push(true);
-                    } else {
-                        valid.push(false);
-                    }
+                Ok(Self::resolve_row_addrs(row_id_array, |id| {
+                    row_id_index.get(id).map(u64::from)
+                }))
+            } else {
+                let row_id_array = row_id_array.as_primitive::<UInt64Type>();
+                let fragments = row_id_array
+                    .values()
+                    .iter()
+                    .map(|id| RowAddress::from(*id).fragment_id())
+                    .collect();
+                if let Some(mask) =
+                    DatasetPreFilter::create_deletion_mask(self.dataset.clone(), fragments)
+                {
+                    let mask = mask.await?;
+                    Ok(Self::resolve_row_addrs(row_id_array, |id| {
+                        mask.selected(id).then_some(id)
+                    }))
+                } else {
+                    Ok((Arc::new(row_id_array.clone()), None))
                 }
+            }
+        }
+    }
 
-                let mask = if addresses.len() < row_id_array.len() {
-                    Some(BooleanArray::from(valid))
-                } else {
-                    None
-                };
-                Ok((Arc::new(UInt64Array::from(addresses)), mask))
+    fn resolve_row_addrs(
+        row_ids: &UInt64Array,
+        mut resolve: impl FnMut(u64) -> Option<u64>,
+    ) -> (Arc<dyn Array>, Option<BooleanArray>) {
+        let mut addresses = Vec::with_capacity(row_ids.len());
+        let mut valid = Vec::with_capacity(row_ids.len());
+
+        for id in row_ids.values().iter() {
+            if let Some(address) = resolve(*id) {
+                addresses.push(address);
+                valid.push(true);
             } else {
-                Ok((row_id_array.clone(), None))
+                valid.push(false);
             }
         }
+
+        let mask = if addresses.len() < row_ids.len() {
+            Some(BooleanArray::from(valid))
+        } else {
+            None
+        };
+        (Arc::new(UInt64Array::from(addresses)), mask)
     }
 
     async fn map_batch(
@@ -221,9 +247,9 @@ impl TakeStream {
         let compute_timer = self.metrics.baseline_metrics.elapsed_compute().timer();
         let (row_addrs_arr, validity_mask) = self.get_row_addrs(&batch).await?;
 
-        // Filter out rows whose row IDs no longer exist (e.g. stale FTS/vector
-        // index entries pointing to deleted rows). Without this, the downstream
-        // merge would fail with a row-count mismatch.
+        // Filter stale index entries before reading so the input batch and taken
+        // columns remain aligned. Otherwise, the downstream merge would fail with
+        // a row-count mismatch.
         let batch = if let Some(mask) = validity_mask {
             arrow::compute::filter_record_batch(&batch, &mask)?
         } else {
@@ -873,6 +899,70 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn test_take_filters_stale_physical_row_ids() {
+        let TestFixture {
+            dataset,
+            _tmp_dir_guard,
+        } = test_fixture().await;
+        let mut dataset = dataset.as_ref().clone();
+        dataset.delete("i = 1").await.unwrap();
+        let dataset = Arc::new(dataset);
+
+        // Simulate stale index results for a deleted row and a removed fragment.
+        let missing_fragment_row_id = u64::from(RowAddress::new_from_parts(99, 0));
+        let row_ids = Arc::new(UInt64Array::from(vec![
+            0_u64,
+            1,
+            2,
+            missing_fragment_row_id,
+        ]));
+        let scores = Arc::new(Int32Array::from(vec![10, 11, 12, 13]));
+        let input_batch = RecordBatch::try_from_iter(vec![
+            (ROW_ID, row_ids as ArrayRef),
+            ("score", scores as ArrayRef),
+        ])
+        .unwrap();
+        let schema = input_batch.schema();
+        let input_stream = futures::stream::iter(vec![Ok(input_batch)]);
+        let input_stream = Box::pin(RecordBatchStreamAdapter::new(schema, input_stream));
+        let input = Arc::new(OneShotExec::new(input_stream));
+
+        let projection = dataset
+            .empty_projection()
+            .union_column("s", OnMissing::Error)
+            .unwrap();
+        let take_exec = TakeExec::try_new(dataset, input, projection)
+            .unwrap()
+            .unwrap();
+        let result = take_exec
+            .execute(0, Arc::new(TaskContext::default()))
+            .unwrap()
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+
+        let result = concat_batches(&result[0].schema(), &result).unwrap();
+        assert_eq!(
+            result[ROW_ID]
+                .as_any()
+                .downcast_ref::<UInt64Array>()
+                .unwrap(),
+            &UInt64Array::from(vec![0_u64, 2])
+        );
+        assert_eq!(
+            result["score"]
+                .as_any()
+                .downcast_ref::<Int32Array>()
+                .unwrap(),
+            &Int32Array::from(vec![10, 12])
+        );
+        assert_eq!(
+            result["s"].as_any().downcast_ref::<StringArray>().unwrap(),
+            &StringArray::from(vec!["str-0", "str-2"])
+        );
+    }
+
     #[tokio::test(flavor = "current_thread")]
     async fn test_take_records_output_and_io_metrics() {
         use datafusion::physical_plan::metrics::MetricValue;