perf(scan): resolve deletions from row offsets instead of row addresses

python/python/benchmarks/test_deletion_scan.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright The Lance Authors
+"""Scans of a fragment that carries a deletion vector.
+
+Every row a scan reads has to be resolved against the deletion vector, so deleting
+even a handful of rows changes what a full scan costs.  These compare a clean
+fragment against one with ten rows removed, laid out both as one run and scattered.
+
+A single process can load only one native pylance.  To compare official main
+against this change, install each wheel into its own interpreter and run this
+file twice:
+
+    base/bin/python -m pytest python/python/benchmarks/test_deletion_scan.py
+    patched/bin/python -m pytest python/python/benchmarks/test_deletion_scan.py
+
+Read the Min column.  ``no_deletions`` should stay about the same; the two
+deletion cases should drop from ~3x that baseline down to about 1x.
+"""
+
+from pathlib import Path
+
+import lance
+import pyarrow as pa
+import pyarrow.compute as pc
+import pytest
+
+NUM_ROWS = 1_000_000
+NUM_DELETED = 10
+
+
+@pytest.mark.parametrize(
+    "predicate",
+    [
+        None,
+        f"id >= {NUM_ROWS // 2} AND id < {NUM_ROWS // 2 + NUM_DELETED}",
+        f"id % {NUM_ROWS // NUM_DELETED} == 0",
+    ],
+    ids=["no_deletions", "contiguous_deletions", "scattered_deletions"],
+)
+@pytest.mark.benchmark(group="scan_with_deletions")
+def test_scan_with_deletions(tmp_path: Path, benchmark, predicate):
+    table = pa.table(
+        {
+            "id": pa.array(range(NUM_ROWS), type=pa.int64()),
+            "value": pc.random(NUM_ROWS),
+        }
+    )
+    # One fragment, so the whole scan runs under a single deletion vector.
+    dataset = lance.write_dataset(table, tmp_path, max_rows_per_file=NUM_ROWS)
+    if predicate is not None:
+        dataset.delete(predicate)
+        dataset = lance.dataset(tmp_path)
+        assert dataset.get_fragments()[0].metadata.deletion_file is not None
+
+    benchmark.name = benchmark.param
+    result = benchmark(dataset.to_table)
+
+    expected = NUM_ROWS if predicate is None else NUM_ROWS - NUM_DELETED
+    assert result.num_rows == expected

rust/lance-core/src/utils/deletion.rs

@@ -5,6 +5,7 @@ use std::{collections::HashSet, ops::Range, sync::Arc};
 
 use crate::deepsize::{Context, DeepSizeOf};
 use arrow_array::BooleanArray;
+use arrow_buffer::BooleanBufferBuilder;
 use roaring::RoaringBitmap;
 
 /// Threshold for when a DeletionVector::Set should be promoted to a DeletionVector::Bitmap.
@@ -102,6 +103,66 @@ impl DeletionVector {
         }
     }
 
