refactor: delete row selection uses roaring treemap iterator and advance_to

Author: sdd

Cargo.lock

@@ -83,7 +83,7 @@ port_scanner = "0.1.5"
 rand = "0.8.5"
 regex = "1.10.5"
 reqwest = { version = "0.12.2", default-features = false, features = ["json"] }
-roaring = "0.10"
+roaring = { version = "0.10", git = "https://github.com/RoaringBitmap/roaring-rs.git" }
 rust_decimal = "1.31"
 serde = { version = "1.0.204", features = ["rc"] }
 serde_bytes = "0.11.15"

Cargo.toml

@@ -370,17 +370,19 @@ impl ArrowReader {
     /// as having been deleted by a positional delete, taking into account any row groups that have
     /// been skipped entirely by the filter predicate
     fn build_deletes_row_selection(
-        row_group_metadata: &[RowGroupMetaData],
+        row_group_metadata_list: &[RowGroupMetaData],
         selected_row_groups: &Option<Vec<usize>>,
-        mut positional_deletes: &DeleteVector,
+        positional_deletes: &DeleteVector,
     ) -> Result<RowSelection> {
         let mut results: Vec<RowSelector> = Vec::new();
         let mut selected_row_groups_idx = 0;
-        let mut current_page_base_idx: u64 = 0;
+        let mut current_row_group_base_idx: u64 = 0;
+        let mut delete_vector_iter = positional_deletes.iter();
+        let mut next_deleted_row_idx_opt = delete_vector_iter.next();
 
-        for (idx, row_group_metadata) in row_group_metadata.iter().enumerate() {
-            let page_num_rows = row_group_metadata.num_rows() as u64;
-            let next_page_base_idx = current_page_base_idx + page_num_rows;
+        for (idx, row_group_metadata) in row_group_metadata_list.iter().enumerate() {
+            let row_group_num_rows = row_group_metadata.num_rows() as u64;
+            let next_row_group_base_idx = current_row_group_base_idx + row_group_num_rows;
 
             // if row group selection is enabled,
             if let Some(selected_row_groups) = selected_row_groups {
@@ -397,36 +399,37 @@ impl ArrowReader {
                 } else {
                     // remove any positional deletes from the skipped page so that
                     // `positional.deletes.min()` can be used
-                    positional_deletes.remove_range(current_page_base_idx..next_page_base_idx);
+                    delete_vector_iter.advance_to(next_row_group_base_idx);
+                    next_deleted_row_idx_opt = delete_vector_iter.next();
 
                     // still increment the current page base index but then skip to the next row group
                     // in the file
-                    current_page_base_idx += page_num_rows;
+                    current_row_group_base_idx += row_group_num_rows;
                     continue;
                 }
             }
 
-            let mut next_deleted_row_idx = match positional_deletes.min() {
+            let mut next_deleted_row_idx = match next_deleted_row_idx_opt {
                 Some(next_deleted_row_idx) => {
-                    // if the index of the next deleted row is beyond this page, add a selection for
-                    // the remainder of this page and skip to the next page
-                    if next_deleted_row_idx >= next_page_base_idx {
-                        results.push(RowSelector::select(page_num_rows as usize));
+                    // if the index of the next deleted row is beyond this row group, add a selection for
+                    // the remainder of this row group and skip to the next row group
+                    if next_deleted_row_idx >= next_row_group_base_idx {
+                        results.push(RowSelector::select(row_group_num_rows as usize));
                         continue;
                     }
 
                     next_deleted_row_idx
                 }
 
-                // If there are no more pos deletes, add a selector for the entirety of this page.
+                // If there are no more pos deletes, add a selector for the entirety of this row group.
                 _ => {
-                    results.push(RowSelector::select(page_num_rows as usize));
+                    results.push(RowSelector::select(row_group_num_rows as usize));
                     continue;
                 }
             };
 
-            let mut current_idx = current_page_base_idx;
-            'chunks: while next_deleted_row_idx < next_page_base_idx {
+            let mut current_idx = current_row_group_base_idx;
+            'chunks: while next_deleted_row_idx < next_row_group_base_idx {
                 // `select` all rows that precede the next delete index
                 if current_idx < next_deleted_row_idx {
                     let run_length = next_deleted_row_idx - current_idx;
@@ -437,18 +440,18 @@ impl ArrowReader {
                 // `skip` all consecutive deleted rows in the current row group
                 let mut run_length = 0;
                 while next_deleted_row_idx == current_idx
-                    && next_deleted_row_idx < next_page_base_idx
+                    && next_deleted_row_idx < next_row_group_base_idx
                 {
                     run_length += 1;
                     current_idx += 1;
-                    positional_deletes.remove(next_deleted_row_idx);
 
-                    next_deleted_row_idx = match positional_deletes.min() {
+                    next_deleted_row_idx_opt = delete_vector_iter.next();
+                    next_deleted_row_idx = match next_deleted_row_idx_opt {
                         Some(next_deleted_row_idx) => next_deleted_row_idx,
                         _ => {
                             // We've processed the final positional delete.
                             // Conclude the skip and then break so that we select the remaining
-                            // rows in the page and move on to the next row group
+                            // rows in the row group and move on to the next row group
                             results.push(RowSelector::skip(run_length));
                             break 'chunks;
                         }
@@ -457,13 +460,13 @@ impl ArrowReader {
                 results.push(RowSelector::skip(run_length));
             }
 
-            if current_idx < next_page_base_idx {
+            if current_idx < next_row_group_base_idx {
                 results.push(RowSelector::select(
-                    (next_page_base_idx - current_idx) as usize,
+                    (next_row_group_base_idx - current_idx) as usize,
                 ));
             }
 
-            current_page_base_idx += page_num_rows;
+            current_row_group_base_idx += row_group_num_rows;
         }
 
         Ok(results.into())
@@ -1400,18 +1403,19 @@ mod tests {
     use arrow_array::{ArrayRef, RecordBatch, StringArray};
     use arrow_schema::{DataType, Field, Schema as ArrowSchema, TimeUnit};
     use futures::TryStreamExt;
+    use parquet::arrow::arrow_reader::{RowSelection, RowSelector};
     use parquet::arrow::{ArrowWriter, ProjectionMask};
     use parquet::basic::Compression;
-    use parquet::file::properties::WriterProperties;
-    use parquet::arrow::arrow_reader::{RowSelection, RowSelector};
     use parquet::file::metadata::{ColumnChunkMetaData, RowGroupMetaData};
+    use parquet::file::properties::WriterProperties;
     use parquet::schema::parser::parse_message_type;
-    use tempfile::TempDir;
     use parquet::schema::types::{SchemaDescPtr, SchemaDescriptor};
     use roaring::RoaringTreemap;
+    use tempfile::TempDir;
 
     use crate::arrow::reader::{CollectFieldIdVisitor, PARQUET_FIELD_ID_META_KEY};
     use crate::arrow::{ArrowReader, ArrowReaderBuilder};
+    use crate::delete_vector::DeleteVector;
     use crate::expr::visitors::bound_predicate_visitor::visit;
     use crate::expr::{Bind, Predicate, Reference};
     use crate::io::FileIO;
@@ -1758,16 +1762,14 @@ message schema {
             2999, // single item at end of selected rg3 (1)
             3000, // single item at start of skipped rg4
         ]);
-        
-        let positional_deletes = DeleteVector {
-            inner: positional_deletes
-        };
+
+        let positional_deletes = DeleteVector::new(positional_deletes);
 
         // using selected row groups 1 and 3
         let result = ArrowReader::build_deletes_row_selection(
             &row_groups_metadata,
             &selected_row_groups,
-            positional_deletes.clone(),
+            &positional_deletes,
         )
         .unwrap();
 
@@ -1791,7 +1793,7 @@ message schema {
         let result = ArrowReader::build_deletes_row_selection(
             &row_groups_metadata,
             &None,
-            positional_deletes,
+            &positional_deletes,
         )
         .unwrap();
 

crates/iceberg/src/arrow/reader.rs

@@ -17,6 +17,8 @@
 
 use std::ops::BitOrAssign;
 
+use roaring::bitmap::Iter;
+use roaring::treemap::BitmapIter;
 use roaring::RoaringTreemap;
 
 #[derive(Debug, Default)]
@@ -25,22 +27,18 @@ pub struct DeleteVector {
 }
 
 impl DeleteVector {
-    pub fn iter(&self) -> DeleteVectorIterator {
-        let mut iter = self.inner.bitmaps();
-        match iter.next() {
-            Some((high_bits, bitmap)) => {
-                DeleteVectorIterator {
-                    inner: Some(DeleteVectorIteratorInner {
-                        // iter,
-                        high_bits: (high_bits as u64) << 32,
-                        bitmap_iter: bitmap.iter(),
-                    }),
-                }
-            }
-            _ => DeleteVectorIterator { inner: None },
+    #[allow(unused)]
+    pub(crate) fn new(roaring_treemap: RoaringTreemap) -> DeleteVector {
+        DeleteVector {
+            inner: roaring_treemap,
         }
     }
 
+    pub fn iter(&self) -> DeleteVectorIterator {
+        let outer = self.inner.bitmaps();
+        DeleteVectorIterator { outer, inner: None }
+    }
+
     pub fn intersect_assign(&mut self, other: &DeleteVector) {
         self.inner.bitor_assign(&other.inner);
     }
@@ -54,51 +52,68 @@ impl DeleteVector {
     }
 }
 
+// Ideally, we'd just wrap `roaring::RoaringTreemap`'s iterator, `roaring::treemap::Iter` here.
+// But right now, it does not have a corresponding implementation of `roaring::bitmap::Iter::advance_to`,
+// which is very handy in ArrowReader::build_deletes_row_selection.
+// There is a PR open on roaring to add this (https://github.com/RoaringBitmap/roaring-rs/pull/314)
+// and if that gets merged then we can simplify `DeleteVectorIterator` here, refactoring `advance_to`
+// to just a wrapper around the underlying iterator's method.
 pub struct DeleteVectorIterator<'a> {
+    // NB: `BitMapIter` was only exposed publicly in https://github.com/RoaringBitmap/roaring-rs/pull/316
+    // which is not yet released. As a consequence our Cargo.toml temporarily uses a git reference for
+    // the roaring dependency.
+    outer: BitmapIter<'a>,
     inner: Option<DeleteVectorIteratorInner<'a>>,
 }
 
 struct DeleteVectorIteratorInner<'a> {
-    // TODO: roaring::treemap::iter::BitmapIter is currently private.
-    // See https://github.com/RoaringBitmap/roaring-rs/issues/312
-    // iter: roaring::treemap::iter::BitmapIter<'a>,
-    high_bits: u64,
-    bitmap_iter: roaring::bitmap::Iter<'a>,
+    high_bits: u32,
+    bitmap_iter: Iter<'a>,
 }
 
 impl Iterator for DeleteVectorIterator<'_> {
     type Item = u64;
 
     fn next(&mut self) -> Option<Self::Item> {
-        let Some(ref mut inner) = &mut self.inner else {
-            return None;
-        };
+        if let Some(ref mut inner) = &mut self.inner {
+            if let Some(inner_next) = inner.bitmap_iter.next() {
+                return Some(u64::from(inner.high_bits) << 32 | u64::from(inner_next));
+            }
+        }
 
-        if let Some(lower) = inner.bitmap_iter.next() {
-            return Some(inner.high_bits & lower as u64);
-        };
+        if let Some((high_bits, next_bitmap)) = self.outer.next() {
+            self.inner = Some(DeleteVectorIteratorInner {
+                high_bits,
+                bitmap_iter: next_bitmap.iter(),
+            })
+        } else {
+            return None;
+        }
 
-        // TODO: roaring::treemap::iter::BitmapIter is currently private.
-        // See https://github.com/RoaringBitmap/roaring-rs/issues/312
-
-        // replace with commented-out code below once BitmapIter is pub,
-        // or use RoaringTreemap::iter if `advance_to` gets implemented natively
-        None
-
-        // let Some((high_bits, bitmap)) = inner.iter.next() else {
-        //     self.inner = None;
-        //     return None;
-        // };
-        //
-        // inner.high_bits = (high_bits as u64) << 32;
-        // inner.bitmap_iter = bitmap.iter();
-        //
-        // self.next()
+        self.next()
     }
 }
 
 impl<'a> DeleteVectorIterator<'a> {
-    pub fn advance_to(&'a mut self, _pos: u64) {
-        // TODO
+    pub fn advance_to(&mut self, pos: u64) {
+        let hi = (pos >> 32) as u32;
+        let lo = pos as u32;
+
+        let Some(ref mut inner) = self.inner else {
+            return;
+        };
+
+        while inner.high_bits < hi {
+            let Some((next_hi, next_bitmap)) = self.outer.next() else {
+                return;
+            };
+
+            *inner = DeleteVectorIteratorInner {
+                high_bits: next_hi,
+                bitmap_iter: next_bitmap.iter(),
+            }
+        }
+
+        inner.bitmap_iter.advance_to(lo);
     }
 }