[BUG] Fix reading partition key columns in DeltaLake

[jaychia]

Fixes pushdowns for column selection on partition keys in DeltaLake and Iceberg.

When table formats such as Iceberg and Delta Lake store the data for a partition column, they will strip the column from the actual Parquet data files that they write out. Then when we want to select only specific columns, our Parquet reader fails when it is not able to find those columns in the file.

NOTE: Seems like Iceberg only does this for identity transformed partition columns

Follow-on issue for selection of only the partition keys: #2129

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 85.15%. Comparing base (0fe3df7) to head (f05a535).

❗ Current head f05a535 differs from pull request most recent head 4b1ef7d. Consider uploading reports for the commit 4b1ef7d to get more accurate results

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2118   +/-   ##
=======================================
  Coverage   85.15%   85.15%           
=======================================
  Files          68       68           
  Lines        7356     7356           
=======================================
  Hits         6264     6264           
  Misses       1092     1092           

[samster25]

you likely need to apply this logic to the eager code path. Say you passed in a predicate.

[jaychia]

This should be fixed in this section: https://github.com/Eventual-Inc/Daft/pull/2118/files#diff-1d03b58906d515cc349a8a414595727e787d60af7adb3ab9fb3db3558c9e7e27R789-R808

Unfortunately our code is a little jank right now, with 2 main codepaths:

1. MicroPartition::from_scan_task either eagerly produces loaded micropartitions or lazily produces unloaded micropartitions using read_parquet_into_micropartition
2. Lazy micropartitions then load data using materialize_scan_task

So any logic often has to be duplicated across those two functions (read_parquet_into_micropartition and materialize_scan_task). Perhaps there is some opportunity to consolidate this by having read_parquet_into_micropartition call materialize_scan_task under the hood, but that might be a refactor for another PR?

[samster25]

I think a simpler way would be to just wrap read_parquet_bulk in this crate

[jaychia]

Could you elaborate? Both read_parquet_into_micropartition and materialize_scan_task do currently wrap read_parquet_bulk.

[clarkzinzow]

I think that @samster25 was talking about creating a read_parquet_bulk_with_virtual (or similarly named) utility in this micropartition.rs module that does all of the virtual column + pushdown handling before calling read_parquet_bulk, shared (called) by both the eager and lazy paths, so we don't have to duplicate this handling logic.

However, that would require duplicating the read_parquet_bulk API, which is pretty large, and we also wouldn't have coverage for other file types (e.g. if we added Avro or ORC support to Iceberg reads). @samster25 am I misunderstanding your suggestion?

[clarkzinzow]

We do already have a shared _read_parquet_into_loaded_micropartition utility for read_parquet_into_micropartition, maybe we should refactor materialize_scan_task to call out to similar helpers for each file format.

[samster25]

I think that @samster25 was talking about creating a read_parquet_bulk_with_virtual (or similarly named) utility in this micropartition.rs module that does all of the virtual column + pushdown handling before calling read_parquet_bulk, shared (called) by both the eager and lazy paths, so we don't have to duplicate this handling logic.

Yeah exactly!

[samster25]

add 2 tests:

• select only the partition column with filter on it
• select only the partition column with filter on another column

[jaychia]

Added the two requested tests:

1. This fails, because reading just the partition column only still fails after this PR (we dont have a great way of "reading [] columns from a file" at the moment
2. This passes