+    /// Build a mask of the rows in a batch that survive deletion.
+    ///
+    /// `ranges` are the fragment row offsets the batch covers, in the order the batch
+    /// returns them, and must sum to `num_rows`.  Bit `i` of the result is set when the
+    /// batch's `i`th row is still live.
+    ///
+    /// Returns `None` when the batch contains no deleted row, letting the caller pass it
+    /// along untouched.
+    ///
+    /// Unlike [`Self::build_predicate`] this looks deleted rows up by offset instead of
+    /// asking about every row in turn, so the cost follows the deletions a batch actually
+    /// covers rather than the number of rows it holds.  A [`HashSet`] cannot be queried by
+    /// range, so for that representation the two loops are compared by length and the
+    /// shorter one is walked.
+    pub fn build_keep_mask(&self, ranges: &[Range<u32>], num_rows: u32) -> Option<BooleanArray> {
+        debug_assert_eq!(
+            ranges.iter().map(|r| r.end - r.start).sum::<u32>(),
+            num_rows,
+            "ranges must cover exactly num_rows rows"
+        );
+
+        let mut mask = KeepMaskBuilder::new(num_rows);
+        match self {
+            Self::NoDeletions => {}
+            Self::Bitmap(bitmap) => {
+                let mut base = 0;
+                for range in ranges {
+                    for deleted in bitmap.range(range.clone()) {
+                        mask.delete(base + deleted - range.start);
+                    }
+                    base += range.end - range.start;
+                }
+            }
+            // Walking the set costs one step per entry for every range, while probing
+            // costs one lookup per row.  A set never grows past `BITMAP_THRESDHOLD`, so
+            // on a batch-sized read the first is usually the shorter walk.
+            Self::Set(set) if set.len().saturating_mul(ranges.len()) <= num_rows as usize => {
+                let mut base = 0;
+                for range in ranges {
+                    for deleted in set.iter().copied().filter(|offset| range.contains(offset)) {
+                        mask.delete(base + deleted - range.start);
+                    }
+                    base += range.end - range.start;
+                }
+            }
+            Self::Set(set) => {
+                let mut position = 0;
+                for range in ranges {
+                    for offset in range.clone() {
+                        if set.contains(&offset) {
+                            mask.delete(position);
+                        }
+                        position += 1;
+                    }
+                }
+            }
+        }
+        mask.finish()
+    }
+
     // Note: deletion vectors are based on 32-bit offsets.  However, this function works
     // even when given 64-bit row addresses.  That is because `id as u32` returns the lower
     // 32 bits (the row offset) and the upper 32 bits are ignored.
@@ -123,6 +184,39 @@ impl DeletionVector {
     }
 }
 
+/// Collects deleted positions into a keep mask, allocating only once one shows up.
+///
+/// Most batches of a scan contain no deleted row at all, and those cost nothing here.
+struct KeepMaskBuilder {
+    num_rows: u32,
+    builder: Option<BooleanBufferBuilder>,
+}
+
+impl KeepMaskBuilder {
+    fn new(num_rows: u32) -> Self {
+        Self {
+            num_rows,
+            builder: None,
+        }
+    }
+
+    fn delete(&mut self, position: u32) {
+        let num_rows = self.num_rows as usize;
+        self.builder
+            .get_or_insert_with(|| {
+                let mut builder = BooleanBufferBuilder::new(num_rows);
+                builder.append_n(num_rows, true);
+                builder
+            })
+            .set_bit(position as usize, false);
+    }
+
+    fn finish(self) -> Option<BooleanArray> {
+        self.builder
+            .map(|mut builder| BooleanArray::new(builder.finish(), None))
+    }
+}
+
 /// Maps a naive offset into a fragment to the local row offset that is
 /// not deleted.
 ///
@@ -421,6 +515,96 @@ mod test {
         );
     }
 
