feat: introduce delete file manager skeleton. Use in ArrowReader

Author: sdd

Cargo.lock

@@ -86,6 +86,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"
 rust_decimal = "1.31"
 serde = { version = "1.0.204", features = ["rc"] }
 serde_bytes = "0.11.15"

Cargo.toml

@@ -72,6 +72,7 @@ parquet = { workspace = true, features = ["async"] }
 paste = { workspace = true }
 rand = { workspace = true }
 reqwest = { workspace = true }
+roaring = { workspace = true }
 rust_decimal = { workspace = true }
 serde = { workspace = true }
 serde_bytes = { workspace = true }

crates/iceberg/Cargo.toml

@@ -0,0 +1,64 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use roaring::RoaringTreemap;
+
+use crate::expr::BoundPredicate;
+use crate::io::FileIO;
+use crate::scan::FileScanTaskDeleteFile;
+use crate::spec::SchemaRef;
+use crate::{Error, ErrorKind, Result};
+
+pub(crate) struct DeleteFileManager {}
+
+#[allow(unused_variables)]
+impl DeleteFileManager {
+    pub(crate) async fn load_deletes(
+        delete_file_entries: Vec<FileScanTaskDeleteFile>,
+        file_io: FileIO,
+        concurrency_limit_data_files: usize,
+    ) -> Result<DeleteFileManager> {
+        // TODO
+
+        if !delete_file_entries.is_empty() {
+            Err(Error::new(
+                ErrorKind::FeatureUnsupported,
+                "Reading delete files is not yet supported",
+            ))
+        } else {
+            Ok(DeleteFileManager {})
+        }
+    }
+
+    pub(crate) fn build_delete_predicate(
+        &self,
+        snapshot_schema: SchemaRef,
+    ) -> Result<Option<BoundPredicate>> {
+        // TODO
+
+        Ok(None)
+    }
+
+    pub(crate) fn get_positional_delete_indexes_for_data_file(
+        &self,
+        data_file_path: &str,
+    ) -> Option<RoaringTreemap> {
+        // TODO
+
+        None
+    }
+}

crates/iceberg/src/arrow/delete_file_manager.rs

@@ -19,6 +19,7 @@
 
 mod schema;
 pub use schema::*;
+pub(crate) mod delete_file_manager;
 mod reader;
 pub(crate) mod record_batch_projector;
 pub(crate) mod record_batch_transformer;

crates/iceberg/src/arrow/mod.rs

@@ -36,9 +36,11 @@ use futures::{try_join, FutureExt, StreamExt, TryFutureExt, TryStreamExt};
 use parquet::arrow::arrow_reader::{ArrowPredicateFn, ArrowReaderOptions, RowFilter, RowSelection};
 use parquet::arrow::async_reader::AsyncFileReader;
 use parquet::arrow::{ParquetRecordBatchStreamBuilder, ProjectionMask, PARQUET_FIELD_ID_META_KEY};
-use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader};
+use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader, RowGroupMetaData};
 use parquet::schema::types::{SchemaDescriptor, Type as ParquetType};
+use roaring::RoaringTreemap;
 
+use crate::arrow::delete_file_manager::DeleteFileManager;
 use crate::arrow::record_batch_transformer::RecordBatchTransformer;
 use crate::arrow::{arrow_schema_to_schema, get_arrow_datum};
 use crate::error::Result;
