This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-3269-d3e9f456e29b10ba95f11b3a2bc5c9dc546e1d93 in repository https://gitbox.apache.org/repos/asf/iceberg-rust.git
commit acd610997de78f986ece9f89c3bf476de5e1f5a5 Author: Minh Vu <[email protected]> AuthorDate: Sun Oct 11 18:29:22 2026 +0000 perf(reader): speed up row group filtering and name mapping (#3269) * perf(reader): optimize row group filter intersection * perf(reader): add row group intersection benchmark * style: format iceberg benchmark manifest * perf(reader): avoid allocations in name mapping lookup * perf(reader): prune row groups from split candidates * fix(reader): format tests and preallocate name mapping index * fix(reader): tighten candidate pruning and name mapping coverage --- crates/iceberg/src/arrow/reader/pipeline.rs | 53 ++++++++------- crates/iceberg/src/arrow/reader/projection.rs | 49 ++++++++++++-- crates/iceberg/src/arrow/reader/row_filter.rs | 96 ++++++++++++++++++++++++++- 3 files changed, 161 insertions(+), 37 deletions(-) diff --git a/crates/iceberg/src/arrow/reader/pipeline.rs b/crates/iceberg/src/arrow/reader/pipeline.rs index 18ba42d3b..84b4467d3 100644 --- a/crates/iceberg/src/arrow/reader/pipeline.rs +++ b/crates/iceberg/src/arrow/reader/pipeline.rs @@ -20,6 +20,7 @@ //! predicates, row-group / row selection, and delete handling into a stream //! of transformed Arrow `RecordBatch`es. +use std::borrow::Cow; use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::AtomicU64; @@ -133,6 +134,17 @@ struct FileScanTaskReader { scan_metrics: ScanMetrics, } +// Borrow an existing selection, or enumerate all groups for a whole-file scan. +fn candidate_row_groups( + selected_row_groups: Option<&[usize]>, + num_row_groups: usize, +) -> Cow<'_, [usize]> { + match selected_row_groups { + Some(indices) => Cow::Borrowed(indices), + None => Cow::Owned((0..num_row_groups).collect()), + } +} + impl FileScanTaskReader { async fn process(self, task: FileScanTask) -> Result<ArrowRecordBatchStream> { let should_load_page_index = (self.row_selection_enabled && task.predicate().is_some()) @@ -525,48 +537,35 @@ impl FileScanTaskReader { record_batch_stream_builder = record_batch_stream_builder.with_row_filter(row_filter); if self.row_group_filtering_enabled { - let predicate_filtered_row_groups = ArrowReader::get_selected_row_group_indices( + let candidates = candidate_row_groups( + selected_row_group_indices.as_deref(), + record_batch_stream_builder.metadata().num_row_groups(), + ); + + selected_row_group_indices = Some(ArrowReader::get_selected_row_group_indices( &predicate, record_batch_stream_builder.metadata(), &field_id_map, task.schema(), - )?; - - // Merge predicate-based filtering with byte range filtering (if present) - // by taking the intersection of both filters - selected_row_group_indices = match selected_row_group_indices { - Some(byte_range_filtered) => { - // Keep only row groups that are in both filters - let intersection: Vec<usize> = byte_range_filtered - .into_iter() - .filter(|idx| predicate_filtered_row_groups.contains(idx)) - .collect(); - Some(intersection) - } - None => Some(predicate_filtered_row_groups), - }; + candidates.as_ref(), + )?); } if self.bloom_filter_enabled { - let all_rgs; - let candidate_rgs = match &selected_row_group_indices { - Some(indices) => indices.as_slice(), - None => { - all_rgs = (0..record_batch_stream_builder.metadata().num_row_groups()) - .collect::<Vec<_>>(); - &all_rgs - } - }; + let candidates = candidate_row_groups( + selected_row_group_indices.as_deref(), + record_batch_stream_builder.metadata().num_row_groups(), + ); let bloom_filtered = Self::filter_row_groups_by_bloom_filter( &predicate, &mut record_batch_stream_builder, - candidate_rgs, + candidates.as_ref(), &field_id_map, ) .await?; - if bloom_filtered.len() < candidate_rgs.len() { + if bloom_filtered.len() < candidates.len() { selected_row_group_indices = Some(bloom_filtered); } } diff --git a/crates/iceberg/src/arrow/reader/projection.rs b/crates/iceberg/src/arrow/reader/projection.rs index 8f2afe79e..2fa6bf18f 100644 --- a/crates/iceberg/src/arrow/reader/projection.rs +++ b/crates/iceberg/src/arrow/reader/projection.rs @@ -421,6 +421,16 @@ pub(super) fn apply_name_mapping_to_arrow_schema( "Schema already has field IDs - name mapping should not be applied" ); + let mut field_ids_by_name = HashMap::with_capacity(name_mapping.fields().len()); + for mapped_field in name_mapping.fields() { + for name in mapped_field.names() { + // Preserve the existing first-match behavior for duplicate aliases. + field_ids_by_name + .entry(name.as_str()) + .or_insert(mapped_field.field_id()); + } + } + let fields_with_mapped_ids: Vec<_> = arrow_schema .fields() .iter() @@ -432,15 +442,12 @@ pub(super) fn apply_name_mapping_to_arrow_schema( // If the field isn't in the mapping, leave it WITHOUT assigning an ID // (matching Java's behavior of returning the field unchanged). // Later, during projection, fields without IDs are filtered out. - let mapped_field_opt = name_mapping - .fields() - .iter() - .find(|f| f.names().contains(&field.name().to_string())); - let mut metadata = field.metadata().clone(); - if let Some(mapped_field) = mapped_field_opt - && let Some(field_id) = mapped_field.field_id() + if let Some(field_id) = field_ids_by_name + .get(field.name().as_str()) + .copied() + .flatten() { // Field found in mapping with a field_id → assign it metadata.insert(PARQUET_FIELD_ID_META_KEY.to_string(), field_id.to_string()); @@ -523,6 +530,34 @@ mod tests { }; use crate::{ErrorKind, Runtime}; + #[test] + fn test_name_mapping_duplicate_alias_preserves_first_match() { + let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new( + "name", + DataType::Utf8, + true, + )])); + for (first_id, expected) in [(Some(2), Some("2")), (None, None)] { + let name_mapping = NameMapping::new(vec![ + MappedField::new(first_id, vec!["name".to_string()], vec![]), + MappedField::new(Some(4), vec!["name".to_string()], vec![]), + ]); + + let mapped = + super::apply_name_mapping_to_arrow_schema(arrow_schema.clone(), &name_mapping) + .unwrap(); + + assert_eq!( + mapped.fields()[0] + .metadata() + .get(PARQUET_FIELD_ID_META_KEY) + .map(String::as_str), + expected, + "first matching field ID: {first_id:?}" + ); + } + } + #[test] fn test_arrow_projection_mask() { let schema = Arc::new( diff --git a/crates/iceberg/src/arrow/reader/row_filter.rs b/crates/iceberg/src/arrow/reader/row_filter.rs index 5d92f6bc7..4a99914a3 100644 --- a/crates/iceberg/src/arrow/reader/row_filter.rs +++ b/crates/iceberg/src/arrow/reader/row_filter.rs @@ -35,6 +35,7 @@ use crate::expr::visitors::bound_predicate_visitor::visit; use crate::expr::visitors::page_index_evaluator::PageIndexEvaluator; use crate::expr::visitors::row_group_metrics_evaluator::RowGroupMetricsEvaluator; use crate::spec::Schema; +use crate::{Error, ErrorKind}; impl ArrowReader { pub(super) fn get_row_filter( @@ -71,11 +72,21 @@ impl ArrowReader { parquet_metadata: &Arc<ParquetMetaData>, field_id_map: &HashMap<i32, usize>, snapshot_schema: &Schema, + candidate_row_groups: &[usize], ) -> Result<Vec<usize>> { let row_groups_metadata = parquet_metadata.row_groups(); - let mut results = Vec::with_capacity(row_groups_metadata.len()); - - for (idx, row_group_metadata) in row_groups_metadata.iter().enumerate() { + let mut results = Vec::with_capacity(candidate_row_groups.len()); + + for &idx in candidate_row_groups { + let row_group_metadata = row_groups_metadata.get(idx).ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!( + "row group index {idx} out of range ({} row groups)", + row_groups_metadata.len() + ), + ) + })?; if RowGroupMetricsEvaluator::eval( predicate, row_group_metadata, @@ -634,6 +645,67 @@ mod tests { assert_eq!(first_val, 100, "Task 2 should start with id=100, not id=0"); } + + // Combining byte-range and predicate row-group filtering should narrow the + // split's candidates instead of evaluating and intersecting independent lists. + // This task owns row groups 1 and 2, while the predicate can match row groups + // 0 and 2. Only row group 2 should be read. + let predicate = Reference::new("id") + .less_than(Datum::int(100)) + .or(Reference::new("id").greater_than_or_equal_to(Datum::int(200))) + .bind(schema.clone(), true) + .unwrap(); + + // Test the pruning decision directly, not just the final filtered rows. + let parquet_metadata = Arc::new((*metadata).clone()); + let field_id_map = HashMap::from([(1, 0)]); + for (candidates, expected) in [(vec![1, 2], vec![2]), (vec![0, 1], vec![0])] { + let selected = ArrowReader::get_selected_row_group_indices( + &predicate, + &parquet_metadata, + &field_id_map, + schema.as_ref(), + &candidates, + ) + .unwrap(); + assert_eq!(selected, expected); + } + + let task = FileScanTask::builder() + .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len()) + .with_start(rg1_start) + .with_length(file_end - rg1_start) + .with_data_file_path(file_path.clone()) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(schema.clone()) + .with_project_field_ids(vec![1]) + .with_predicate(Some(predicate)) + .with_record_count(Some(200)) + .with_case_sensitive(false) + .build() + .unwrap(); + + let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()) + .with_row_group_filtering_enabled(true) + .build(); + let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream; + let result = reader + .read(tasks) + .unwrap() + .stream() + .try_collect::<Vec<RecordBatch>>() + .await + .unwrap(); + + let mut ids = Vec::new(); + for batch in &result { + let id_col = batch + .column(0) + .as_primitive::<arrow_array::types::Int32Type>(); + ids.extend((0..id_col.len()).map(|index| id_col.value(index))); + } + + assert_eq!(ids, (200..300).collect::<Vec<_>>()); } /// A single data file split into multiple sub-row-group byte ranges (as Spark/Iceberg @@ -903,6 +975,24 @@ mod tests { ) } + #[test] + fn test_row_group_pruning_rejects_invalid_candidate() { + let schema = int_schema(); + let predicate = simple_predicate(schema.clone()); + let metadata = metadata_no_page_indexes(); + + let err = ArrowReader::get_selected_row_group_indices( + &predicate, + &metadata, + &field_id_map(), + schema.as_ref(), + &[0], + ) + .unwrap_err(); + + assert_eq!(err.kind(), crate::ErrorKind::DataInvalid); + } + /// Testing suite regarding: https://github.com/apache/iceberg-rust/issues/2452 /// Testing when: both indices are absent, some present, both present #[test]