+    /// A missing mask means every row survived, so spell that out as a full mask.
+    fn keep_mask_values(dv: &DeletionVector, ranges: &[Range<u32>], num_rows: u32) -> Vec<bool> {
+        match dv.build_keep_mask(ranges, num_rows) {
+            Some(mask) => mask.iter().map(|v| v.unwrap()).collect(),
+            None => vec![true; num_rows as usize],
+        }
+    }
+
+    #[rstest]
+    #[case::set(set_dv([1, 3]))]
+    #[case::bitmap(bitmap_dv([1, 3]))]
+    fn test_build_keep_mask(#[case] dv: DeletionVector) {
+        assert_eq!(
+            keep_mask_values(&dv, &[0..5], 5),
+            [true, false, true, false, true]
+        );
+    }
+
+    #[rstest]
+    #[case::set(set_dv([1, 6, 9]))]
+    #[case::bitmap(bitmap_dv([1, 6, 9]))]
+    fn test_build_keep_mask_maps_offsets_to_batch_positions(#[case] dv: DeletionVector) {
+        // Offsets 0..3 land at positions 0..3 and offsets 6..10 at positions 3..7, so the
+        // deleted offsets 1, 6 and 9 must show up at positions 1, 3 and 6.
+        assert_eq!(
+            keep_mask_values(&dv, &[0..3, 6..10], 7),
+            [true, false, true, false, true, true, false]
+        );
+    }
+
+    #[rstest]
+    #[case::no_deletions(DeletionVector::NoDeletions)]
+    #[case::set(set_dv([100, 200]))]
+    #[case::bitmap(bitmap_dv([100, 200]))]
+    fn test_build_keep_mask_skips_batches_without_deletions(#[case] dv: DeletionVector) {
+        assert!(dv.build_keep_mask(&[0..10], 10).is_none());
+    }
+
+    /// The row addresses `build_predicate` would be handed for a batch covering `ranges`.
+    fn batch_row_addrs(ranges: &[Range<u32>]) -> Vec<u64> {
+        ranges
+            .iter()
+            .flat_map(|range| range.clone())
+            .map(u64::from)
+            .collect()
+    }
+
+    #[rstest]
+    // Few enough deleted rows that walking the set is the shorter loop.
+    #[case::set_walked(set_dv([2, 5, 11]), vec![0..8, 10..14], 12)]
+    // More deleted rows than the batch holds, so the per-row probe is used instead.
+    #[case::set_probed(set_dv([0, 1, 2, 3, 4, 5, 6, 7]), vec![2..5], 3)]
+    #[case::bitmap(bitmap_dv([2, 5, 11]), vec![0..8, 10..14], 12)]
+    fn test_build_keep_mask_matches_build_predicate(
+        #[case] dv: DeletionVector,
+        #[case] ranges: Vec<Range<u32>>,
+        #[case] num_rows: u32,
+    ) {
+        let row_addrs = batch_row_addrs(&ranges);
+        let expected = dv
+            .build_predicate(row_addrs.iter())
+            .map(|mask| mask.iter().map(|v| v.unwrap()).collect::<Vec<_>>())
+            .unwrap_or_else(|| vec![true; num_rows as usize]);
+        assert_eq!(keep_mask_values(&dv, &ranges, num_rows), expected);
+    }
+
+    /// Bitmap lookups always go through `range()`, including short runs, runs that
+    /// start mid-fragment, and a batch that mixes several of those in one call.
+    #[rstest]
+    #[case::short_run(vec![0..3])]
+    #[case::long_run(vec![0..64])]
+    #[case::offset_run(vec![32..80])]
+    #[case::mixed_runs(vec![0..3, 10..26, 100..104])]
+    fn test_build_keep_mask_range_lookup(#[case] ranges: Vec<Range<u32>>) {
+        let num_rows = ranges.iter().map(|range| range.end - range.start).sum();
+        // Spread across the fragment and onto run edges, so every case drops a row.
+        let dv = bitmap_dv([0, 3, 11, 25, 40, 41, 42, 79, 100, 103]);
+
+        let row_addrs = batch_row_addrs(&ranges);
+        let expected = dv
+            .build_predicate(row_addrs.iter())
+            .map(|mask| mask.iter().map(|v| v.unwrap()).collect::<Vec<_>>())
+            .unwrap_or_else(|| vec![true; num_rows as usize]);
+        assert!(
+            expected.contains(&false),
+            "case would pass without resolving any deletion"
+        );
+        assert_eq!(keep_mask_values(&dv, &ranges, num_rows), expected);
+    }
+
     #[rstest]
     #[case::no_deletions(DeletionVector::NoDeletions, 0)]
     #[case::set(set_dv([1, 2, 3]), 3)]

rust/lance-table/src/utils/stream.rs

