From 0c5c7d83b6d3c39baaa6b4c337cd2754d15c639e Mon Sep 17 00:00:00 2001 From: Huaijin Date: Sat, 8 Aug 2026 23:35:39 +0800 Subject: [PATCH 1/4] feat: preserve bitmap-backed RowSelection through ParquetAccessPlan Parquet 59 added a bitmap (BooleanBuffer) backing for RowSelection in addition to the RLE selector form. Previously ParquetAccessPlan always materialized selectors, so a caller-provided bitmap selection was converted to RLE before reaching the parquet reader, which is wasteful for fragmented selections. This change keeps bitmap-backed selections bitmap-backed end to end: * `try_new_from_overall_row_selection` splits a mask-backed selection per row group with `split_off`, which slices the BooleanBuffer instead of materializing selectors. * `into_overall_row_selection` builds a bitmap-backed overall selection when any row group selection is bitmap-backed, promoting selector-backed groups (e.g. ones page index pruning intersected). * `scan_selection` promotes the incoming selection to a bitmap when intersecting with an existing bitmap-backed selection, so the intersection is a bitwise AND and stays mask-backed. * `reverse_row_selection` slices and re-concatenates the bitmap for reverse scans, with a debug_assert guarding the row-count contract. --- .../tests/parquet/external_access_plan.rs | 80 ++++++ .../datasource-parquet/src/access_plan.rs | 272 +++++++++++++++--- datafusion/datasource-parquet/src/sort.rs | 66 +++++ 3 files changed, 381 insertions(+), 37 deletions(-) diff --git a/datafusion/core/tests/parquet/external_access_plan.rs b/datafusion/core/tests/parquet/external_access_plan.rs index 8fd9689ae3a8d..f7aa522f0b976 100644 --- a/datafusion/core/tests/parquet/external_access_plan.rs +++ b/datafusion/core/tests/parquet/external_access_plan.rs @@ -23,6 +23,7 @@ use std::sync::Arc; use crate::parquet::utils::MetricsFinder; use crate::parquet::{Scenario, create_data_batch}; +use arrow::buffer::BooleanBuffer; use arrow::datatypes::SchemaRef; use arrow::util::pretty::pretty_format_batches; use datafusion::common::Result; @@ -218,6 +219,85 @@ async fn row_selection_extension_spanning_row_groups() { assert_ne!(bytes_scanned, 0, "metrics : {parquet_metrics:#?}",); } +#[tokio::test] +async fn bitmap_row_selection_extension_spanning_row_groups() { + // Keep the same cross-row-group pattern as the selector-backed test, but + // provide it as a packed bitmap. The bitmap should remain packed through + // ParquetAccessPlan and reach the parquet reader unchanged. + let parquet_metrics = TestFull { + access_plan: None, + row_selection: Some(ParquetRowSelection::new(RowSelection::from_boolean_buffer( + BooleanBuffer::from(vec![ + false, false, false, false, true, true, true, false, false, false, + ]), + ))), + expected_rows: 3, + expected_output: Some(&[ + "+------+------------+", + "| utf8 | large_utf8 |", + "+------+------------+", + "| | |", + "| e | e |", + "| f | f |", + "+------+------------+", + ]), + predicate: None, + } + .run() + .await + .unwrap(); + + let bytes_scanned = metric_value(&parquet_metrics, "bytes_scanned").unwrap(); + assert_ne!(bytes_scanned, 0, "metrics : {parquet_metrics:#?}",); +} + +#[tokio::test] +async fn bitmap_row_selection_extension_with_predicate() { + // The bitmap selects rows 2 (c) and 4 (null) in row group 0 and rows 5 + // and 6 (e, f) in row group 1. The predicate `utf8 = 'e'` prunes row + // group 0 via statistics (its max value "d" sorts before "e"), so only + // the bitmap-selected rows of row group 1 are returned. + let parquet_metrics = TestFull { + access_plan: None, + row_selection: Some(ParquetRowSelection::new(RowSelection::from_boolean_buffer( + BooleanBuffer::from(vec![ + false, false, true, false, true, true, true, false, false, false, + ]), + ))), + expected_rows: 2, + expected_output: Some(&[ + "+------+------------+", + "| utf8 | large_utf8 |", + "+------+------------+", + "| e | e |", + "| f | f |", + "+------+------------+", + ]), + predicate: Some(col("utf8").eq(lit("e"))), + } + .run() + .await + .unwrap(); + + // Row group 0 was pruned by statistics even though the bitmap selected + // rows in it. + let row_groups_pruned_statistics = parquet_metrics + .sum_by_name("row_groups_pruned_statistics") + .unwrap(); + if let MetricValue::PruningMetrics { + pruning_metrics, .. + } = row_groups_pruned_statistics + { + assert_eq!( + pruning_metrics.pruned(), + 1, + "metrics : {parquet_metrics:#?}" + ); + } else { + unreachable!("metrics `row_groups_pruned_statistics` should exist") + } +} + #[tokio::test] async fn bad_row_selection_extension() { // selection specifies fewer rows than the file actually contains diff --git a/datafusion/datasource-parquet/src/access_plan.rs b/datafusion/datasource-parquet/src/access_plan.rs index 1e9bae0ff6ba3..682c950906a61 100644 --- a/datafusion/datasource-parquet/src/access_plan.rs +++ b/datafusion/datasource-parquet/src/access_plan.rs @@ -16,6 +16,7 @@ // under the License. use crate::sort::reverse_row_selection; +use arrow::array::BooleanBufferBuilder; use arrow::datatypes::Schema; use datafusion_common::{Result, assert_eq_or_internal_err, exec_err}; use datafusion_physical_expr::expressions::Column; @@ -203,11 +204,6 @@ impl OverallRowSelectionCursor { skip: sel.skip, }) } - - fn remaining_rows(self) -> usize { - self.current.map_or(0, |s| s.row_count) - + self.selector_iter.map(|s| s.row_count).sum::() - } } /// Accumulates the selector fragments that belong to one row group. @@ -256,6 +252,21 @@ impl RowGroupAccessBuilder { } } +/// Convert `selection` to a bitmap-backed [`RowSelection`]. +/// +/// No-op if the selection is already bitmap-backed. +fn into_mask_backed(selection: RowSelection) -> RowSelection { + if selection.as_mask().is_some() { + return selection; + } + let total_rows = selection.row_count() + selection.skipped_row_count(); + let mut mask = BooleanBufferBuilder::new(total_rows); + for selector in selection.iter() { + mask.append_n(selector.row_count, !selector.skip); + } + RowSelection::from_boolean_buffer(mask.finish()) +} + impl ParquetAccessPlan { /// Create a new `ParquetAccessPlan` that scans all row groups pub fn new_all(row_group_count: usize) -> Self { @@ -295,44 +306,67 @@ impl ParquetAccessPlan { /// Returns an error if the selection does not specify exactly the same /// number of rows as the file metadata. pub fn try_new_from_overall_row_selection( - selection: RowSelection, + mut selection: RowSelection, row_group_meta_data: &[RowGroupMetaData], ) -> Result { + let selection_rows = selection.row_count() + selection.skipped_row_count(); + let file_rows = row_group_meta_data + .iter() + .map(|rg| rg.num_rows() as usize) + .sum::(); + + if selection_rows != file_rows { + return exec_err!( + "Invalid Parquet RowSelection. File has {file_rows} rows, \ + but selection specifies {selection_rows} rows." + ); + } + + // `split_off` slices bitmap-backed selections without converting them + // to selectors. Keep partially selected row groups bitmap-backed so + // they can reach the parquet reader without an intermediate RLE. + if selection.as_mask().is_some() { + let row_groups = row_group_meta_data + .iter() + .map(|rg_meta| { + let rg_rows = rg_meta.num_rows() as usize; + let rg_selection = selection.split_off(rg_rows); + let selected_rows = rg_selection.row_count(); + + if selected_rows == 0 { + RowGroupAccess::Skip + } else if selected_rows == rg_rows { + RowGroupAccess::Scan + } else { + RowGroupAccess::Selection(rg_selection) + } + }) + .collect(); + + return Ok(Self::new(row_groups)); + } + // Keep this as a single pass over the selector stream rather than // repeatedly calling `RowSelection::split_off` per row group. The // `split_off` version is simpler, but it clones/retains substantially // more selector buffer capacity for highly fragmented selections. let mut cursor = OverallRowSelectionCursor::new(selection); - let mut selection_rows = 0usize; - let mut file_rows = 0usize; - let mut row_groups = Vec::with_capacity(row_group_meta_data.len()); for rg_meta in row_group_meta_data { let rg_rows = rg_meta.num_rows() as usize; - file_rows += rg_rows; let mut builder = RowGroupAccessBuilder::new(rg_rows); while builder.remaining > 0 { let Some(selector) = cursor.take(builder.remaining) else { break; }; - selection_rows += selector.row_count; builder.push(selector); } row_groups.push(builder.into_access()); } - selection_rows += cursor.remaining_rows(); - - if selection_rows != file_rows { - return exec_err!( - "Invalid Parquet RowSelection. File has {file_rows} rows, \ - but selection specifies {selection_rows} rows." - ); - } - Ok(Self::new(row_groups)) } @@ -392,6 +426,16 @@ impl ParquetAccessPlan { RowGroupAccess::Skip => RowGroupAccess::Skip, RowGroupAccess::Scan => RowGroupAccess::Selection(selection), RowGroupAccess::Selection(existing_selection) => { + // `RowSelection::intersection` only stays bitmap-backed when + // both sides are bitmap-backed, so promote the incoming + // selection (e.g. from page index pruning) to match an + // existing bitmap: the intersection is then a bitwise AND + // instead of a selector merge. + let selection = if existing_selection.as_mask().is_some() { + into_mask_backed(selection) + } else { + selection + }; RowGroupAccess::Selection(existing_selection.intersection(&selection)) } } @@ -419,6 +463,11 @@ impl ParquetAccessPlan { /// is returned for *all* the rows in the row groups that are not skipped. /// Thus it includes a `Select` selection for any [`RowGroupAccess::Scan`]. /// + /// If any [`RowGroupAccess::Selection`] is bitmap-backed + /// ([`RowSelection::as_mask`] returns `Some`), the overall selection is + /// bitmap-backed as well and any selector-backed selections are promoted + /// to bitmaps; otherwise the overall selection is selector-backed. + /// /// # Errors /// /// Returns an error if any specified row selection does not specify @@ -494,10 +543,7 @@ impl ParquetAccessPlan { let RowGroupAccess::Selection(selection) = rg else { continue; }; - let rows_in_selection = selection - .iter() - .map(|selection| selection.row_count) - .sum::(); + let rows_in_selection = selection.row_count() + selection.skipped_row_count(); let row_group_row_count = rg_meta.num_rows(); assert_eq_or_internal_err!( @@ -509,24 +555,61 @@ impl ParquetAccessPlan { ); } - let total_selection: RowSelection = self - .row_groups - .into_iter() - .zip(row_group_meta_data.iter()) - .flat_map(|(rg, rg_meta)| { + // A bitmap-backed group selection signals the plan originated from a + // bitmap (e.g. a mask-backed `ParquetRowSelection`), even if later + // pruning such as the page index intersected some groups back to + // selectors. Promote the selector-backed groups in that case so the + // reader still receives a bitmap-backed selection. + let any_selection_mask_backed = self.row_groups.iter().any(|rg| { + matches!(rg, RowGroupAccess::Selection(selection) if selection.as_mask().is_some()) + }); + + let total_selection = if any_selection_mask_backed { + let total_rows = self + .row_groups + .iter() + .zip(row_group_meta_data.iter()) + .filter(|(rg, _)| rg.should_scan()) + .map(|(_, rg_meta)| rg_meta.num_rows() as usize) + .sum(); + let mut mask = BooleanBufferBuilder::new(total_rows); + + for (rg, rg_meta) in self.row_groups.into_iter().zip(row_group_meta_data) { match rg { + // Skipped row groups are not passed to the parquet reader. + RowGroupAccess::Skip => {} + // Represent scanned row groups as all-set bitmap ranges. + RowGroupAccess::Scan => { + mask.append_n(rg_meta.num_rows() as usize, true) + } + RowGroupAccess::Selection(selection) => match selection.as_mask() { + Some(buffer) => mask.append_buffer(buffer), + // Promote selector-backed groups to bitmap ranges. + None => { + for selector in selection.iter() { + mask.append_n(selector.row_count, !selector.skip); + } + } + }, + } + } + + RowSelection::from_boolean_buffer(mask.finish()) + } else { + // Preserve the existing selector path when no row group + // selection is bitmap-backed. + self.row_groups + .into_iter() + .zip(row_group_meta_data.iter()) + .flat_map(|(rg, rg_meta)| match rg { RowGroupAccess::Skip => vec![], RowGroupAccess::Scan => { - // need a row group access to scan the entire row group (need row group counts) vec![RowSelector::select(rg_meta.num_rows() as usize)] } - RowGroupAccess::Selection(selection) => { - let selection: Vec = selection.into(); - selection - } - } - }) - .collect(); + RowGroupAccess::Selection(selection) => selection.into(), + }) + .collect() + }; Ok(Some(total_selection)) } @@ -787,6 +870,7 @@ impl PreparedAccessPlan { #[cfg(test)] mod test { use super::*; + use arrow::buffer::BooleanBuffer; use datafusion_common::assert_contains; use parquet::basic::LogicalType; use parquet::file::metadata::ColumnChunkMetaData; @@ -910,6 +994,63 @@ mod test { ); } + #[test] + fn test_mixed_backings_promote_to_mask() { + // One bitmap-backed group next to a selector-backed group (e.g. one + // that page index pruning intersected): the overall selection is + // promoted to a bitmap so the caller-provided mask backing survives + // to the reader. + let access_plan = ParquetAccessPlan::new(vec![ + RowGroupAccess::Selection(RowSelection::from_boolean_buffer( + BooleanBuffer::from(vec![ + true, false, false, false, false, false, false, false, false, false, + ]), + )), + RowGroupAccess::Selection(RowSelection::from(vec![ + RowSelector::select(10), + RowSelector::skip(10), + ])), + RowGroupAccess::Skip, + RowGroupAccess::Skip, + ]); + + let row_selection = access_plan + .into_overall_row_selection(&ROW_GROUP_METADATA) + .unwrap() + .unwrap(); + + let mut expected = vec![true]; + expected.extend(vec![false; 9]); + expected.extend(vec![true; 10]); + expected.extend(vec![false; 10]); + let expected = BooleanBuffer::from(expected); + assert_eq!(row_selection.as_mask(), Some(&expected)); + } + + #[test] + fn test_scan_selection_preserves_mask_backing() { + let mut access_plan = ParquetAccessPlan::new(vec![RowGroupAccess::Selection( + RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![ + true, true, false, false, true, true, false, false, true, true, + ])), + )]); + + // Intersect with a selector-backed selection, as produced by page + // index pruning. + access_plan.scan_selection( + 0, + RowSelection::from(vec![RowSelector::select(5), RowSelector::skip(5)]), + ); + + let RowGroupAccess::Selection(selection) = &access_plan.inner()[0] else { + panic!("expected a selection for row group 0"); + }; + let expected = BooleanBuffer::from(vec![ + true, true, false, false, true, false, false, false, false, false, + ]); + assert_eq!(selection.as_mask(), Some(&expected)); + } + #[test] fn test_new_from_overall_row_selection() { let row_selection = RowSelection::from(vec![ @@ -994,6 +1135,63 @@ mod test { ); } + #[test] + fn test_new_from_overall_mask_preserves_bitmap_backing() { + let partial_mask = vec![ + false, true, false, true, false, true, false, true, false, true, false, true, + false, true, false, true, false, true, false, true, false, true, false, true, + false, true, false, true, false, true, + ]; + let mut file_mask = vec![true; 10]; + file_mask.extend(vec![false; 20]); + file_mask.extend(&partial_mask); + file_mask.extend(vec![true; 40]); + + let access_plan = ParquetAccessPlan::try_new_from_overall_row_selection( + RowSelection::from_boolean_buffer(BooleanBuffer::from(file_mask)), + &ROW_GROUP_METADATA, + ) + .unwrap(); + + assert!(matches!(access_plan.inner()[0], RowGroupAccess::Scan)); + assert!(matches!(access_plan.inner()[1], RowGroupAccess::Skip)); + let RowGroupAccess::Selection(selection) = &access_plan.inner()[2] else { + panic!("expected a partial selection for row group 2"); + }; + let expected_partial_mask = BooleanBuffer::from(partial_mask.clone()); + assert_eq!(selection.as_mask(), Some(&expected_partial_mask)); + assert!(matches!(access_plan.inner()[3], RowGroupAccess::Scan)); + + // The fully skipped row group is omitted from the reader selection; + // scanned groups are represented as all-set bitmap ranges. + let overall = access_plan + .into_overall_row_selection(&ROW_GROUP_METADATA) + .unwrap() + .unwrap(); + let mut expected_overall = vec![true; 10]; + expected_overall.extend(partial_mask); + expected_overall.extend(vec![true; 40]); + let expected_overall = BooleanBuffer::from(expected_overall); + assert_eq!(overall.as_mask(), Some(&expected_overall)); + } + + #[test] + fn test_new_from_overall_mask_invalid_row_count() { + let row_selection = RowSelection::from_boolean_buffer(BooleanBuffer::new_set(99)); + + let err = ParquetAccessPlan::try_new_from_overall_row_selection( + row_selection, + &ROW_GROUP_METADATA, + ) + .unwrap_err() + .to_string(); + + assert_contains!( + err, + "Invalid Parquet RowSelection. File has 100 rows, but selection specifies 99 rows" + ); + } + #[test] fn test_invalid_too_few() { let access_plan = ParquetAccessPlan::new(vec![ diff --git a/datafusion/datasource-parquet/src/sort.rs b/datafusion/datasource-parquet/src/sort.rs index ea33fb0e2ecb2..e8cd2ea910f57 100644 --- a/datafusion/datasource-parquet/src/sort.rs +++ b/datafusion/datasource-parquet/src/sort.rs @@ -54,6 +54,29 @@ pub fn reverse_row_selection( ) -> Result { let rg_metadata = parquet_metadata.row_groups(); + // Splitting a bitmap-backed selection slices its BooleanBuffer without + // materializing selectors. Collecting the reversed slices concatenates + // them back into a bitmap-backed selection. + if row_selection.as_mask().is_some() { + let mut remaining = row_selection.clone(); + let mut row_group_selections = Vec::with_capacity(row_groups_to_scan.len()); + + for &rg_idx in row_groups_to_scan { + let num_rows = rg_metadata[rg_idx].num_rows() as usize; + row_group_selections.push(remaining.split_off(num_rows)); + } + + // `split_off` silently returns short slices when the selection runs + // out early, which would misalign the reversed row group boundaries. + debug_assert_eq!( + remaining.row_count() + remaining.skipped_row_count(), + 0, + "row selection covers more rows than the scanned row groups" + ); + + return Ok(row_group_selections.into_iter().rev().collect()); + } + // Build a mapping of row group index to its row range, but ONLY for // the row groups that are actually being scanned. // @@ -244,6 +267,7 @@ fn file_min_value(file: &PartitionedFile, col_idx: usize) -> Option mod tests { use crate::ParquetAccessPlan; use crate::RowGroupAccess; + use arrow::buffer::BooleanBuffer; use arrow::datatypes::{DataType, Field, Schema}; use bytes::Bytes; use parquet::arrow::ArrowWriter; @@ -358,6 +382,48 @@ mod tests { ); } + #[test] + fn test_prepared_access_plan_reverse_preserves_bitmap_backing() { + let metadata = create_test_metadata(vec![4, 3, 5, 2]); + let first_mask = vec![true, false, true, false]; + let third_mask = vec![false, true, true, false, true]; + let access_plan = ParquetAccessPlan::new(vec![ + RowGroupAccess::Selection(RowSelection::from_boolean_buffer( + BooleanBuffer::from(first_mask.clone()), + )), + RowGroupAccess::Skip, + RowGroupAccess::Selection(RowSelection::from_boolean_buffer( + BooleanBuffer::from(third_mask.clone()), + )), + RowGroupAccess::Scan, + ]); + + let prepared_plan = access_plan.prepare(metadata.row_groups()).unwrap(); + assert!( + prepared_plan + .row_selection + .as_ref() + .unwrap() + .as_mask() + .is_some() + ); + + let reversed_plan = prepared_plan.reverse(&metadata).unwrap(); + assert_eq!(reversed_plan.row_group_indexes, vec![3, 2, 0]); + let reversed_selection = reversed_plan.row_selection.unwrap(); + + // The fully scanned row group 3 becomes an all-set range at the + // front of the reversed selection. + let expected = BooleanBuffer::from( + [true, true] + .into_iter() + .chain(third_mask) + .chain(first_mask) + .collect::>(), + ); + assert_eq!(reversed_selection.as_mask(), Some(&expected)); + } + #[test] fn test_prepared_access_plan_reverse_multi_row_group_selection() { // Test: row selection spanning multiple row groups From 63a06a96a12f5c0abfecd49905a38d97309ce8fa Mon Sep 17 00:00:00 2001 From: Huaijin Date: Sun, 9 Aug 2026 12:51:10 +0800 Subject: [PATCH 2/4] update --- .../datasource-parquet/src/access_plan.rs | 155 ++++++++++++------ 1 file changed, 103 insertions(+), 52 deletions(-) diff --git a/datafusion/datasource-parquet/src/access_plan.rs b/datafusion/datasource-parquet/src/access_plan.rs index 682c950906a61..b036d99c59487 100644 --- a/datafusion/datasource-parquet/src/access_plan.rs +++ b/datafusion/datasource-parquet/src/access_plan.rs @@ -252,19 +252,29 @@ impl RowGroupAccessBuilder { } } +/// Return the total number of rows described by `selection`. +fn row_selection_len(selection: &RowSelection) -> usize { + match selection.as_mask() { + Some(mask) => mask.len(), + None => selection.iter().map(|selector| selector.row_count).sum(), + } +} + /// Convert `selection` to a bitmap-backed [`RowSelection`]. /// /// No-op if the selection is already bitmap-backed. fn into_mask_backed(selection: RowSelection) -> RowSelection { - if selection.as_mask().is_some() { - return selection; - } - let total_rows = selection.row_count() + selection.skipped_row_count(); - let mut mask = BooleanBufferBuilder::new(total_rows); - for selector in selection.iter() { - mask.append_n(selector.row_count, !selector.skip); + match selection.as_mask() { + Some(_) => selection, + None => { + let total_rows = row_selection_len(&selection); + let mut mask = BooleanBufferBuilder::new(total_rows); + for selector in selection.iter() { + mask.append_n(selector.row_count, !selector.skip); + } + RowSelection::from_boolean_buffer(mask.finish()) + } } - RowSelection::from_boolean_buffer(mask.finish()) } impl ParquetAccessPlan { @@ -309,7 +319,7 @@ impl ParquetAccessPlan { mut selection: RowSelection, row_group_meta_data: &[RowGroupMetaData], ) -> Result { - let selection_rows = selection.row_count() + selection.skipped_row_count(); + let selection_rows = row_selection_len(&selection); let file_rows = row_group_meta_data .iter() .map(|rg| rg.num_rows() as usize) @@ -322,11 +332,12 @@ impl ParquetAccessPlan { ); } - // `split_off` slices bitmap-backed selections without converting them - // to selectors. Keep partially selected row groups bitmap-backed so - // they can reach the parquet reader without an intermediate RLE. - if selection.as_mask().is_some() { - let row_groups = row_group_meta_data + let row_groups = match selection.as_mask() { + // `split_off` slices bitmap-backed selections without converting + // them to selectors. Keep partially selected row groups + // bitmap-backed so they can reach the parquet reader without an + // intermediate RLE. + Some(_) => row_group_meta_data .iter() .map(|rg_meta| { let rg_rows = rg_meta.num_rows() as usize; @@ -341,31 +352,32 @@ impl ParquetAccessPlan { RowGroupAccess::Selection(rg_selection) } }) - .collect(); - - return Ok(Self::new(row_groups)); - } - - // Keep this as a single pass over the selector stream rather than - // repeatedly calling `RowSelection::split_off` per row group. The - // `split_off` version is simpler, but it clones/retains substantially - // more selector buffer capacity for highly fragmented selections. - let mut cursor = OverallRowSelectionCursor::new(selection); + .collect(), + None => { + // Keep this as a single pass over the selector stream rather + // than repeatedly calling `RowSelection::split_off` per row + // group. The `split_off` version is simpler, but it + // clones/retains substantially more selector buffer capacity + // for highly fragmented selections. + let mut cursor = OverallRowSelectionCursor::new(selection); + + let mut row_groups = Vec::with_capacity(row_group_meta_data.len()); + for rg_meta in row_group_meta_data { + let rg_rows = rg_meta.num_rows() as usize; - let mut row_groups = Vec::with_capacity(row_group_meta_data.len()); - for rg_meta in row_group_meta_data { - let rg_rows = rg_meta.num_rows() as usize; + let mut builder = RowGroupAccessBuilder::new(rg_rows); + while builder.remaining > 0 { + let selector = cursor + .take(builder.remaining) + .expect("row selection length was validated"); + builder.push(selector); + } - let mut builder = RowGroupAccessBuilder::new(rg_rows); - while builder.remaining > 0 { - let Some(selector) = cursor.take(builder.remaining) else { - break; - }; - builder.push(selector); + row_groups.push(builder.into_access()); + } + row_groups } - - row_groups.push(builder.into_access()); - } + }; Ok(Self::new(row_groups)) } @@ -431,10 +443,9 @@ impl ParquetAccessPlan { // selection (e.g. from page index pruning) to match an // existing bitmap: the intersection is then a bitwise AND // instead of a selector merge. - let selection = if existing_selection.as_mask().is_some() { - into_mask_backed(selection) - } else { - selection + let selection = match existing_selection.as_mask() { + Some(_) => into_mask_backed(selection), + None => selection, }; RowGroupAccess::Selection(existing_selection.intersection(&selection)) } @@ -543,7 +554,7 @@ impl ParquetAccessPlan { let RowGroupAccess::Selection(selection) = rg else { continue; }; - let rows_in_selection = selection.row_count() + selection.skipped_row_count(); + let rows_in_selection = row_selection_len(selection); let row_group_row_count = rg_meta.num_rows(); assert_eq_or_internal_err!( @@ -1041,16 +1052,53 @@ mod test { 0, RowSelection::from(vec![RowSelector::select(5), RowSelector::skip(5)]), ); + // Intersect with an already mask-backed selection to exercise the + // no-op path in `into_mask_backed`. + access_plan.scan_selection( + 0, + RowSelection::from_boolean_buffer(BooleanBuffer::new_set(10)), + ); - let RowGroupAccess::Selection(selection) = &access_plan.inner()[0] else { - panic!("expected a selection for row group 0"); - }; + let selection = access_plan + .into_overall_row_selection(&ROW_GROUP_METADATA[..1]) + .unwrap() + .unwrap(); let expected = BooleanBuffer::from(vec![ true, true, false, false, true, false, false, false, false, false, ]); assert_eq!(selection.as_mask(), Some(&expected)); } + #[test] + fn test_scan_selection_preserves_selector_backing() { + let mut access_plan = ParquetAccessPlan::new(vec![RowGroupAccess::Selection( + RowSelection::from(vec![RowSelector::select(6), RowSelector::skip(4)]), + )]); + + access_plan.scan_selection( + 0, + RowSelection::from(vec![ + RowSelector::skip(2), + RowSelector::select(5), + RowSelector::skip(3), + ]), + ); + + let selection = access_plan + .into_overall_row_selection(&ROW_GROUP_METADATA[..1]) + .unwrap() + .unwrap(); + assert_eq!(selection.as_mask(), None); + assert_eq!( + selection, + RowSelection::from(vec![ + RowSelector::skip(2), + RowSelector::select(4), + RowSelector::skip(4), + ]) + ); + } + #[test] fn test_new_from_overall_row_selection() { let row_selection = RowSelection::from(vec![ @@ -1153,14 +1201,17 @@ mod test { ) .unwrap(); - assert!(matches!(access_plan.inner()[0], RowGroupAccess::Scan)); - assert!(matches!(access_plan.inner()[1], RowGroupAccess::Skip)); - let RowGroupAccess::Selection(selection) = &access_plan.inner()[2] else { - panic!("expected a partial selection for row group 2"); - }; - let expected_partial_mask = BooleanBuffer::from(partial_mask.clone()); - assert_eq!(selection.as_mask(), Some(&expected_partial_mask)); - assert!(matches!(access_plan.inner()[3], RowGroupAccess::Scan)); + assert_eq!( + access_plan, + ParquetAccessPlan::new(vec![ + RowGroupAccess::Scan, + RowGroupAccess::Skip, + RowGroupAccess::Selection(RowSelection::from_boolean_buffer( + BooleanBuffer::from(partial_mask.clone()), + )), + RowGroupAccess::Scan, + ]) + ); // The fully skipped row group is omitted from the reader selection; // scanned groups are represented as all-set bitmap ranges. From 62182a452dec88521a40e84f543b5daf4d21d164 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Sun, 9 Aug 2026 13:16:08 +0800 Subject: [PATCH 3/4] update --- .../datasource-parquet/src/access_plan.rs | 201 +++++++++++------- 1 file changed, 126 insertions(+), 75 deletions(-) diff --git a/datafusion/datasource-parquet/src/access_plan.rs b/datafusion/datasource-parquet/src/access_plan.rs index b036d99c59487..c1506a88eb888 100644 --- a/datafusion/datasource-parquet/src/access_plan.rs +++ b/datafusion/datasource-parquet/src/access_plan.rs @@ -17,6 +17,7 @@ use crate::sort::reverse_row_selection; use arrow::array::BooleanBufferBuilder; +use arrow::buffer::BooleanBuffer; use arrow::datatypes::Schema; use datafusion_common::{Result, assert_eq_or_internal_err, exec_err}; use datafusion_physical_expr::expressions::Column; @@ -161,16 +162,74 @@ impl RowGroupAccess { } } -/// Single-pass cursor over a file-level [`RowSelection`]. +/// Single-pass cursor that partitions a file-level [`RowSelection`] into row +/// groups while preserving its backing representation. +enum OverallRowSelectionCursor { + Mask { mask: BooleanBuffer, offset: usize }, + Selectors(SelectorRowSelectionCursor), +} + +impl OverallRowSelectionCursor { + fn new(selection: RowSelection) -> Self { + match selection.as_mask() { + Some(mask) => Self::Mask { + mask: mask.clone(), + offset: 0, + }, + None => Self::Selectors(SelectorRowSelectionCursor::new(selection)), + } + } + + /// Take the selection for the next row group. + /// + /// Returns `None` if the selection contains fewer than `row_group_rows` + /// remaining rows. + fn take_row_group(&mut self, row_group_rows: usize) -> Option { + match self { + Self::Mask { mask, offset } => { + if row_group_rows > mask.len() - *offset { + return None; + } + + let end = *offset + row_group_rows; + let row_group_mask = mask.slice(*offset, row_group_rows); + *offset = end; + let selected_rows = row_group_mask.count_set_bits(); + + Some(if selected_rows == 0 { + RowGroupAccess::Skip + } else if selected_rows == row_group_rows { + RowGroupAccess::Scan + } else { + RowGroupAccess::Selection(RowSelection::from_boolean_buffer( + row_group_mask, + )) + }) + } + Self::Selectors(cursor) => cursor.take_row_group(row_group_rows), + } + } + + /// Return the total number of rows in the original selection. + fn total_rows(self) -> usize { + match self { + Self::Mask { mask, .. } => mask.len(), + Self::Selectors(cursor) => cursor.total_rows(), + } + } +} + +/// Cursor over a selector-backed [`RowSelection`]. /// /// `take` returns the next selector fragment capped to the requested row count, /// splitting the current selector when it straddles a row group boundary. -struct OverallRowSelectionCursor { +struct SelectorRowSelectionCursor { selector_iter: std::vec::IntoIter, current: Option, + consumed_rows: usize, } -impl OverallRowSelectionCursor { +impl SelectorRowSelectionCursor { fn new(selection: RowSelection) -> Self { let selectors: Vec = selection.into(); let mut selector_iter = selectors.into_iter(); @@ -178,6 +237,7 @@ impl OverallRowSelectionCursor { Self { selector_iter, current, + consumed_rows: 0, } } @@ -190,6 +250,7 @@ impl OverallRowSelectionCursor { fn take(&mut self, max_rows: usize) -> Option { let sel = self.current?; let row_count = sel.row_count.min(max_rows); + self.consumed_rows += row_count; self.current = if row_count < sel.row_count { Some(RowSelector { row_count: sel.row_count - row_count, @@ -204,6 +265,23 @@ impl OverallRowSelectionCursor { skip: sel.skip, }) } + + fn take_row_group(&mut self, row_group_rows: usize) -> Option { + let mut builder = RowGroupAccessBuilder::new(row_group_rows); + while builder.remaining > 0 { + builder.push(self.take(builder.remaining)?); + } + Some(builder.into_access()) + } + + fn total_rows(self) -> usize { + self.consumed_rows + + self.current.map_or(0, |selector| selector.row_count) + + self + .selector_iter + .map(|selector| selector.row_count) + .sum::() + } } /// Accumulates the selector fragments that belong to one row group. @@ -316,69 +394,30 @@ impl ParquetAccessPlan { /// Returns an error if the selection does not specify exactly the same /// number of rows as the file metadata. pub fn try_new_from_overall_row_selection( - mut selection: RowSelection, + selection: RowSelection, row_group_meta_data: &[RowGroupMetaData], ) -> Result { - let selection_rows = row_selection_len(&selection); let file_rows = row_group_meta_data .iter() .map(|rg| rg.num_rows() as usize) .sum::(); - if selection_rows != file_rows { + let mut cursor = OverallRowSelectionCursor::new(selection); + let row_groups = row_group_meta_data + .iter() + .map_while(|rg_meta| cursor.take_row_group(rg_meta.num_rows() as usize)) + .collect::>(); + + // For selector-backed selections this computes the total while + // consuming the stream above, rather than requiring a separate pass. + let selection_rows = cursor.total_rows(); + if row_groups.len() != row_group_meta_data.len() || selection_rows != file_rows { return exec_err!( "Invalid Parquet RowSelection. File has {file_rows} rows, \ but selection specifies {selection_rows} rows." ); } - let row_groups = match selection.as_mask() { - // `split_off` slices bitmap-backed selections without converting - // them to selectors. Keep partially selected row groups - // bitmap-backed so they can reach the parquet reader without an - // intermediate RLE. - Some(_) => row_group_meta_data - .iter() - .map(|rg_meta| { - let rg_rows = rg_meta.num_rows() as usize; - let rg_selection = selection.split_off(rg_rows); - let selected_rows = rg_selection.row_count(); - - if selected_rows == 0 { - RowGroupAccess::Skip - } else if selected_rows == rg_rows { - RowGroupAccess::Scan - } else { - RowGroupAccess::Selection(rg_selection) - } - }) - .collect(), - None => { - // Keep this as a single pass over the selector stream rather - // than repeatedly calling `RowSelection::split_off` per row - // group. The `split_off` version is simpler, but it - // clones/retains substantially more selector buffer capacity - // for highly fragmented selections. - let mut cursor = OverallRowSelectionCursor::new(selection); - - let mut row_groups = Vec::with_capacity(row_group_meta_data.len()); - for rg_meta in row_group_meta_data { - let rg_rows = rg_meta.num_rows() as usize; - - let mut builder = RowGroupAccessBuilder::new(rg_rows); - while builder.remaining > 0 { - let selector = cursor - .take(builder.remaining) - .expect("row selection length was validated"); - builder.push(selector); - } - - row_groups.push(builder.into_access()); - } - row_groups - } - }; - Ok(Self::new(row_groups)) } @@ -1135,19 +1174,25 @@ mod test { #[test] fn test_new_from_overall_row_selection_invalid_row_count() { - let row_selection = RowSelection::from(vec![RowSelector::select(99)]); + for selection_rows in [99, 101] { + let row_selection = + RowSelection::from(vec![RowSelector::select(selection_rows)]); - let err = ParquetAccessPlan::try_new_from_overall_row_selection( - row_selection, - &ROW_GROUP_METADATA, - ) - .unwrap_err() - .to_string(); + let err = ParquetAccessPlan::try_new_from_overall_row_selection( + row_selection, + &ROW_GROUP_METADATA, + ) + .unwrap_err() + .to_string(); - assert_contains!( - err, - "Invalid Parquet RowSelection. File has 100 rows, but selection specifies 99 rows" - ); + assert_contains!( + err, + format!( + "Invalid Parquet RowSelection. File has 100 rows, \ + but selection specifies {selection_rows} rows" + ) + ); + } } #[test] @@ -1228,19 +1273,25 @@ mod test { #[test] fn test_new_from_overall_mask_invalid_row_count() { - let row_selection = RowSelection::from_boolean_buffer(BooleanBuffer::new_set(99)); + for selection_rows in [99, 101] { + let row_selection = + RowSelection::from_boolean_buffer(BooleanBuffer::new_set(selection_rows)); - let err = ParquetAccessPlan::try_new_from_overall_row_selection( - row_selection, - &ROW_GROUP_METADATA, - ) - .unwrap_err() - .to_string(); + let err = ParquetAccessPlan::try_new_from_overall_row_selection( + row_selection, + &ROW_GROUP_METADATA, + ) + .unwrap_err() + .to_string(); - assert_contains!( - err, - "Invalid Parquet RowSelection. File has 100 rows, but selection specifies 99 rows" - ); + assert_contains!( + err, + format!( + "Invalid Parquet RowSelection. File has 100 rows, \ + but selection specifies {selection_rows} rows" + ) + ); + } } #[test] From 605d2fb83b2c5b6ba6b26684d37ddb4f4963b41f Mon Sep 17 00:00:00 2001 From: Huaijin Date: Sun, 9 Aug 2026 13:32:31 +0800 Subject: [PATCH 4/4] update comment --- .../tests/parquet/external_access_plan.rs | 12 +--- .../datasource-parquet/src/access_plan.rs | 57 +++++-------------- datafusion/datasource-parquet/src/sort.rs | 10 +--- 3 files changed, 20 insertions(+), 59 deletions(-) diff --git a/datafusion/core/tests/parquet/external_access_plan.rs b/datafusion/core/tests/parquet/external_access_plan.rs index f7aa522f0b976..27387448296e2 100644 --- a/datafusion/core/tests/parquet/external_access_plan.rs +++ b/datafusion/core/tests/parquet/external_access_plan.rs @@ -221,9 +221,7 @@ async fn row_selection_extension_spanning_row_groups() { #[tokio::test] async fn bitmap_row_selection_extension_spanning_row_groups() { - // Keep the same cross-row-group pattern as the selector-backed test, but - // provide it as a packed bitmap. The bitmap should remain packed through - // ParquetAccessPlan and reach the parquet reader unchanged. + // Repeat the cross-row-group case with a bitmap-backed selection. let parquet_metrics = TestFull { access_plan: None, row_selection: Some(ParquetRowSelection::new(RowSelection::from_boolean_buffer( @@ -253,10 +251,7 @@ async fn bitmap_row_selection_extension_spanning_row_groups() { #[tokio::test] async fn bitmap_row_selection_extension_with_predicate() { - // The bitmap selects rows 2 (c) and 4 (null) in row group 0 and rows 5 - // and 6 (e, f) in row group 1. The predicate `utf8 = 'e'` prunes row - // group 0 via statistics (its max value "d" sorts before "e"), so only - // the bitmap-selected rows of row group 1 are returned. + // Pruning removes row group 0, leaving the bitmap-selected rows in group 1. let parquet_metrics = TestFull { access_plan: None, row_selection: Some(ParquetRowSelection::new(RowSelection::from_boolean_buffer( @@ -279,8 +274,7 @@ async fn bitmap_row_selection_extension_with_predicate() { .await .unwrap(); - // Row group 0 was pruned by statistics even though the bitmap selected - // rows in it. + // Verify that statistics pruned row group 0. let row_groups_pruned_statistics = parquet_metrics .sum_by_name("row_groups_pruned_statistics") .unwrap(); diff --git a/datafusion/datasource-parquet/src/access_plan.rs b/datafusion/datasource-parquet/src/access_plan.rs index c1506a88eb888..544aa9a2e0554 100644 --- a/datafusion/datasource-parquet/src/access_plan.rs +++ b/datafusion/datasource-parquet/src/access_plan.rs @@ -162,8 +162,7 @@ impl RowGroupAccess { } } -/// Single-pass cursor that partitions a file-level [`RowSelection`] into row -/// groups while preserving its backing representation. +/// Splits an overall selection into row groups in one pass. enum OverallRowSelectionCursor { Mask { mask: BooleanBuffer, offset: usize }, Selectors(SelectorRowSelectionCursor), @@ -180,10 +179,7 @@ impl OverallRowSelectionCursor { } } - /// Take the selection for the next row group. - /// - /// Returns `None` if the selection contains fewer than `row_group_rows` - /// remaining rows. + /// Takes the next row group, or `None` if too few rows remain. fn take_row_group(&mut self, row_group_rows: usize) -> Option { match self { Self::Mask { mask, offset } => { @@ -210,7 +206,6 @@ impl OverallRowSelectionCursor { } } - /// Return the total number of rows in the original selection. fn total_rows(self) -> usize { match self { Self::Mask { mask, .. } => mask.len(), @@ -219,7 +214,7 @@ impl OverallRowSelectionCursor { } } -/// Cursor over a selector-backed [`RowSelection`]. +/// Cursor for a selector-backed selection. /// /// `take` returns the next selector fragment capped to the requested row count, /// splitting the current selector when it straddles a row group boundary. @@ -330,7 +325,7 @@ impl RowGroupAccessBuilder { } } -/// Return the total number of rows described by `selection`. +/// Returns the selection length. fn row_selection_len(selection: &RowSelection) -> usize { match selection.as_mask() { Some(mask) => mask.len(), @@ -338,9 +333,7 @@ fn row_selection_len(selection: &RowSelection) -> usize { } } -/// Convert `selection` to a bitmap-backed [`RowSelection`]. -/// -/// No-op if the selection is already bitmap-backed. +/// Converts to bitmap backing if needed. fn into_mask_backed(selection: RowSelection) -> RowSelection { match selection.as_mask() { Some(_) => selection, @@ -408,8 +401,7 @@ impl ParquetAccessPlan { .map_while(|rg_meta| cursor.take_row_group(rg_meta.num_rows() as usize)) .collect::>(); - // For selector-backed selections this computes the total while - // consuming the stream above, rather than requiring a separate pass. + // Count selector rows during traversal. let selection_rows = cursor.total_rows(); if row_groups.len() != row_group_meta_data.len() || selection_rows != file_rows { return exec_err!( @@ -477,11 +469,7 @@ impl ParquetAccessPlan { RowGroupAccess::Skip => RowGroupAccess::Skip, RowGroupAccess::Scan => RowGroupAccess::Selection(selection), RowGroupAccess::Selection(existing_selection) => { - // `RowSelection::intersection` only stays bitmap-backed when - // both sides are bitmap-backed, so promote the incoming - // selection (e.g. from page index pruning) to match an - // existing bitmap: the intersection is then a bitwise AND - // instead of a selector merge. + // Keep intersections bitmap-backed when the existing selection is. let selection = match existing_selection.as_mask() { Some(_) => into_mask_backed(selection), None => selection, @@ -513,10 +501,7 @@ impl ParquetAccessPlan { /// is returned for *all* the rows in the row groups that are not skipped. /// Thus it includes a `Select` selection for any [`RowGroupAccess::Scan`]. /// - /// If any [`RowGroupAccess::Selection`] is bitmap-backed - /// ([`RowSelection::as_mask`] returns `Some`), the overall selection is - /// bitmap-backed as well and any selector-backed selections are promoted - /// to bitmaps; otherwise the overall selection is selector-backed. + /// The result is bitmap-backed if any row-group selection is bitmap-backed. /// /// # Errors /// @@ -605,11 +590,7 @@ impl ParquetAccessPlan { ); } - // A bitmap-backed group selection signals the plan originated from a - // bitmap (e.g. a mask-backed `ParquetRowSelection`), even if later - // pruning such as the page index intersected some groups back to - // selectors. Promote the selector-backed groups in that case so the - // reader still receives a bitmap-backed selection. + // Preserve bitmap backing across mixed row-group selections. let any_selection_mask_backed = self.row_groups.iter().any(|rg| { matches!(rg, RowGroupAccess::Selection(selection) if selection.as_mask().is_some()) }); @@ -626,15 +607,12 @@ impl ParquetAccessPlan { for (rg, rg_meta) in self.row_groups.into_iter().zip(row_group_meta_data) { match rg { - // Skipped row groups are not passed to the parquet reader. RowGroupAccess::Skip => {} - // Represent scanned row groups as all-set bitmap ranges. RowGroupAccess::Scan => { mask.append_n(rg_meta.num_rows() as usize, true) } RowGroupAccess::Selection(selection) => match selection.as_mask() { Some(buffer) => mask.append_buffer(buffer), - // Promote selector-backed groups to bitmap ranges. None => { for selector in selection.iter() { mask.append_n(selector.row_count, !selector.skip); @@ -646,8 +624,7 @@ impl ParquetAccessPlan { RowSelection::from_boolean_buffer(mask.finish()) } else { - // Preserve the existing selector path when no row group - // selection is bitmap-backed. + // Keep selector backing when possible. self.row_groups .into_iter() .zip(row_group_meta_data.iter()) @@ -1046,10 +1023,7 @@ mod test { #[test] fn test_mixed_backings_promote_to_mask() { - // One bitmap-backed group next to a selector-backed group (e.g. one - // that page index pruning intersected): the overall selection is - // promoted to a bitmap so the caller-provided mask backing survives - // to the reader. + // Mixed backing produces a bitmap-backed overall selection. let access_plan = ParquetAccessPlan::new(vec![ RowGroupAccess::Selection(RowSelection::from_boolean_buffer( BooleanBuffer::from(vec![ @@ -1085,14 +1059,12 @@ mod test { ])), )]); - // Intersect with a selector-backed selection, as produced by page - // index pruning. + // Simulate selector-backed page pruning. access_plan.scan_selection( 0, RowSelection::from(vec![RowSelector::select(5), RowSelector::skip(5)]), ); - // Intersect with an already mask-backed selection to exercise the - // no-op path in `into_mask_backed`. + // Exercise the already bitmap-backed path. access_plan.scan_selection( 0, RowSelection::from_boolean_buffer(BooleanBuffer::new_set(10)), @@ -1258,8 +1230,7 @@ mod test { ]) ); - // The fully skipped row group is omitted from the reader selection; - // scanned groups are represented as all-set bitmap ranges. + // Skipped groups are omitted; scanned groups become set bits. let overall = access_plan .into_overall_row_selection(&ROW_GROUP_METADATA) .unwrap() diff --git a/datafusion/datasource-parquet/src/sort.rs b/datafusion/datasource-parquet/src/sort.rs index e8cd2ea910f57..ec846168ca074 100644 --- a/datafusion/datasource-parquet/src/sort.rs +++ b/datafusion/datasource-parquet/src/sort.rs @@ -54,9 +54,7 @@ pub fn reverse_row_selection( ) -> Result { let rg_metadata = parquet_metadata.row_groups(); - // Splitting a bitmap-backed selection slices its BooleanBuffer without - // materializing selectors. Collecting the reversed slices concatenates - // them back into a bitmap-backed selection. + // Reverse bitmap slices without materializing selectors. if row_selection.as_mask().is_some() { let mut remaining = row_selection.clone(); let mut row_group_selections = Vec::with_capacity(row_groups_to_scan.len()); @@ -66,8 +64,7 @@ pub fn reverse_row_selection( row_group_selections.push(remaining.split_off(num_rows)); } - // `split_off` silently returns short slices when the selection runs - // out early, which would misalign the reversed row group boundaries. + // No rows should remain after splitting all row groups. debug_assert_eq!( remaining.row_count() + remaining.skipped_row_count(), 0, @@ -412,8 +409,7 @@ mod tests { assert_eq!(reversed_plan.row_group_indexes, vec![3, 2, 0]); let reversed_selection = reversed_plan.row_selection.unwrap(); - // The fully scanned row group 3 becomes an all-set range at the - // front of the reversed selection. + // Fully scanned row group 3 becomes the leading set bits. let expected = BooleanBuffer::from( [true, true] .into_iter()