Refactor parquet row group pruning into a struct (use new statistics API, part 1) #10607

alamb · 2024-05-21T18:27:52Z

Which issue does this PR close?

Rationale for this change

This PR prepares switching the parquet reader to use the new statistics extraction API added in #10537.

To avoid a massive, I broke it into its own PR for easier review. This PR refactors the code around, but doesn't change any behavior. The next PR will change the pruning API to use the new statistics extraction API.

I also think this change makes the code easier to understand and maintain as it uses a documented struct to represent the set of RowGroups to scan rather than a Vec<usize>.

What changes are included in this PR?

Add RowGroupSet to encapsulate the set of row groups which is being pruned
Move pruning code into RowGroupSet methods
Update code + tests to use RowGroupSet

Are these changes tested?

Yes, existing tests are updated to use the new structures

Are there any user-facing changes?

No these are all internal changes

alamb

Reviewing with whitespace blind diff makes the changes in this PR more obvious I think

https://github.com/apache/datafusion/pull/10607/files?w=1

alamb · 2024-05-21T18:54:36Z

datafusion/core/src/datasource/physical_plan/parquet/mod.rs

-            );
+            let rg_metadata = file_metadata.row_groups();
+            // track which row groups to actually read
+            let mut row_groups = RowGroupSet::new(rg_metadata.len());


The whole point of this PR is to move the set of RowGroups that are still under consideration into a RowGroupSet rather than a Vec<usize>

This makes it clearer what is going on in my opinion, as well as sets up up for vectorized pruning evaluation in the next PR

alamb · 2024-05-21T19:03:33Z

datafusion/core/src/datasource/physical_plan/parquet/row_groups.rs

