Support hive partitioned parquet files (except HybridScan) in cudf-polars - #23966
Open
mroeschke wants to merge 28 commits into
Open
Support hive partitioned parquet files (except HybridScan) in cudf-polars#23966mroeschke wants to merge 28 commits into
mroeschke wants to merge 28 commits into
Conversation
Polars parses hive keys out of file paths and hands cudf-polars the surviving file list plus the partition values for those files. Those values are not stored in the parquet files, so they have to be materialized onto the rows we read. HivePartitions holds the per-path values and knows how to turn them into columns, either by broadcasting when every path shares a value or by gathering with the reader's source index column. It is frozen and hashable so it can take part in Scan's IR identity.
Thread HivePartitions through Scan so the parquet reader can produce the hive columns polars expects. The projection polars gives us may name hive columns, so those are removed before asking the reader for file columns and added back afterwards from the partition values. Rows are matched to their partition values with the reader's prepended source index column, which is synthesized before filtering and so stays correct when a predicate is pushed down. When every path shares the same values the index is not read at all and the values are broadcast instead. Predicates over hive columns are held back from the parquet filter, since those columns do not exist in the files, and applied post-read once the columns have been materialized. Nothing constructs a hive-partitioned Scan yet; translation still raises.
Selecting only hive columns leaves the parquet reader with an empty projection, and it then returns no rows at all, so neither the row count nor a source index can come from the read. Take the per-path row counts from the file metadata instead and repeat each path's partition values that many times. This is only sound because polars pushes a predicate down alongside the file columns it references, so an empty projection always means no filter was pushed and the metadata counts are exact.
Translate polars' hive_parts frame into HivePartitions instead of refusing the scan. Polars has already pruned the file list against any predicate it could evaluate on the partition keys by this point, so what arrives is the surviving paths and the partition values for them. Tests cover the shapes polars produces: hive columns projected alone or mixed with file columns, predicates over hive columns, over file columns and over both, pruning down to a single file or to none, a slice or row index pushed into a hive-only projection, dtype overrides via hive_schema, null partition values, and a file column shadowed by a hive key. These run on the in-memory executor; the streaming executors still drop hive columns and are enabled next.
SplitScan and FusedScan each read a subset of the base scan's files, so they need the partition values for just those files. Slice them when the StreamingScan is built: a split always covers one file, so its values are uniform and get broadcast, while a fused group needs the source index. The hybrid reader builds its own projection and filter and has no notion of columns that are absent from the file, so hive scans fall back to the regular reader for now. The hive scan tests now run against the full engine matrix rather than just the in-memory executor.
libcudf clears its per-source row counts once the reader applies a
filter, so repeating each path by its row count produced a length
mismatch and the read failed outright:
CUDF failure at repeat.cu:101: in and count must have equal size
Gather the paths with the source index instead, which the reader
synthesizes before filtering. The row counts are still used when nothing
is filtered, to avoid reading a column that is not needed.
Add the cases the earlier commits left untested: the chunked reader,
partition keys of each dtype polars can parse out of a path, and several
files under one partition so the values are uniform across more than one
path.
Four delta tests were xfailed as needing hive partitioning. They still
fail with hive scans enabled, and identically so on main:
CUDF failure: All non-empty sources must have the same number of
columns
These tables evolve their schema between versions, so the files have
differing column counts, which the reader cannot handle. Correct the
reasons rather than removing the entries.
Polars' own tests/unit/io/test_hive.py goes from 7 to 33 passing, and
the rest of tests/unit/io fails identically to main.
Scan is only ever constructed in translate.py, where the paths and the partition values come from the same polars node, and reconstruct() and __reduce__() carry both through together. Nothing could trip this without editing that one call site, and no test reached the line. Unlike the neighbouring cached parquet info check, which guards state attached by a later pass and does have a test behind it.
Give _parquet_rows_per_path and _pop_source_index the Parameters and Returns sections the rest of the module uses.
Once the budget reaches zero every later path takes zero rows, so the loop can leave the pre-filled zeros in place and stop. This also covers the cases that start with no budget at all, where skip_rows swallows the whole dataset or n_rows is zero.
Keep what polars hands over rather than unpacking it into parallel name, dtype and value tuples, and derive the properties from it on demand. Conversion to device goes through Table.from_arrow, which also lets broadcast reuse filling.repeat instead of building a scalar per key, so nulls and temporal types need no special handling. Identity needs care. The default dataclass __eq__ would compare two frames with ==, which yields a frame rather than a bool, and DataFrame.equals ignores the schema, so Int32 and Int64 partition values would compare equal. Worse, node digests fall back to repr() and polars elides the middle of a tall frame's repr, so two scans differing only in a middle partition value would digest identically. All three dunders are therefore driven off a cached key holding the schema and every row.
These three tests translate a query to inspect the IR, and were passing a locally built GPUEngine rather than the shared fixture. Take the engine fixture instead so they see each real backend, and skip the in-memory variant, which never builds a SplitScan or FusedScan.
The row counts were asserted by calling the private helper directly. Slicing a hive-only projection reaches it through the public API with the same skip_rows and n_rows combinations, and checks the rows that come back rather than just the counts, so the whole helper stays covered including the paths where the budget starts empty.
__repr__ still unpacked the removed _key, so every repr() raised AttributeError. __hash__ hashed the result of list.append, which is None, so every instance hashed alike; hash the schema alongside a tuple of the row hashes instead. Since __repr__ now defers to the frame, which elides the middle of a tall one, the stable node digest no longer separates partition values that differ only in an elided row. That digest feeds explain output and actor ids rather than result reuse, and is truncated to 32 bits anyway, while CSE goes through __hash__ and __eq__, which do see every row. The test that covered the digest now covers those instead, joined by cases for column names and path order, which hash_rows alone does not distinguish.
mroeschke
added
improvement
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
github-actions
Bot
added
Python
Rename cudf_polars/dsl/utils/hive.py to per_path.py and HivePartitions to PerPathValues, then use it in Scan.add_file_paths so the gather/repeat expansion of per-path values lives in one place.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 SummarySummary by CodeRabbit
WalkthroughHive-partitioned Parquet scans now preserve per-path partition metadata, defer filters on unreadable Hive columns, reconstruct partition and file-path columns, and support split and fused streaming scans. Tests cover filtering, projections, slicing, partition types, nulls, and fallback behavior. ChangesHive-partitioned Parquet scans
Estimated code review effort: 4 (Complex) | ~60 minutes Merge Risk: ⚪ Minimal · up to Delta tests now classify mismatched source-schema cases as unsupported schema evolution rather than Hive partitioning. No current merge-blocking behavior or production risk is identified. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
vyasr
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Towards #17832
This PR enables reading hive partitioned parquet files in
pl.scan_parquet. This PR doesn't intentionally omits support for the hybrid scan reader as a follow up.Checklist