From 9f463f6389e23c93f0b39b1c4e9d130227a1d1a9 Mon Sep 17 00:00:00 2001 From: jerrytao Date: Thu, 17 Sep 2026 17:10:09 +0800 Subject: [PATCH] perf(scan): resolve deletions from row offsets instead of row addresses A scan over a fragment with deletions probed the vector once per row and materialized a u64 address for every row just to recover the offset. Cost followed rows scanned, not rows deleted: 10 deletes in a 1M-row fragment made a full scan ~3x slower. Build the keep mask from the batch's offset ranges. Stop fetching row addresses only because the fragment has deletions. --- .../python/benchmarks/test_deletion_scan.py | 59 ++++++ rust/lance-core/src/utils/deletion.rs | 184 ++++++++++++++++++ rust/lance-table/src/utils/stream.rs | 91 +++++++-- 3 files changed, 320 insertions(+), 14 deletions(-) create mode 100644 python/python/benchmarks/test_deletion_scan.py diff --git a/python/python/benchmarks/test_deletion_scan.py b/python/python/benchmarks/test_deletion_scan.py new file mode 100644 index 00000000000..7905d851c6b --- /dev/null +++ b/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 diff --git a/rust/lance-core/src/utils/deletion.rs b/rust/lance-core/src/utils/deletion.rs index c7f8b142464..a975a44bd71 100644 --- a/rust/lance-core/src/utils/deletion.rs +++ b/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], num_rows: u32) -> Option { + debug_assert_eq!( + ranges.iter().map(|r| r.end - r.start).sum::(), + 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, +} + +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 { + 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], num_rows: u32) -> Vec { + 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]) -> Vec { + 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>, + #[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::>()) + .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>) { + 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::>()) + .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)] diff --git a/rust/lance-table/src/utils/stream.rs b/rust/lance-table/src/utils/stream.rs index a45d0f3c139..da749cee3ea 100644 --- a/rust/lance-table/src/utils/stream.rs +++ b/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::>()), + 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::()) + .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::>().await.unwrap(); + + let actual = batches + .iter() + .flat_map(|batch| batch[ROW_ID].as_primitive::().values().to_vec()) + .collect::>(); + 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::>(); + assert_eq!(actual, expected); + } + #[tokio::test] async fn test_deletes() { let no_deletes: Option> = None;