@@ -487,22 +487,37 @@ fn apply_row_id_and_deletes_with_system_columns(
     debug_assert!(batch.num_columns() > 0 || config.has_system_cols() || has_deletions);
 
     // If row id sequence is None, then row id IS row address.
-    let should_fetch_row_addr = config.with_row_addr
-        || (config.with_row_id && config.row_id_sequence.is_none())
-        || has_deletions;
+    //
+    // Deletions are resolved from row offsets, so they no longer make row addresses
+    // worth materializing on their own.
+    let should_fetch_row_addr =
+        config.with_row_addr || (config.with_row_id && config.row_id_sequence.is_none());
 
     let num_rows = batch.num_rows() as u32;
 
+    // Row addresses and the deletion mask are both derived from where this batch's rows
+    // sit in the fragment, so resolve that once and hand it to both.
+    let batch_selection = if should_fetch_row_addr || has_deletions {
+        Some(
+            config
+                .params
+                .slice(batch_offset as usize, num_rows as usize)?,
+        )
+    } else {
+        None
+    };
+
+    // Only the mask needs the runs up front; row addresses can stream them.
+    let offset_ranges = match batch_selection.as_ref().filter(|_| has_deletions) {
+        Some(selection) => Some(selection.iter_offset_ranges()?.collect::<Vec<_>>()),
+        None => None,
+    };
+
     let row_addrs =
-        if should_fetch_row_addr {
+        if let Some(selection) = batch_selection.as_ref().filter(|_| should_fetch_row_addr) {
             let _rowaddrs = tracing::span!(tracing::Level::DEBUG, "fetch_row_addrs").entered();
             let mut row_addrs = Vec::with_capacity(num_rows as usize);
-            for offset_range in config
-                .params
-                .slice(batch_offset as usize, num_rows as usize)
-                .unwrap()
-                .iter_offset_ranges()?
-            {
+            for offset_range in selection.iter_offset_ranges()? {
                 row_addrs.extend(offset_range.map(|row_offset| {
                     u64::from(RowAddress::new_from_parts(fragment_id, row_offset))
                 }));
@@ -545,10 +560,12 @@ fn apply_row_id_and_deletes_with_system_columns(
 
     let span = tracing::span!(tracing::Level::DEBUG, "apply_deletions");
     let _enter = span.enter();
-    let deletion_mask = deletion_vector.and_then(|v| {
-        let row_addrs: &[u64] = row_addrs.as_ref().unwrap().values();
-        v.build_predicate(row_addrs.iter())
-    });
+    let deletion_mask =
+        deletion_vector
+            .zip(offset_ranges.as_ref())
+            .and_then(|(deletion_vector, offset_ranges)| {
+                deletion_vector.build_keep_mask(offset_ranges, num_rows)
+            });
 
     let mut system_columns: Vec<(Field, ArrayRef)> = Vec::with_capacity(4);
     if config.with_row_id {
@@ -1490,6 +1507,52 @@ mod tests {
         .await;
     }
 
+    /// A filtered scan reads a fragment as a set of ranges, which is the shape where a
+    /// batch's position and its fragment row offset drift apart.
+    #[tokio::test]
+    async fn test_deletes_with_range_selection() {
+        // 100 rows over 10 batches, selected as two ranges of a 110 row fragment.
+        let data = batch_task_stream(
+            lance_datagen::gen_batch()
+                .col("x", lance_datagen::array::rand::<Int32Type>())
+                .into_reader_stream(RowCount::from(10), BatchCount::from(10))
+                .0,
+        );
+
+        // Offset 55 is deleted but never selected, so it must not shift anything.
+        let deletion_vector = Arc::new(DeletionVector::Bitmap(RoaringBitmap::from_iter([
+            0, 49, 55, 60, 109,
+        ])));
+
+        let config = RowIdAndDeletesConfig {
+            params: ReadBatchParams::Ranges(Arc::from(vec![0..50_u64, 60..110])),
+            with_row_id: true,
+            with_row_addr: false,
+            with_row_last_updated_at_version: false,
+            with_row_created_at_version: false,
+            deletion_vector: Some(deletion_vector),
+            row_id_sequence: None,
+            last_updated_at_sequence: None,
+            created_at_sequence: None,
+            make_deletions_null: false,
+            total_num_rows: 110,
+        };
+
+        let stream = super::wrap_with_row_id_and_delete(data, 0, config);
+        let batches = stream.buffered(1).try_collect::<Vec<_>>().await.unwrap();
+
+        let actual = batches
+            .iter()
+            .flat_map(|batch| batch[ROW_ID].as_primitive::<UInt64Type>().values().to_vec())
+            .collect::<Vec<u64>>();
+        let expected = (0..50)
+            .chain(60..110)
+            .filter(|offset| ![0, 49, 60, 109].contains(offset))
+            .map(|offset| RowAddress::new_from_parts(0, offset).into())
+            .collect::<Vec<u64>>();
+        assert_eq!(actual, expected);
+    }
+
     #[tokio::test]
     async fn test_deletes() {
         let no_deletes: Option<Arc<DeletionVector>> = None;