+    ///
+    /// Updates this set to mark row groups that should not be scanned
+    pub fn prune_by_range(&mut self, groups: &[RowGroupMetaData], range: &FileRange) {
+        for (idx, metadata) in groups.iter().enumerate() {


I moved the range filtering into its own function to make it clearer what is going on (it has nothing to do with statistics pruning)

+1 this looks more readable

alamb · 2024-05-21T19:04:04Z

datafusion/core/src/datasource/physical_plan/parquet/row_groups.rs

-                &metrics
-            ),
-            vec![1]
+        let mut row_groups = RowGroupSet::new(2);


I updated the tests to use RowGroupSet which I think actually makes them easier to read

alamb · 2024-05-21T19:04:41Z

datafusion/core/src/datasource/physical_plan/parquet/row_groups.rs

@@ -1196,14 +1223,60 @@ mod tests {
            .await
    }

+    // What row groups are expected to be left after pruning


I was finding vec![0] hard to understand what was being tested, so I made this structure to encapsulate it

alamb · 2024-05-22T09:50:00Z

@Ted-Jiang I wonder if you might have time to review this PR at some point?

Ted-Jiang · 2024-05-23T05:07:58Z

@Ted-Jiang I wonder if you might have time to review this PR at some point?

Thanks ❤️ @alamb i will review this carefully later today

Ted-Jiang

LGTM add RowGroupSet really make this pruning code more easy to read 👍 Thanks @alamb

alamb · 2024-05-23T11:16:01Z

Thank you very much for the review @Ted-Jiang 🙏

…API, part 1) (apache#10607) * Refactor parquet row group pruning into a struct * Port tests * improve docs * fix msrv

advancedxy · 2024-05-23T14:31:05Z

datafusion/core/src/datasource/physical_plan/parquet/row_groups.rs

+    /// Prune remaining row groups to only those  within the specified range.
+    ///
+    /// Updates this set to mark row groups that should not be scanned
+    pub fn prune_by_range(&mut self, groups: &[RowGroupMetaData], range: &FileRange) {


Sorry for being late and nit-picking here.
There is a potential for misuse of this API if the number of row groups provided to RowGroupSet does not align with the passed groups. Should we first verify that the number of row groups matches, or do you believe the row group metadata array should be included as a field in RowGroupSet?

Sorry for being late and nit-picking here.

No that is great. Thank you . I appreciate the comments.

should we first verify that the number of row groups matches,

Yes this is a good idea. I will do so in the next PR

or do you believe the row group metadata array should be included as a field in RowGroupSet?

I actually tried this in an earlier version of this code. One challenge I found is that prune_by_bloom_filter needs to have the builder (as it may need to potentially fetch bloom filters) which also has the metadata. So if the RowGroupSet had the metadata then there was an opportunity for inconsistency (between the metadata in the RowGroupSet and the builder)

I do think there might be an opportunity for imporvement here though

Here is a PR to add some asserts: #10641

Yes this is a good idea. I will do so in the next PR

Thanks for the quick update. I agree we can at least add check first.

I actually tried this in an earlier version of this code. One challenge I found is that prune_by_bloom_filter needs to have the builder (as it may need to potentially fetch bloom filters) which also has the metadata. So if the RowGroupSet had the metadata then there was an opportunity for inconsistency (between the metadata in the RowGroupSet and the builder)

I do think there might be an opportunity for improvement here though

Then, I think we can leave it as it is for now. We can revise the impl later if some issues popped up.

…API, part 1) (apache#10607) * Refactor parquet row group pruning into a struct * Port tests * improve docs * fix msrv

github-actions bot added the core Core DataFusion crate label May 21, 2024

alamb changed the title ~~Refactor parquet row group pruning into a struct~~ Refactor parquet row group pruning into a struct (use new statistics API, pt 1) May 21, 2024

alamb added 2 commits May 21, 2024 14:53

Refactor parquet row group pruning into a struct

a9444d8

Port tests

f6f9586

alamb force-pushed the alamb/use_new_api branch from 45d4813 to f6f9586 Compare May 21, 2024 18:53

alamb changed the title ~~Refactor parquet row group pruning into a struct (use new statistics API, pt 1)~~ Refactor parquet row group pruning into a struct (use new statistics API, part 1) May 21, 2024

improve docs

0805226

alamb commented May 21, 2024

View reviewed changes

fix msrv

4160326

alamb marked this pull request as ready for review May 21, 2024 21:11

Ted-Jiang approved these changes May 23, 2024

View reviewed changes

alamb merged commit 656da83 into apache:main May 23, 2024
23 checks passed

alamb added a commit to alamb/datafusion that referenced this pull request May 23, 2024

Refactor parquet row group pruning into a struct (use new statistics …

8458025

…API, part 1) (apache#10607) * Refactor parquet row group pruning into a struct * Port tests * improve docs * fix msrv

advancedxy reviewed May 23, 2024

View reviewed changes

alamb deleted the alamb/use_new_api branch May 23, 2024 16:05

alamb mentioned this pull request May 23, 2024

Minor: add runtime asserts to RowGroup #10641

Merged

alamb mentioned this pull request Jun 5, 2024

Prune Parquet RowGroup in a single call to PruningPredicate::prune, update StatisticsExtractor API #10802

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor parquet row group pruning into a struct (use new statistics API, part 1) #10607

Refactor parquet row group pruning into a struct (use new statistics API, part 1) #10607

alamb commented May 21, 2024 •

edited

Loading

alamb left a comment

alamb May 21, 2024

alamb May 21, 2024

Ted-Jiang May 23, 2024

alamb May 21, 2024

alamb May 21, 2024

alamb commented May 22, 2024

Ted-Jiang commented May 23, 2024

Ted-Jiang left a comment

alamb commented May 23, 2024

advancedxy May 23, 2024

alamb May 23, 2024

alamb May 23, 2024

advancedxy May 24, 2024

Refactor parquet row group pruning into a struct (use new statistics API, part 1) #10607

Refactor parquet row group pruning into a struct (use new statistics API, part 1) #10607

Conversation

alamb commented May 21, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb commented May 22, 2024

Ted-Jiang commented May 23, 2024

Ted-Jiang left a comment

Choose a reason for hiding this comment

alamb commented May 23, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb commented May 21, 2024 •

edited

Loading