diff --git a/rust/lance/src/dataset/rowids.rs b/rust/lance/src/dataset/rowids.rs index 082b041d40a..bb934787a84 100644 --- a/rust/lance/src/dataset/rowids.rs +++ b/rust/lance/src/dataset/rowids.rs @@ -90,12 +90,12 @@ pub async fn get_row_id_index( /// Map a set of physical row addresses to their stable row ids /// -/// For each fragment present in `addrs`, the live rows in physical order carry -/// the stable ids yielded by the fragment's [`RowIdSequence`] in the same -/// order. Zipping the two (skipping deleted physical offsets) gives the -/// `physical offset -> stable id` mapping. Addresses that point at deleted rows -/// have no live counterpart and are dropped, which is correct: those rows are -/// not part of the answer. +/// A fragment's [`RowIdSequence`] holds one id per physical row in offset order; +/// deletions are tracked by the deletion vector and do not compact the sequence, +/// so a physical offset indexes the sequence directly (see [`RowIdIndex`], which +/// maps ids to `start_address + position` and filters deletions separately). +/// Addresses that point at deleted rows have no live counterpart and are dropped, +/// which is correct: those rows are not part of the answer. pub(crate) async fn translate_addr_treemap_to_row_ids( dataset: &Dataset, addrs: &RowAddrTreeMap, @@ -116,55 +116,30 @@ pub(crate) async fn translate_addr_treemap_to_row_ids( row_ids |= RowAddrTreeMap::from(sequence.as_ref()); } RowAddrSelection::Partial(offsets) => { - let Some(max_offset) = offsets.max() else { - continue; - }; - let (deletion_vector, num_physical_rows) = futures::try_join!( - file_fragment.get_deletion_vector(), - file_fragment.physical_rows() - )?; - let num_physical_rows = u32::try_from(num_physical_rows).map_err(|_| { - Error::internal(format!( - "fragment_id={fragment_id} has num_physical_rows={num_physical_rows}, \ - which exceeds the maximum representable physical offset" - )) - })?; - if max_offset >= num_physical_rows { - return Err(Error::internal(format!( - "fragment_id={fragment_id} selection has max_offset={max_offset}, \ - but num_physical_rows={num_physical_rows}" - ))); - } - let mut ids = sequence.iter(); - for physical_offset in 0..num_physical_rows { - if physical_offset > max_offset { - break; - } + let deletion_vector = file_fragment.get_deletion_vector().await?; + for physical_offset in offsets.iter() { + // A deletion does not compact the row id sequence; the deleted row keeps + // its slot (the deletion vector tracks it separately). So a physical offset + // is a direct index into the sequence, regardless of any deletions below it. + // A stale offset that points at a deleted row has no live counterpart and is + // not part of any answer, so it contributes no id to the block/take set. let deleted = deletion_vector .as_ref() .is_some_and(|dv| dv.contains(physical_offset)); if deleted { continue; } - match ids.next() { - Some(id) => { - if offsets.contains(physical_offset) { - row_ids.insert(id); - } - } - // The sequence yields one id per live row, so it can only - // run dry before `max_offset` if it holds fewer ids than the - // fragment has live rows. Breaking would silently drop the - // remaining selected offsets and let stale index results - // escape masking, so treat the mismatch as corruption. - None => { - return Err(Error::internal(format!( - "fragment_id={fragment_id} row-id sequence exhausted at \ - physical_offset={physical_offset} before reaching \ - max_offset={max_offset} (num_physical_rows={num_physical_rows})" - ))); - } - } + // A selected offset with no sequence entry points past the fragment's rows. + // Silently dropping it would let a stale index result escape masking, so + // treat the mismatch as corruption. + let id = sequence.get(physical_offset as usize).ok_or_else(|| { + Error::internal(format!( + "fragment_id={fragment_id} row-id sequence has no entry at \ + physical_offset={physical_offset} (sequence len={})", + sequence.len() + )) + })?; + row_ids.insert(id); } } } diff --git a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs index cfabdc13266..27e2cad16d4 100644 --- a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs +++ b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs @@ -406,6 +406,55 @@ async fn test_overlay_multi_fragment(#[values(false, true)] stable_row_ids: bool assert_eq!(ids_matching(&dataset, "age = 30").await, vec![3]); } +/// A deletion below an overlaid row must not corrupt the physical-offset → stable-row-id +/// translation used to build the overlay block mask. +/// +/// Under stable row ids the stale-row block/take set is computed by mapping each stale +/// *physical offset* to its stable row id via the fragment's `RowIdSequence`. The sequence +/// keeps one entry per physical row (deleted rows are tracked separately by the deletion +/// vector, not compacted out), so the correct mapping is `sequence.get(offset)`. A regression +/// that instead advanced a `sequence.iter()` cursor only for non-deleted offsets desynced the +/// cursor after any deletion at an offset *below* the stale one, blocking/taking the wrong row +/// id: the stale index hit then leaked and the new value was never surfaced. +/// +/// Setup (stable row ids): fragment 1 holds ids 6..12 at offsets 0..6. Delete id=6 (offset 0), +/// then overlay offset 2 (id=8, age 80 → 999). The deletion at offset 0 sits below the stale +/// offset 2, so a cursor-based translation would map offset 2 to id=7 instead of id=8. +/// +/// Parametrized over `stable_row_ids`: only the stable-row-id path translates offsets to row +/// ids, so the bug is specific to it; the non-stable case (addresses are row ids) is a control. +#[rstest] +#[tokio::test] +async fn test_btree_overlay_stale_row_with_prior_deletion( + #[values(false, true)] stable_row_ids: bool, +) { + let mut dataset = create_base_dataset_with(stable_row_ids).await; + build_age_index(&mut dataset).await; + + // Delete id=6 (fragment 1, offset 0) — a deletion hole below the row the overlay marks stale. + dataset.delete("id = 6").await.unwrap(); + + // Fragment 1, offset 2 is id=8 (age 80). The overlay (committed after the index) → age 999. + let dataset = commit_overlay( + dataset, + "age_overlay_del", + 1, + &[1], + OverlayCoverage::dense(RoaringBitmap::from_iter([2])), + vec![i32_array([Some(999)])], + ) + .await; + + // Stale-drop: id=8's old age=80 index entry must not be returned. + assert_eq!(ids_matching(&dataset, "age = 80").await, Vec::::new()); + // New-match: id=8's current age=999 is found by re-evaluating the stale row. + assert_eq!(ids_matching(&dataset, "age = 999").await, vec![8]); + // A non-stale row in the same deletion-bearing fragment is still served by the index. + assert_eq!(ids_matching(&dataset, "age = 70").await, vec![7]); + // The deleted row is gone. + assert_eq!(ids_matching(&dataset, "age = 60").await, Vec::::new()); +} + const VEC_DIM: i32 = 8; fn vec_query() -> Vec {