diff --git a/rust/lance/src/io/exec/take.rs b/rust/lance/src/io/exec/take.rs index fea9e45b3ad..6921e20203d 100644 --- a/rust/lance/src/io/exec/take.rs +++ b/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::(); - 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::(); + 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, + ) -> (Arc, Option) { + 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::>() + .await + .unwrap(); + + let result = concat_batches(&result[0].schema(), &result).unwrap(); + assert_eq!( + result[ROW_ID] + .as_any() + .downcast_ref::() + .unwrap(), + &UInt64Array::from(vec![0_u64, 2]) + ); + assert_eq!( + result["score"] + .as_any() + .downcast_ref::() + .unwrap(), + &Int32Array::from(vec![10, 12]) + ); + assert_eq!( + result["s"].as_any().downcast_ref::().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;