Prefetch parquet metadata for scan tasks - #22700
Conversation
|
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
Add a parquet metadata prefetch cache in IR execution context and plumb it through in-memory and streaming scan paths to reuse footer metadata across related reads.
392b7f0 to
80d92b7
Compare
TomAugspurger
added
improvement
|
#22739 implements a change to (py)libcudf to accept a file size in addition to a path in I suspect that we can use that if we explicitly list all the file sizes (using something like Finally, #22734 would let us skip that explicit list file sizes call and just use the file sizes directly from polars. This will require changes to polars, so is not a blocker for this PR. |
|
109ad43 demonstrates using the new
Note that the But for now, I've reverted that commit so this isn't blocked by #22739. |
This reverts commit 109ad43.
This refactors how we dynamically generate `Scan` nodes with the rapidsmpf streaming runtime. Previously, these nodes were generated inside `cudf_polars.streaming.actor_grpah.io.scan_node`. The lowered `ir` node would have a basic `Scan` (or maybe a `SplitScan` node) that was never actually executed. Instead, *after* we created the Actor responsible for that node, we'd generate new `Scan` nodes to do the work of reading some subset of the files / row groups. This PR does the same thing (generate these `Scan`/`SplitScan` nodes) but it does it at lowering time, rather than inside the `Scan` task. The actual logic of generating these new nodes is unchanged. We should end up with exactly the same execution as before; we just make the `list[Scan | SplitScan]` earlier. This is primarily motivated by awkwardness in using the dynamically generated Scan nodes in a couple spots: 1. prefetching parquet metadata(#22700) 2. Quent tracing I've added a new `StreamingScan` IR node that just holds references to the `list[Scan]` nodes. This might be overkill for what we need, but it's at least consistent with how we do the main lowering, which I like. Authors: - Tom Augspurger (https://github.com/TomAugspurger) Approvers: - Matthew Murray (https://github.com/Matt711) URL: #22758
|
/ok to test f3fa17f |
We treat this as an optimization, so exclude prefetched metadata from hashing & equality of Scan-type nodes. This makes it safe to mutate the Scan nodes after theyre' created by lowering, which simplifies updating the IR graph after metadata have been prefetched.
Matt711
left a comment
Matt711
left a comment
There was a problem hiding this comment.
Thanks @TomAugspurger ! Pretty happy where this PR is at.
These are the next steps I'm aware of
- Separate Scan class by file type #23063
- Prefetch inside the streaming network (maybe do). Can you open an issue?
- Turn on prefetching by default (this is a maybe I think) #22826
- Use prefetched footers for hybrid scan in #22857
- Reuse prefetched parquet metadata for decimal predicate pushdown #22940
Anything missing?
| def _parquet_physical_types( | ||
| paths: list[str], columns: list[str] | None | ||
| ) -> dict[str, plc.DataType]: | ||
| # This may not be able use prefetched metadata, since we don't (currently) have a Schema. |
There was a problem hiding this comment.
nitpick: Can you make this a TODO and link the issue?
|
/merge |
| with kvikio.RemoteFile.open(path) as remote_file: | ||
| sizes.append(remote_file.nbytes()) | ||
| else: | ||
| sizes.append(None) |
There was a problem hiding this comment.
Here you make a size value of None.
| CachedParquetInfo(path, size, file_metadata) | ||
| for path, size, file_metadata in zip(paths, sizes, metadata, strict=True) |
There was a problem hiding this comment.
But here you are claiming that size is an int. What gives?
There was a problem hiding this comment.
Whoops, I think Rick might have flagged this earlier too but I misread it.
It's fine for us to pass an int | None to FilepathSource, but CachedParquetInfo currently states that .size is an int.
I'll get this fixed.
In rapidsai#22700 we needed a way to get the equivalent of `ParquetMetadata.columnchunk_metadata` from a list of parquet footers. See the discussion at https://github.com/rapidsai/cudf/pull/22700/files/1d9e2147f8001a6e607da6f79c8125be5d1bb3da#diff-e5ac2ccddbef69c522213586ce23621cc87f3c21f123291c17a0d4dda628def8 for details. This implements that in libcudf / pylibcudf, and we'll make use of it in rapidsai#22700.
In rapidsai#22700 we needed a way to get the equivalent of `ParquetMetadata.columnchunk_metadata` from a list of parquet footers. See the discussion at https://github.com/rapidsai/cudf/pull/22700/files/1d9e2147f8001a6e607da6f79c8125be5d1bb3da#diff-e5ac2ccddbef69c522213586ce23621cc87f3c21f123291c17a0d4dda628def8 for details. This implements that in libcudf / pylibcudf, and we'll make use of it in rapidsai#22700.
5 participants
Description
To improve the performance of cudf-polars' remote IO, we'd like to prefetch all parquet metadata ahead of time, rather than inside a
Scantask. This ensures we read metadata at most once perpaths, and allows us to cleanly separate setup things like metadata reads from the actual job of reading data. This is also a necessary step for adopting the Hybrid Scan reader.To accomplish this, we have a new stage of execution that prefetches parquet metadata for a given IR query plan. This occurs immediately after lowering. The
FileMetadatafor those paths is read and cached on a new field inIRExecutionContext. Once we get to the actualScan.do_evaluatethat's reading the parquet metadata, we access the pre-fetched parquet file metadata from theIRExecutionContextfor ourpaths.I've implemented this as a new option on
ParquetOptions, mostly for ease of benchmarking. I suspect we'll want to enable it by default but I want to confirm that.Closes #22667
Closes #22753