alamb commented on code in PR #10738: URL: https://github.com/apache/datafusion/pull/10738#discussion_r1623202606
########## datafusion/core/src/datasource/physical_plan/parquet/mod.rs: ########## @@ -152,8 +154,9 @@ pub use writer::plan_to_parquet; /// the file. /// /// * Step 3: The `ParquetOpener` gets the [`ParquetMetaData`] (file metadata) -/// via [`ParquetFileReaderFactory`] and applies any predicates and projections -/// to determine what pages must be read. +/// via [`ParquetFileReaderFactory`], creating a [`ParquetAccessPlan`] by Review Comment: The main change / reason for this PR is to encapsulate the decision of what row groups / what row selection to scan into a single struct (in this case `ParquetAccessPlan`) so that that the information can be passed in from the outside as well Aka `ParquetAccessPlan` will be the struct passed in the API described in https://github.com/apache/datafusion/issues/9929 ########## datafusion/core/src/datasource/physical_plan/parquet/page_filter.rs: ########## @@ -282,86 +289,70 @@ fn find_column_index( parquet_column(parquet_schema, arrow_schema, column.name()).map(|x| x.0) } -/// Intersects the [`RowSelector`]s -/// -/// For exampe, given: -/// * `RowSelector1: [ Skip(0~199), Read(200~299)]` -/// * `RowSelector2: [ Skip(0~99), Read(100~249), Skip(250~299)]` -/// -/// The final selection is the intersection of these `RowSelector`s: -/// * `final_selection:[ Skip(0~199), Read(200~249), Skip(250~299)]` -fn combine_multi_col_selection(row_selections: Vec<Vec<RowSelector>>) -> RowSelection { Review Comment: this intersection is now done in `ParquetAccessPlan::scan_selection` ########## datafusion/core/src/datasource/physical_plan/parquet/opener.rs: ########## @@ -163,32 +166,38 @@ impl FileOpener for ParquetOpener { } } + let mut access_plan = row_groups.build(); + // page index pruning: if all data on individual pages can // be ruled using page metadata, rows from other columns // with that range can be skipped as well - if enable_page_index && !row_groups.is_empty() { + if enable_page_index && !access_plan.is_empty() { if let Some(p) = page_pruning_predicate { - let pruned = p.prune( + access_plan = p.prune_plan_with_page_index( + access_plan, &file_schema, builder.parquet_schema(), - &row_groups, file_metadata.as_ref(), &file_metrics, - )?; - if let Some(row_selection) = pruned { - builder = builder.with_row_selection(row_selection); - } + ); } } + let row_group_indexes = access_plan.row_group_indexes(); Review Comment: while the rationale is to pass in scan restrictions from outside, I actually think encapsulating the logic for scan / skip into `ParquetAccessPlan` has made the existing logic easier to follow as well ########## datafusion/core/src/datasource/physical_plan/parquet/page_filter.rs: ########## @@ -129,105 +129,94 @@ impl PagePruningPredicate { Ok(Self { predicates }) } - /// Returns a [`RowSelection`] for the given file - pub fn prune( + /// Returns an updated [`ParquetAccessPlan`] by applying predicates to the + /// parquet page index, if any + pub fn prune_plan_with_page_index( &self, + mut access_plan: ParquetAccessPlan, arrow_schema: &Schema, parquet_schema: &SchemaDescriptor, - row_groups: &RowGroupSet, file_metadata: &ParquetMetaData, file_metrics: &ParquetFileMetrics, - ) -> Result<Option<RowSelection>> { + ) -> ParquetAccessPlan { // scoped timer updates on drop let _timer_guard = file_metrics.page_index_eval_time.timer(); if self.predicates.is_empty() { - return Ok(None); + return access_plan; } let page_index_predicates = &self.predicates; let groups = file_metadata.row_groups(); if groups.is_empty() { - return Ok(None); + return access_plan; Review Comment: I reworked this logic in two ways: 1. The loop is now for each row group (rather than for each predicate) 2. The row selection is incrementally built up for each row group rather than being created all at the end. I think this logic is clearer and also is no less performant (the same work is still done) ########## datafusion/core/src/datasource/physical_plan/parquet/row_groups.rs: ########## @@ -36,58 +36,34 @@ use crate::datasource::physical_plan::parquet::statistics::{ }; use crate::physical_optimizer::pruning::{PruningPredicate, PruningStatistics}; -use super::ParquetFileMetrics; +use super::{ParquetAccessPlan, ParquetFileMetrics}; -/// Tracks which RowGroups within a parquet file should be scanned. +/// Calculates which RowGroups within a parquet file should be scanned. /// -/// This struct encapsulates the various types of pruning that can be applied to +/// This struct implements the various types of pruning that are applied to /// a set of row groups within a parquet file, progressively narrowing down the /// set of row groups that should be scanned. -#[derive(Debug, PartialEq)] -pub struct RowGroupSet { - /// `row_groups[i]` is true if the i-th row group should be scanned - row_groups: Vec<bool>, +#[derive(Debug, Clone, PartialEq)] Review Comment: this extracts the plan of what to scan from the logic that prunes it based on row group statistic metadat ########## datafusion/core/src/datasource/physical_plan/parquet/access_plan.rs: ########## @@ -0,0 +1,399 @@ +// 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 parquet::arrow::arrow_reader::{RowSelection, RowSelector}; +use parquet::file::metadata::RowGroupMetaData; + +/// A selection of rows and row groups within a ParquetFile to decode. +/// +/// A `ParquetAccessPlan` is used to limits the row groups and data pages a `ParquetExec` +/// will read and decode and this improve performance. +/// +/// Note that page level pruning based on ArrowPredicate is applied after all of +/// these selections +/// +/// # Example +/// +/// For example, given a Parquet file with 4 row groups, a `ParquetAccessPlan` +/// can be used to specify skipping row group 0 and 2, scanning a range of rows +/// in row group 1, and scanning all rows in row group 3 as follows: +/// +/// ```rust +/// # use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; +/// # use datafusion::datasource::physical_plan::parquet::ParquetAccessPlan; +/// // Default to scan all row groups +/// let mut access_plan = ParquetAccessPlan::new_all(4); +/// access_plan.skip(0); // skip row group +/// // Use parquet reader RowSelector to specify scanning rows 100-200 and 350-400 +/// let row_selection = RowSelection::from(vec![ +/// RowSelector::skip(100), +/// RowSelector::select(100), +/// RowSelector::skip(150), +/// RowSelector::select(50), +/// ]); +/// access_plan.scan_selection(1, row_selection); +/// access_plan.skip(2); // skip row group 2 +/// // row group 3 is scanned by default +/// ``` +/// +/// The resulting plan would look like: +/// +/// ```text +/// ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐ +/// +/// │ │ SKIP +/// +/// └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘ +/// Row Group 0 +/// ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐ +/// ┌────────────────┐ SCAN ONLY ROWS +/// │└────────────────┘ │ 100-200 +/// ┌────────────────┐ 350-400 +/// │└────────────────┘ │ +/// ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ +/// Row Group 1 +/// ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐ +/// SKIP +/// │ │ +/// +/// └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘ +/// Row Group 2 +/// ┌───────────────────┐ +/// │ │ SCAN ALL ROWS +/// │ │ +/// │ │ +/// └───────────────────┘ +/// Row Group 3 +/// ``` +#[derive(Debug, Clone, PartialEq)] +pub struct ParquetAccessPlan { + /// How to access the i-th row group + row_groups: Vec<RowGroupAccess>, +} + +/// Describes how the parquet reader will access a row group +#[derive(Debug, Clone, PartialEq)] +pub enum RowGroupAccess { + /// Do not read the row group at all + Skip, + /// Read all rows from the row group + Scan, + /// Scan only the specified rows within the row group + Selection(RowSelection), +} + +impl RowGroupAccess { + /// Return true if this row group should be scanned + pub fn should_scan(&self) -> bool { + match self { + RowGroupAccess::Skip => false, + RowGroupAccess::Scan | RowGroupAccess::Selection(_) => true, + } + } +} + +impl ParquetAccessPlan { + /// Create a new `ParquetAccessPlan` that scans all row groups + pub fn new_all(row_group_count: usize) -> Self { + Self { + row_groups: vec![RowGroupAccess::Scan; row_group_count], + } + } + + /// Create a new `ParquetAccessPlan` that scans no row groups + pub fn new_none(row_group_count: usize) -> Self { + Self { + row_groups: vec![RowGroupAccess::Skip; row_group_count], + } + } + + /// Create a new `ParquetAccessPlan` from the specified [`RowGroupAccess`]es + pub fn new(row_groups: Vec<RowGroupAccess>) -> Self { + Self { row_groups } + } + + /// Set the i-th row group to the specified [`RowGroupAccess`] + pub fn set(&mut self, idx: usize, access: RowGroupAccess) { + self.row_groups[idx] = access; + } + + /// skips the i-th row group (should not be scanned) + pub fn skip(&mut self, idx: usize) { + self.set(idx, RowGroupAccess::Skip); + } + + /// Return true if the i-th row group should be scanned + pub fn should_scan(&self, idx: usize) -> bool { + self.row_groups[idx].should_scan() + } + + /// Set to scan only the [`RowSelection`] in the specified row group. + /// + /// Behavior is different depending on the existing access + /// * [`RowGroupAccess::Skip`]: does nothing + /// * [`RowGroupAccess::Scan`]: Updates to scan only the rows in the `RowSelection` + /// * [`RowGroupAccess::Selection`]: Updates to scan only the intersection of the existing selection and the new selection + pub fn scan_selection(&mut self, idx: usize, selection: RowSelection) { + self.row_groups[idx] = match &self.row_groups[idx] { + // already skipping the entire row group + RowGroupAccess::Skip => RowGroupAccess::Skip, + RowGroupAccess::Scan => RowGroupAccess::Selection(selection), + RowGroupAccess::Selection(existing_selection) => { + RowGroupAccess::Selection(existing_selection.intersection(&selection)) + } + } + } + + /// Return the overall `RowSelection` for all scanned row groups + /// + /// This is used to compute the row selection for the parquet reader. See + /// [`ArrowReaderBuilder::with_row_selection`] for more details. + /// + /// Returns + /// * `None` if there are no [`RowGroupAccess::Selection`] + /// * `Some(selection)` if there are [`RowGroupAccess::Selection`]s + /// + /// The returned selection represents which rows to scan across any row + /// row groups which are not skipped. + /// + /// # Example + /// + /// Given an access plan like this: + /// + /// ```text + /// Scan (scan all row group 0) + /// Skip (skip row group 1) + /// Select 50-100 (scan rows 50-100 in row group 2) + /// ``` + /// + /// Assuming each row group has 1000 rows, the resulting row selection would + /// be the rows to scan in row group 0 and 2: + /// + /// ```text + /// Select 1000 (scan all rows in row group 0) + /// Select 50-100 (scan rows 50-100 in row group 2) + /// ``` + /// + /// Note there is no entry for the (entirely) skipped row group 1. + /// + /// [`ArrowReaderBuilder::with_row_selection`]: parquet::arrow::arrow_reader::ArrowReaderBuilder::with_row_selection + pub fn into_overall_row_selection( + self, + row_group_meta_data: &[RowGroupMetaData], + ) -> Option<RowSelection> { + assert_eq!(row_group_meta_data.len(), self.row_groups.len()); + if !self + .row_groups + .iter() + .any(|rg| matches!(rg, RowGroupAccess::Selection(_))) + { + return None; + } + + let total_selection: RowSelection = 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<RowSelector> = selection.into(); + selection + } + } + }) + .collect(); + + Some(total_selection) + } + + /// Return an iterator over the row group indexes that should be scanned + pub fn row_group_index_iter(&self) -> impl Iterator<Item = usize> + '_ { + self.row_groups.iter().enumerate().filter_map(|(idx, b)| { + if b.should_scan() { + Some(idx) + } else { + None + } + }) + } + + /// Return a vec of all row group indexes to scan + pub fn row_group_indexes(&self) -> Vec<usize> { + self.row_group_index_iter().collect() + } + + /// Return the total number of row groups (not the total number or groups to + /// scan) + pub fn len(&self) -> usize { + self.row_groups.len() + } + + /// Return true if there are no row groups + pub fn is_empty(&self) -> bool { + self.row_groups.is_empty() + } + + /// Get a reference to the inner accesses + pub fn inner(&self) -> &[RowGroupAccess] { + &self.row_groups + } + + /// Covert into the inner row group accesses + pub fn into_inner(self) -> Vec<RowGroupAccess> { + self.row_groups + } +} + +#[cfg(test)] +mod test { + use super::*; + use parquet::basic::LogicalType; + use parquet::file::metadata::ColumnChunkMetaData; + use parquet::schema::types::{SchemaDescPtr, SchemaDescriptor}; + use std::sync::{Arc, OnceLock}; + + #[test] + fn test_overall_row_selection_only_scans() { Review Comment: I added these tests because I found the existing tests don't cover the case where there is a mix of scan some entire row groups and only scan the rows of another (happens when the statistics can't be extracted or there is some error evaluating the page pruning predicate only on some row groups) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
