From 46b0a0fd39edbd80d96745beb4eaa57376c96b67 Mon Sep 17 00:00:00 2001 From: Chakshu Dhannawat Date: Tue, 14 Jul 2026 17:30:18 +0900 Subject: [PATCH] fix: keep row ids and addresses aligned in filter_deleted_ids optimize_indices could fail with "batch.num_rows() != chunk.len()" after deleting rows from a stable-row-id dataset indexed with IVF_FLAT. When stable row ids are enabled, deleting rows also removes them from the row id index, so those ids no longer resolve to an address. filter_deleted_ids built the address list with a filter_map, which silently dropped the unresolvable ids and left the addresses vec shorter than the ids vec. It then handed the full ids vec plus the shortened addresses vec to filter_addr_or_ids, which relies on the two being positionally aligned (it sorts the addresses, checks each against its fragment's deletion vector, then zips the result back against the input ids). Once the lengths diverge the zip pairs ids with the wrong liveness verdict, so some already-deleted ids get reported as live. take_vectors then asks take_rows for those ids, gets back fewer rows than the chunk it requested, and errors out. Build ids and addresses together and unzip them so the two slices stay the same length and positionally correspond. Ids that no longer resolve are dropped from both, which is correct since those rows are already gone. Fixes #7701 --- rust/lance/src/dataset.rs | 26 ++++-- rust/lance/src/index/vector/builder.rs | 117 +++++++++++++++++++++++++ 2 files changed, 134 insertions(+), 9 deletions(-) diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index dc04095462a..f9bacf68748 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -2743,17 +2743,25 @@ impl Dataset { } pub(crate) async fn filter_deleted_ids(&self, ids: &[u64]) -> Result> { - let addresses = if let Some(row_id_index) = get_row_id_index(self).await? { - let addresses = ids + if let Some(row_id_index) = get_row_id_index(self).await? { + // Resolve each stable row id to its address. Some ids may no longer + // resolve (e.g. their rows were already removed), so we drop those + // and keep `ids` and `addresses` positionally aligned via unzip. + // `filter_addr_or_ids` relies on that correspondence; passing the + // full `ids` alongside a shorter, filter_map'd `addresses` would + // misalign the two and mis-filter the results. + let (ids, addresses): (Vec, Vec) = ids .iter() - .filter_map(|id| row_id_index.get(*id).map(|address| address.into())) - .collect::>(); - Cow::Owned(addresses) + .filter_map(|id| { + row_id_index + .get(*id) + .map(|address| (*id, u64::from(address))) + }) + .unzip(); + self.filter_addr_or_ids(&ids, &addresses).await } else { - Cow::Borrowed(ids) - }; - - self.filter_addr_or_ids(ids, &addresses).await + self.filter_addr_or_ids(ids, ids).await + } } /// Gets the number of files that are so small they don't even have a full diff --git a/rust/lance/src/index/vector/builder.rs b/rust/lance/src/index/vector/builder.rs index f881d28e745..5a53cfbd01d 100644 --- a/rust/lance/src/index/vector/builder.rs +++ b/rust/lance/src/index/vector/builder.rs @@ -2433,6 +2433,123 @@ mod tests { assert_eq!(total_rows, 19_500, "all vectors preserved across split"); } + #[tokio::test] + async fn optimize_after_delete_with_stable_row_ids_joins_partitions() { + // Regression test for https://github.com/lancedb/lance/issues/7701 + // + // With stable row ids enabled, deleting rows removes them from the + // row id index, so those ids no longer resolve to an address. When + // `optimize_indices` later joins undersized partitions it re-reads the + // stored (stable) row ids for a partition and asks the dataset to drop + // the deleted ones. If the surviving ids and their addresses get out of + // alignment, some deleted ids leak through and `take_rows` returns + // fewer rows than requested, tripping the + // `batch.num_rows() != chunk.len()` check. This exercises that path. + use crate::dataset::WriteParams; + use crate::index::DatasetIndexExt; + use crate::index::vector::VectorIndexParams; + use arrow_array::types::Int64Type; + use arrow_array::{Int64Array, RecordBatchIterator}; + use arrow_schema::Schema as ArrowSchema; + use lance_index::optimize::OptimizeOptions; + use lance_linalg::distance::MetricType; + + let dim = 32; + let num_rows = 5000usize; + let item_field = Arc::new(Field::new("item", DataType::Float32, true)); + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new( + "vec", + DataType::FixedSizeList(item_field, dim as i32), + false, + ), + ])); + + let mut values = Vec::with_capacity(num_rows * dim); + for i in 0..num_rows { + for j in 0..dim { + values.push(((i as f32) * 0.1 + (j as f32) * 0.3).sin()); + } + } + let fsl = FixedSizeListArray::try_new_from_values(Float32Array::from(values), dim as i32) + .unwrap(); + let ids = Int64Array::from((0..num_rows as i64).collect::>()); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(ids), Arc::new(fsl)]).unwrap(); + + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_str().unwrap(); + + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let mut dataset = crate::Dataset::write( + reader, + uri, + Some(WriteParams { + enable_stable_row_ids: true, + ..Default::default() + }), + ) + .await + .unwrap(); + + // 16 partitions over 5000 rows keeps every partition well under the + // IVF_FLAT join threshold (1024), so the optimize below takes the join + // path for every partition. + let params = VectorIndexParams::ivf_flat(16, MetricType::L2); + dataset + .create_index( + &["vec"], + IndexType::Vector, + Some("idx".into()), + ¶ms, + false, + ) + .await + .unwrap(); + + // Delete 1/5 of the rows, spread across every partition, so each + // partition carries some now-unresolvable stable row ids. + dataset.delete("id % 5 == 0").await.unwrap(); + + // Before the fix this panicked with `batch.num_rows() != chunk.len()`. + dataset + .optimize_indices(&OptimizeOptions::default()) + .await + .unwrap(); + + // The delete is intact and the data is untouched. + assert_eq!(dataset.count_rows(None).await.unwrap(), 4000); + + // The optimized index is still usable and only returns live rows. + let query = dataset + .scan() + .filter("id = 1") + .unwrap() + .project(&["vec"] as &[&str]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let q = query["vec"].as_fixed_size_list().value(0); + let result = dataset + .scan() + .project(&["id"] as &[&str]) + .unwrap() + .nearest("vec", q.as_ref(), 10) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert!( + result.num_rows() > 0, + "search over optimized index returned nothing" + ); + for id in result["id"].as_primitive::().values() { + assert_ne!(id % 5, 0, "search returned a deleted row id {id}"); + } + } + #[tokio::test] async fn take_partition_batches_preserves_partition_order_for_large_fixed_size_list() { let value_length = 1_073_741_824i32;