Refactor dynamic Scan node lowering - #22758
Conversation
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`. Now, we treat it more like a (rank-specific) lowering stage. This is primarily motivated by awkwardness in using the dynamically generated Scan nodes in a couple spots: 1. prefetching parquet metadata(rapidsai#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.
|
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
|
/ok to test e1a4b4c |
TomAugspurger
added
improvement
This refactors the refactor to just lower to this StreamingScan inside the the initial lowering.
| return new_ir, {new_ir: PartitionInfo(count=count, io_plan=plan)} | ||
|
|
||
|
|
||
| class StreamingScan(IR): |
There was a problem hiding this comment.
Now that we're doing this earlier, I worry that the lack of a StreamingScan.do_evalaute is going to cause issues. Like, how would the in-memory executor handle this?
There was a problem hiding this comment.
Maybe this PR should just ago a bit further? Maybe we don't really want StreamingScan to carry around a list[Scan | SplitScan]. Rather, we just collect list[ScanChunkInfo], where:
class ScanChunkInfo:
paths: list[str]
"""List of paths to read from"""
chunk_index: int
"""Index of this chunk (within `paths`)"""
num_chunks: int
"""Total number of chunks being generated from `paths`"""This way, we can add a StreamingScan.do_evaluate that basically looks like SplitScan.do_evaluate (but is actually simpler when len(paths) > 1, because num_chunks is always 1 in that case).
There was a problem hiding this comment.
Taking a look, but can you say a bit more about why avoiding Scan / SplitScan instances might be desirable here?
I'm looking at
cudf/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
Lines 244 to 254 in 0c0f0ea
read_chunk on the Scan / SplitScan that were previously dynamically generated.
There was a problem hiding this comment.
can you say a bit more about why avoiding Scan / SplitScan instances might be desirable here?
I may be wrong, but it seems like you are running into issues with the misdirection between the Scan node in the plan and the actual Scan/SplitScan IR objects that are used in the actor. I'm just pointing out that we don't actually need to create a bunch of independent Scan/SplitScan objects if that's an issue.
There was a problem hiding this comment.
Gotcha. I think it's less about whether or not the Scan nodes are created, and more about when they're created.
In #22700, we want to know all of the parquet paths we're going to read (on each worker) before we start reading any data. I think as long as we create these nodes ahead of time then we're OK.
There was a problem hiding this comment.
Okay, yeah. That's all completely up to us.
The only explicit design choice we probably want to commit to is the idea that we don't need to parse all Parquet metadata up-front to make the general partitioning plan.
Some unnecessary background (feel free to ignore if irrelevant):
This was a mistake that we made in Dask-DataFrame read_parquet API several years ago: We needed to specify specific path and row-group indeices for every read task in order to build the task graph. That ended up being terrible for cloud IO, because we would need to parse all footer metadata up front. The cudf-polars design is different. We either expect each chunk to be a range of complete files, or we expect each chunk to be a slice of a single file. For the latter case, we don't need to know exactly which row-groups a specific chunk will be mapped to ahead of time. We just need to know how many chunks will be mapped from that file. The specific row-group range can be determined by the actor at runtime when the footer metadata for that file is already available.
There was a problem hiding this comment.
Yeah, I think that aligns with my goals. With this and #22700, we should read each file metadata at most once per worker, and we should only read the file metadata of files we actually are responsible for reading.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR threads rank/nranks through IR lowering, adds expand_scan_for_rank and a StreamingScan IR to produce rank-local scan lists, updates the actor graph to consume StreamingScan (including native-parquet gating and simplified scan_node), extends serialization and fast-count paths, and adds unit and integration tests. ChangesRank-aware Streaming IO
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
| return new_ir, {new_ir: PartitionInfo(count=count, io_plan=plan)} | ||
|
|
||
|
|
||
| class StreamingScan(IR): |
There was a problem hiding this comment.
Maybe this PR should just ago a bit further? Maybe we don't really want StreamingScan to carry around a list[Scan | SplitScan]. Rather, we just collect list[ScanChunkInfo], where:
class ScanChunkInfo:
paths: list[str]
"""List of paths to read from"""
chunk_index: int
"""Index of this chunk (within `paths`)"""
num_chunks: int
"""Total number of chunks being generated from `paths`"""This way, we can add a StreamingScan.do_evaluate that basically looks like SplitScan.do_evaluate (but is actually simpler when len(paths) > 1, because num_chunks is always 1 in that case).
I'm looking into this. It's certainly correct that the |
|
b8ca40e attempts to fix that "chunked=False fallback" thing, without really understanding the issue. I'm just trying to replicate the old behavior here. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/cudf_polars/cudf_polars/streaming/io.py (1)
477-496:⚠️ Potential issue | 🟠 Major | ⚡ Quick winFix
StreamingScan_non_child/__init__argument mismatch to avoid pickling/reconstructTypeError
Node.__reduce__/reconstructbuild constructor args fromStreamingScan._non_childin-order (scans,base_scan,schema), butStreamingScan.__init__currently only accepts(scans, base_scan), so reconstruction (e.g. Dask/Ray pickling) would call the constructor with 3 positional args and raiseTypeError. Align_non_childwith the constructor signature (and update thenew_ir = StreamingScan(scans, ir)call-site accordingly).🐛 Proposed fix — make
schemaa constructor arg_non_child = ( + "schema", "scans", "base_scan", - "schema", ) - _n_non_child_args = 2 + _n_non_child_args = 3 scans: list[Scan | SplitScan] base_scan: Scan - def __init__(self, scans: list[Scan | SplitScan], base_scan: Scan): + def __init__(self, schema: Schema, scans: list[Scan | SplitScan], base_scan: Scan): self.scans = scans self.base_scan = base_scan - self.schema = base_scan.schema - self._non_child_args = (scans, base_scan) + self.schema = schema + self._non_child_args = (schema, scans, base_scan) self.children = ()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_polars/cudf_polars/streaming/io.py` around lines 477 - 496, StreamingScan's _non_child lists ("scans", "base_scan", "schema") but __init__ only accepts (scans, base_scan), causing reconstruct/pickling to pass three positional args and raise TypeError; change StreamingScan.__init__ to accept a third parameter (schema) and set self.schema from that param, update _n_non_child_args to 3, keep _non_child in the declared order, and update call-sites such as the new_ir = StreamingScan(scans, ir) invocation to pass the schema (e.g., StreamingScan(scans, ir, ir.schema)) so construction and Node.__reduce__/reconstruct ordering match.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@python/cudf_polars/cudf_polars/streaming/io.py`:
- Around line 477-496: StreamingScan's _non_child lists ("scans", "base_scan",
"schema") but __init__ only accepts (scans, base_scan), causing
reconstruct/pickling to pass three positional args and raise TypeError; change
StreamingScan.__init__ to accept a third parameter (schema) and set self.schema
from that param, update _n_non_child_args to 3, keep _non_child in the declared
order, and update call-sites such as the new_ir = StreamingScan(scans, ir)
invocation to pass the schema (e.g., StreamingScan(scans, ir, ir.schema)) so
construction and Node.__reduce__/reconstruct ordering match.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b0416e34-b2a6-4d24-9031-7900f6dbbe0c
📒 Files selected for processing (2)
python/cudf_polars/cudf_polars/streaming/actor_graph/io.pypython/cudf_polars/cudf_polars/streaming/io.py
Matt711
left a comment
Matt711
left a comment
There was a problem hiding this comment.
This ended up being very clean, thanks!
|
/merge |
|
Some failures in the polars tests: At least some of these are consistent with the node just being out of GPU memory unexpectedly. |
|
Maybe just rerun since the other job passed? |
5 participants
Description
This refactors how we dynamically generate
Scannodes with the rapidsmpf streaming runtime.Previously, these nodes were generated inside
cudf_polars.streaming.actor_grpah.io.scan_node. The loweredirnode would have a basicScan(or maybe aSplitScannode) that was never actually executed. Instead, after we created the Actor responsible for that node, we'd generate newScannodes to do the work of reading some subset of the files / row groups.This PR does the same thing (generate these
Scan/SplitScannodes) but it does it at lowering time, rather than inside theScantask. The actual logic of generating these new nodes is unchanged. We should end up with exactly the same execution as before; we just make thelist[Scan | SplitScan]earlier.This is primarily motivated by awkwardness in using the dynamically generated Scan nodes in a couple spots:
I've added a new
StreamingScanIR node that just holds references to thelist[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.