@@ -145,6 +147,7 @@ impl ArrowReader {
                     file_io,
                     row_group_filtering_enabled,
                     row_selection_enabled,
+                    concurrency_limit_data_files,
                 )
             })
             .map_err(|err| {
@@ -162,30 +165,24 @@ impl ArrowReader {
         file_io: FileIO,
         row_group_filtering_enabled: bool,
         row_selection_enabled: bool,
+        concurrency_limit_data_files: usize,
     ) -> Result<ArrowRecordBatchStream> {
-        // TODO: add support for delete files
-        if !task.deletes.is_empty() {
-            return Err(Error::new(
-                ErrorKind::FeatureUnsupported,
-                "Delete files are not yet supported",
-            ));
-        }
-
-        // Get the metadata for the Parquet file we need to read and build
-        // a reader for the data within
-        let parquet_file = file_io.new_input(&task.data_file_path)?;
-        let (parquet_metadata, parquet_reader) =
-            try_join!(parquet_file.metadata(), parquet_file.reader())?;
-        let parquet_file_reader = ArrowFileReader::new(parquet_metadata, parquet_reader);
-
-        let should_load_page_index = row_selection_enabled && task.predicate.is_some();
-
-        // Start creating the record batch stream, which wraps the parquet file reader
-        let mut record_batch_stream_builder = ParquetRecordBatchStreamBuilder::new_with_options(
-            parquet_file_reader,
-            ArrowReaderOptions::new().with_page_index(should_load_page_index),
-        )
-        .await?;
+        let should_load_page_index =
+            (row_selection_enabled && task.predicate.is_some()) || !task.deletes.is_empty();
+
+        // concurrently retrieve delete files and create RecordBatchStreamBuilder
+        let (delete_file_manager, mut record_batch_stream_builder) = try_join!(
+            DeleteFileManager::load_deletes(
+                task.deletes.clone(),
+                file_io.clone(),
+                concurrency_limit_data_files
+            ),
+            Self::create_parquet_record_batch_stream_builder(
+                &task.data_file_path,
+                file_io.clone(),
+                should_load_page_index,
+            )
+        )?;
 
         // Create a projection mask for the batch stream to select which columns in the
         // Parquet file that we want in the response
@@ -197,7 +194,7 @@ impl ArrowReader {
         )?;
         record_batch_stream_builder = record_batch_stream_builder.with_projection(projection_mask);
 
-        // RecordBatchTransformer performs any required transformations on the RecordBatches
+        // RecordBatchTransformer performs any transformations required on the RecordBatches
         // that come back from the file, such as type promotion, default column insertion
         // and column re-ordering
         let mut record_batch_transformer =
@@ -207,49 +204,102 @@ impl ArrowReader {
             record_batch_stream_builder = record_batch_stream_builder.with_batch_size(batch_size);
         }
 
-        if let Some(predicate) = &task.predicate {
+        let delete_predicate = delete_file_manager.build_delete_predicate(task.schema.clone())?;
+
+        // In addition to the optional predicate supplied in the `FileScanTask`,
+        // we also have an optional predicate resulting from equality delete files.
+        // If both are present, we logical-AND them together to form a single filter
+        // predicate that we can pass to the `RecordBatchStreamBuilder`.
+        let final_predicate = match (&task.predicate, delete_predicate) {
+            (None, None) => None,
+            (Some(predicate), None) => Some(predicate.clone()),
+            (None, Some(ref predicate)) => Some(predicate.clone()),
+            (Some(filter_predicate), Some(delete_predicate)) => {
+                Some(filter_predicate.clone().and(delete_predicate))
+            }
+        };
+
+        // There are two possible sources both for potential lists of selected RowGroup indices,
+        // and for `RowSelection`s.
+        // Selected RowGroup index lists can come from two sources:
+        //   * When there are equality delete files that are applicable;
+        //   * When there is a scan predicate and row_group_filtering_enabled = true.
+        // `RowSelection`s can be created in either or both of the following cases:
+        //   * When there are positional delete files that are applicable;
+        //   * When there is a scan predicate and row_selection_enabled = true
+        // Note that, in the former case we only perform row group filtering when
+        // there is a scan predicate AND row_group_filtering_enabled = true,
+        // but we perform row selection filtering if there are applicable
+        // equality delete files OR (there is a scan predicate AND row_selection_enabled),
+        // since the only implemented method of applying positional deletes is
+        // by using a `RowSelection`.
+        let mut selected_row_group_indices = None;
+        let mut row_selection = None;
+
+        if let Some(predicate) = final_predicate {
             let (iceberg_field_ids, field_id_map) = Self::build_field_id_set_and_map(
                 record_batch_stream_builder.parquet_schema(),
-                predicate,
+                &predicate,
             )?;
 
             let row_filter = Self::get_row_filter(
-                predicate,
+                &predicate,
                 record_batch_stream_builder.parquet_schema(),
                 &iceberg_field_ids,
                 &field_id_map,
             )?;
             record_batch_stream_builder = record_batch_stream_builder.with_row_filter(row_filter);
 
-            let mut selected_row_groups = None;
             if row_group_filtering_enabled {
                 let result = Self::get_selected_row_group_indices(
-                    predicate,
+                    &predicate,
                     record_batch_stream_builder.metadata(),
                     &field_id_map,
                     &task.schema,
                 )?;
 
-                selected_row_groups = Some(result);
+                selected_row_group_indices = Some(result);
             }
 
             if row_selection_enabled {
-                let row_selection = Self::get_row_selection(
-                    predicate,
+                row_selection = Some(Self::get_row_selection_for_filter_predicate(
+                    &predicate,
                     record_batch_stream_builder.metadata(),
-                    &selected_row_groups,
+                    &selected_row_group_indices,
                     &field_id_map,
                     &task.schema,
-                )?;
-
-                record_batch_stream_builder =
-                    record_batch_stream_builder.with_row_selection(row_selection);
+                )?);
             }
+        }
 
-            if let Some(selected_row_groups) = selected_row_groups {
-                record_batch_stream_builder =
-                    record_batch_stream_builder.with_row_groups(selected_row_groups);
-            }
+        let positional_delete_indexes =
+            delete_file_manager.get_positional_delete_indexes_for_data_file(&task.data_file_path);
+
+        if let Some(positional_delete_indexes) = positional_delete_indexes {
+            let delete_row_selection = Self::build_deletes_row_selection(
+                record_batch_stream_builder.metadata().row_groups(),
+                &selected_row_group_indices,
+                positional_delete_indexes,
+            )?;
+
+            // merge the row selection from the delete files with the row selection
+            // from the filter predicate, if there is one from the filter predicate
+            row_selection = match row_selection {
+                None => Some(delete_row_selection),
+                Some(filter_row_selection) => {
+                    Some(filter_row_selection.intersection(&delete_row_selection))
+                }
+            };
+        }
+
+        if let Some(row_selection) = row_selection {
+            record_batch_stream_builder =
+                record_batch_stream_builder.with_row_selection(row_selection);
+        }
+
+        if let Some(selected_row_group_indices) = selected_row_group_indices {
+            record_batch_stream_builder =
+                record_batch_stream_builder.with_row_groups(selected_row_group_indices);
         }
 
         // Build the batch stream and send all the RecordBatches that it generates
@@ -265,6 +315,43 @@ impl ArrowReader {
         Ok(Box::pin(record_batch_stream) as ArrowRecordBatchStream)
     }
 
+    async fn create_parquet_record_batch_stream_builder(
+        data_file_path: &str,
+        file_io: FileIO,
+        should_load_page_index: bool,
+    ) -> Result<ParquetRecordBatchStreamBuilder<ArrowFileReader<impl FileRead + Sized>>> {
+        // Get the metadata for the Parquet file we need to read and build
+        // a reader for the data within
+        let parquet_file = file_io.new_input(data_file_path)?;
+        let (parquet_metadata, parquet_reader) =
+            try_join!(parquet_file.metadata(), parquet_file.reader())?;
+        let parquet_file_reader = ArrowFileReader::new(parquet_metadata, parquet_reader);
+
+        // Create the record batch stream builder, which wraps the parquet file reader
+        let record_batch_stream_builder = ParquetRecordBatchStreamBuilder::new_with_options(
+            parquet_file_reader,
+            ArrowReaderOptions::new().with_page_index(should_load_page_index),
+        )
+        .await?;
+        Ok(record_batch_stream_builder)
+    }
+
+    /// computes a `RowSelection` from positional delete indices.
+    ///
+    /// Using the Parquet page index, we build a `RowSelection` that rejects rows that are indicated
+    /// as having been deleted by a positional delete, taking into account any row groups that have
+    /// been skipped entirely by the filter predicate
+    #[allow(unused)]
+    fn build_deletes_row_selection(
+        row_group_metadata: &[RowGroupMetaData],
+        selected_row_groups: &Option<Vec<usize>>,
+        mut positional_deletes: RoaringTreemap,
+    ) -> Result<RowSelection> {
+        // TODO
+
+        Ok(RowSelection::default())
+    }
+
     fn build_field_id_set_and_map(
         parquet_schema: &SchemaDescriptor,
         predicate: &BoundPredicate,
@@ -475,7 +562,7 @@ impl ArrowReader {
         Ok(results)
     }
 
-    fn get_row_selection(
+    fn get_row_selection_for_filter_predicate(
         predicate: &BoundPredicate,
         parquet_metadata: &Arc<ParquetMetaData>,
         selected_row_groups: &Option<Vec<usize>>,

crates/iceberg/src/arrow/reader.rs

@@ -726,6 +726,12 @@ pub enum BoundPredicate {
     Set(SetExpression<BoundReference>),
 }
 
+impl BoundPredicate {
+    pub(crate) fn and(self, other: BoundPredicate) -> BoundPredicate {
+        BoundPredicate::And(LogicalExpression::new([Box::new(self), Box::new(other)]))
+    }
+}
+
 impl Display for BoundPredicate {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {