feat(rust): Prune row groups before loading all columns

[bchalk101]

This addresses #13608.

This addresses the situation when the row group statistics don't help with filtering out row groups. By loading just the columns required to apply the predicate and filtering out non-required row groups, followed by loading all the data from the leftover row groups.

The main downside of this implementation is that in the bad case, where the number of rows for the predicate and projection are equal the data is being downloaded twice. To get around this, perhaps a feature flag can be added, something like ROW_GROUP_PRUNING, and only if it is turned on will this filtering be applied.

I will note, that on the datasets I am currently working on, as specified in the issue, the filtering went from 25min and 32 GB memory consumption to 25 seconds and negligible memory.

• Is it safe? Handles all situations when there is no predicate
• Feature flagged
• Tested

[bchalk101]

Hey @ritchie46,
Would be great to get your take on this enhancement.
Is this something you will be ok merging in? Is this being implemented in the right direction?

Just want to check before putting more effort into this and implementing the feature toggle and testing.

[ritchie46]

Hi @bchalk101 why was this closed? Didn't it work?

[bchalk101]

Hey @ritchie46,
No, it does work. I just wasn't sure if it would get merged in and wanted to test out some other optimizations, so I closed it. But happy to re-open it again.

[mkleinbort]

Hi. Very interested in this feature - it'd be amazing for some very wide tables.

[ritchie46]

Yes, it is interesting, but I want to think about this a bit more. As the most common case would slow down. We loose an embarrassingly parallel load to a sequential one with a very low probability of being faster.

[bchalk101]

I want to give some further insight into our use case, which perhaps can help with the decision. We are using Parquets as the format for saving data for training ML models, this means that each row can be quite large, even if some columns are small. For example, large compressed numpy arrays or even jpeg images in a column. There is still small metadata in each row, which is what we use with the filters before actually selecting the rows. While I think we would be a small percentage of users, there is definitely a large number of people using Parquets for ML training data.

I would be open to other ways to use this with Polars, potentially a secondary read library suited to such data. The issue is that it makes it very difficult to use Polars with such data without a bunch of changes. Other changes may include, applying limits (is .limit) on the pruned row groups before collecting more data (which again will slow down the general use case), implementing a generator for loading data or allowing setting parquet size when sinking data into files.

The other option, as I mentioned in the description is to use a feature flag, but of course, this can lead to "flag bloat".

[codspeed-hq]

CodSpeed Performance Report

Merging #13746 will not alter performance

_{Comparing bchalk101:optimize_rg_read (c77cf72) with main (11fe9d8)}

Summary

✅ 34 untouched benchmarks

[kszlim]

Is this the same optimization as Late Materialization?

[bchalk101]

Not exactly - it's the same idea, but done at the IO level.
The process is as follows:

1. Download predicate-only columns
2. Apply predicate to decide if the rest of the required columns should be downloaded
3. Mark the Row Group as required and use it to apply the slice.
4. Download the entire Row Group and apply the predicate (At this point late materialisation would help)

This specifically helps when I/O and memory are the blockers, ie wide tables with columns that contain heavy data.

[ritchie46]

I will close this one as this isn't getting merged anymore. I will see us dynamically adapting to such a strategy in the new streaming engine, but for now it's out of scope.