Fetch parquet metadata in rapidsmpf streaming network - #23145
Fetch parquet metadata in rapidsmpf streaming network#23145TomAugspurger wants to merge 23 commits into
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
TomAugspurger
added
improvement
a4215ea to
f89d255
Compare
This moves the parquet metadata prefetching that (optionally) happens between IR lowering and streaming network execution into the streaming network. This ensures that any given Scan node can begin executing as soon as *its* metadata is ready, rather than *all* the metadata in a query. At a high level, this PR adds new Actors to the streaming network that are responsible solely for fetching parquet metadata for some paths. The actors send that metadata to the Actors running the Scan nodes. One wrinkle is that we want to preserve the "no duplicate reads of parquet metadata" feature of cudf-polars' prefetching. If two Scan nodes refer to the same paths, we want a single actor that reads the metadata and sends it to both actors doing the Scans. PDS-H Query 7 exhibits this behavior.
73f00db to
98b2ed5
Compare
|
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 moves parquet footer metadata prefetching into streaming scan actor nodes, threads prefetched metadata through scan execution, and updates tests for shared-path, disjoint-path, and handler validation cases. ChangesParquet metadata prefetch relocation to streaming actor graph
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
♻️ Duplicate comments (1)
python/cudf_polars/cudf_polars/streaming/actor_graph/io.py (1)
704-706: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick winHIGH: Fail closed when prefetched metadata is missing.
Issue: these
.get(...)lookups silently passNoneif a prefetch-enabled scan is missing its metadata key. Why: that masks actor-graph wiring bugs and reintroduces duplicate footer reads instead of preserving the no-duplicate-reads contract.This is the same failure mode as the prior
.getmetadata-channel review.Proposed fix
- prefetched_parquet_metadata.get( - tuple(cast("FusedScan | SplitScan", scan).paths) - ), + ( + prefetched_parquet_metadata[ + tuple(cast("FusedScan | SplitScan", scan).paths) + ] + if metadata_channels_by_key + else None + ), @@ - prefetched_parquet_metadata.get(tuple(scan.paths)), + ( + prefetched_parquet_metadata[tuple(scan.paths)] + if metadata_channels_by_key + else None + ),Also applies to: 734-734
🤖 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/actor_graph/io.py` around lines 704 - 706, The prefetched metadata lookup in the actor-graph I/O path is failing open by using a .get() that can return None for prefetch-enabled scans, which can hide wiring bugs and trigger duplicate footer reads. Update the logic in the relevant scan-handling code in io.py (around the prefetched_parquet_metadata access, including the other matching .get site) to require the metadata entry to exist and raise immediately if it is missing, so the no-duplicate-reads contract is enforced.
🧹 Nitpick comments (2)
python/cudf_polars/tests/streaming/test_scan.py (2)
134-146: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winTest doesn't verify the actual dedup/single-actor behavior.
This test only checks result correctness via
assert_gpu_result_equal, but the PR's stated goal is that a single metadata-fetching actor serves bothScannodes sharing the same paths. As written, the test would pass even if metadata were fetched twice per path (i.e., no dedup at all), since correctness doesn't depend on that.Consider asserting something more direct about the actor graph (e.g., number of
parquet_metadata_prefetch_nodeactors created equals 1 for the shared-path scenario), if such introspection is feasible.🤖 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/tests/streaming/test_scan.py` around lines 134 - 146, The current test in test_scan_parquet_prefetch_metadata_shared_scan_paths only validates result equality and does not prove shared metadata prefetching. Update it to assert the dedup/single-actor behavior directly by inspecting the streaming engine or actor graph after building the query, and verify that only one parquet_metadata_prefetch_node is created for the shared-path Scan nodes. Use the existing test_scan_parquet_prefetch_metadata_shared_scan_paths and assert_gpu_result_equal setup as the entry point, but add a structural assertion instead of relying only on output correctness.
720-742: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMissing success-path coverage for
recv_prefetched_parquet_metadata_handler.Only the two error branches are tested. Since
recv_prefetched_parquet_metadata_handlerinio.pymarks both raise branches with# pragma: no cover; unreachable, but this test now explicitly exercises them, those pragma comments are stale and should be removed inio.pyto keep coverage reporting accurate.Additionally, consider adding a case where
group_keymatches, asserting the handler returnspayload.cached_parquet_infounmodified.🤖 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/tests/streaming/test_scan.py` around lines 720 - 742, The `recv_prefetched_parquet_metadata_handler` coverage is incomplete because `test_recv_prefetched_parquet_metadata_handler_errors` only exercises the error branches. Update the existing test in `test_scan.py` to add a success-path case where the `Message.group_key` matches the scan input and assert the handler returns `payload.cached_parquet_info` unchanged. Then remove the now-stale `# pragma: no cover; unreachable` comments from the corresponding raise branches in `recv_prefetched_parquet_metadata_handler` in `io.py` so coverage reporting reflects the exercised code.
🤖 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.
Duplicate comments:
In `@python/cudf_polars/cudf_polars/streaming/actor_graph/io.py`:
- Around line 704-706: The prefetched metadata lookup in the actor-graph I/O
path is failing open by using a .get() that can return None for prefetch-enabled
scans, which can hide wiring bugs and trigger duplicate footer reads. Update the
logic in the relevant scan-handling code in io.py (around the
prefetched_parquet_metadata access, including the other matching .get site) to
require the metadata entry to exist and raise immediately if it is missing, so
the no-duplicate-reads contract is enforced.
---
Nitpick comments:
In `@python/cudf_polars/tests/streaming/test_scan.py`:
- Around line 134-146: The current test in
test_scan_parquet_prefetch_metadata_shared_scan_paths only validates result
equality and does not prove shared metadata prefetching. Update it to assert the
dedup/single-actor behavior directly by inspecting the streaming engine or actor
graph after building the query, and verify that only one
parquet_metadata_prefetch_node is created for the shared-path Scan nodes. Use
the existing test_scan_parquet_prefetch_metadata_shared_scan_paths and
assert_gpu_result_equal setup as the entry point, but add a structural assertion
instead of relying only on output correctness.
- Around line 720-742: The `recv_prefetched_parquet_metadata_handler` coverage
is incomplete because `test_recv_prefetched_parquet_metadata_handler_errors`
only exercises the error branches. Update the existing test in `test_scan.py` to
add a success-path case where the `Message.group_key` matches the scan input and
assert the handler returns `payload.cached_parquet_info` unchanged. Then remove
the now-stale `# pragma: no cover; unreachable` comments from the corresponding
raise branches in `recv_prefetched_parquet_metadata_handler` in `io.py` so
coverage reporting reflects the exercised code.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 185220c1-75b6-4a04-9e96-8857f651158a
📒 Files selected for processing (6)
python/cudf_polars/cudf_polars/dsl/utils/io.pypython/cudf_polars/cudf_polars/engine/core.pypython/cudf_polars/cudf_polars/streaming/actor_graph/core.pypython/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.pypython/cudf_polars/cudf_polars/streaming/actor_graph/io.pypython/cudf_polars/tests/streaming/test_scan.py
💤 Files with no reviewable changes (1)
- python/cudf_polars/cudf_polars/engine/core.py
| ir_context=ir_context, | ||
| ) as tracer: | ||
| for key, ch_metadata in metadata_channels_by_key.items(): | ||
| msg = await ch_metadata.recv(context) |
There was a problem hiding this comment.
So we wait for all metadata to arrive for the chunks being processed by this actor before we start reading data?
There was a problem hiding this comment.
I believe that's correct.
As you're implying, we can probably structure this so that each chunk's read can begin as soon as its metadata is ready. We just need to ensure that the ordering is consistent. I'll look into that now.
There was a problem hiding this comment.
890094a has that.
One subtle thing to call out: the prefetching channels are keyed by scan.paths: we want to read the parquet metadata just once per set of files. But for SplitScan nodes, we might have many individual scans reusing the same parquet file metadata, and these might be running concurrently. So we have a little helper cached_parquet_info_for_scan for receiving the prefetched metadata that handles
- doing the actual recv and parsing a
Messageinto alist[CachedParquetInfo] - caching, for reuse by other
scannodes sharing the same path. This cache needs to work for both cases where the lookups from the cache are sequential (easy) or concurrent (which is handled by making anasyncio.Taskthat wraps the awaitable receiving the parquet metadata, so that multiple producers can await the same ultimate data.
There was a problem hiding this comment.
I don't think I understand why we need to have so many distinct actors and channels. I would think that we only need a single parquet_metadata_actor for each scan_actor. The parquet_metadata_actor would construct an ordered list of unique list[tuple[str, ...]] groups needed by the down-stream scan_actor, and initiate the metadata reads for all these groups at once. As these metadata groups become available, we just push them into a channel that is connected to the scan_node as metadata (so we have an unbound queue of metadata messages that the scan_node can pull from).
There was a problem hiding this comment.
I suppose we can pass the metadata as "data", but I don't think we need the single-item-capacity back pressure.
There was a problem hiding this comment.
f8001e8 has another pass at this. This implements the "1 Prefetch actor per Scan actor" thing we talked about earlier.
I have, at least temporarily, given up on the requirement that we not reread metadata when two nodes request the same metadata. I think it was creating too many complexities downstream. I'll see whether than can be recovered, but it might not be worth the hassle.
…fetch-metadata-actors-2
…fetch-metadata-actors-2
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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.
Inline comments:
In `@python/cudf_polars/cudf_polars/streaming/actor_graph/io.py`:
- Around line 686-703: The metadata receiver task started by
_start_metadata_receiver in scan_node is never cleaned up, so wrap the
shutdown_on_error block in a try/finally and explicitly cancel and await
_metadata_receiver_task on exit. Use the existing _metadata_receiver_task symbol
alongside shutdown_channels and cached_parquet_info_for_task so the task is
always finalized even on early return or failure.
- Around line 616-623: The _receive helper in the metadata prefetch flow should
not leave pending entries in metadata_futures unresolved if it exits early;
update the logic in _receive and the create_task path so any exception or
premature stop is propagated to all unfinished futures. Use the existing symbols
recv_prefetched_parquet_metadata, recv_prefetched_parquet_metadata_handler, and
futures to locate the fix, and ensure remaining futures are completed with the
same failure so downstream awaiters fail fast instead of hanging.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b70920bf-2a71-4034-94ab-653419f5c5d7
📒 Files selected for processing (4)
python/cudf_polars/cudf_polars/streaming/actor_graph/core.pypython/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.pypython/cudf_polars/cudf_polars/streaming/actor_graph/io.pypython/cudf_polars/tests/streaming/test_scan.py
| return tuple(metadata_scans) | ||
|
|
||
|
|
||
| def _build_metadata_channels_by_scan( |
There was a problem hiding this comment.
This can probably be inlined now.
| await ch_out.drain(context) | ||
|
|
||
|
|
||
| def _start_metadata_receiver( |
There was a problem hiding this comment.
This is kind of annoying / complicated, but I think we need some mechanism to ensure that the right parquet metadata gets to the right SplitScan | FusedScan task when we have multiple producers running concurrently.
I think we do know for sure that metadata prefetcher Actor iterates over .scans and inserts a Message per scan in the original order from the StreamingScan. But I think when we're running multiple scan reader tasks concurrently, they might get to the await reading a message with parquet metadata out of order.
TomAugspurger
changed the title
There was a problem hiding this comment.
♻️ Duplicate comments (2)
python/cudf_polars/cudf_polars/streaming/actor_graph/io.py (2)
708-740: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winVerify
_metadata_receiver_tasklifecycle cleanup.A prior review pass on this same feature flagged that the receiver task started here was never cancelled/awaited on early exit or failure of
scan_node, causing a leaked task whose exception is never retrieved. The AI summary confirmsscan_nodenow "includes the metadata channel in shutdown handling," but it's unclear whether the task itself (as opposed to just the channel) is cancelled in afinallyblock. As per coding guidelines, "ensure asyncio.Tasks are explicitly canceled in a finally block upon failure."#!/bin/bash sed -n '690,800p' python/cudf_polars/cudf_polars/streaming/actor_graph/io.py rg -n '_metadata_receiver_task|cancel\(\)' python/cudf_polars/cudf_polars/streaming/actor_graph/io.py🤖 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/actor_graph/io.py` around lines 708 - 740, The metadata receiver task started by _start_metadata_receiver in scan_node is still not explicitly cleaned up if scan_node exits early or fails. Add a finally block around the shutdown_on_error scope that cancels _metadata_receiver_task when it exists and awaits it to completion, handling cancellation cleanly so no exception is left unretrieved. Keep the fix local to the scan_node path that sets up metadata_futures and _metadata_receiver_task.Source: Path instructions
647-651: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winReceiver failures leave downstream futures unresolved forever.
If
recv_prefetched_parquet_metadataorrecv_prefetched_parquet_metadata_handlerraises partway through the loop,futures[task_idx]for all remaining, not-yet-processed indices are never resolved. Any concurrent producer awaiting one of those futures (viacached_parquet_info_for_task) will hang indefinitely instead of failing fast — this is the same failure mode flagged in a prior review pass on this exact helper (_receive).Wrap the loop and propagate the exception to all unfinished futures on failure.
🩹 Suggested fix
async def _receive() -> None: - for task_idx, scan in enumerate(scans): - msg = await recv_prefetched_parquet_metadata(context, ch_metadata) - cached = recv_prefetched_parquet_metadata_handler(msg, tuple(scan.paths)) - futures[task_idx].set_result(cached) + try: + for task_idx, scan in enumerate(scans): + msg = await recv_prefetched_parquet_metadata(context, ch_metadata) + cached = recv_prefetched_parquet_metadata_handler( + msg, tuple(scan.paths) + ) + futures[task_idx].set_result(cached) + except Exception as exc: + for fut in futures: + if not fut.done(): + fut.set_exception(exc) + raisePlease confirm whether this has already been addressed in the surrounding (unshown) code between lines 652-724, since the visible
_receivebody doesn't include any exception handling.#!/bin/bash ast-grep outline python/cudf_polars/cudf_polars/streaming/actor_graph/io.py --items all --match '_receive|_metadata_receiver_task|scan_node' --view expanded sed -n '600,800p' python/cudf_polars/cudf_polars/streaming/actor_graph/io.py🤖 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/actor_graph/io.py` around lines 647 - 651, The _receive helper in io.py can leave unresolved futures if recv_prefetched_parquet_metadata or recv_prefetched_parquet_metadata_handler throws before the loop finishes. Update _receive to catch exceptions, then fail every unfinished entry in futures (including the current and remaining task_idx values) so cached_parquet_info_for_task does not hang; use the existing _receive, futures, and recv_prefetched_parquet_metadata / recv_prefetched_parquet_metadata_handler symbols to place the fix in the correct async receiver path.
🤖 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.
Duplicate comments:
In `@python/cudf_polars/cudf_polars/streaming/actor_graph/io.py`:
- Around line 708-740: The metadata receiver task started by
_start_metadata_receiver in scan_node is still not explicitly cleaned up if
scan_node exits early or fails. Add a finally block around the shutdown_on_error
scope that cancels _metadata_receiver_task when it exists and awaits it to
completion, handling cancellation cleanly so no exception is left unretrieved.
Keep the fix local to the scan_node path that sets up metadata_futures and
_metadata_receiver_task.
- Around line 647-651: The _receive helper in io.py can leave unresolved futures
if recv_prefetched_parquet_metadata or recv_prefetched_parquet_metadata_handler
throws before the loop finishes. Update _receive to catch exceptions, then fail
every unfinished entry in futures (including the current and remaining task_idx
values) so cached_parquet_info_for_task does not hang; use the existing
_receive, futures, and recv_prefetched_parquet_metadata /
recv_prefetched_parquet_metadata_handler symbols to place the fix in the correct
async receiver path.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 65bfb0e5-e4f1-46da-8117-0ef84ba953a0
📒 Files selected for processing (7)
python/cudf_polars/cudf_polars/dsl/ir.pypython/cudf_polars/cudf_polars/dsl/utils/io.pypython/cudf_polars/cudf_polars/streaming/actor_graph/core.pypython/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.pypython/cudf_polars/cudf_polars/streaming/actor_graph/io.pypython/cudf_polars/cudf_polars/streaming/io.pypython/cudf_polars/tests/streaming/test_scan.py
🚧 Files skipped from review as they are similar to previous changes (3)
- python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py
- python/cudf_polars/tests/streaming/test_scan.py
- python/cudf_polars/cudf_polars/dsl/utils/io.py
|
CI hit a timeout in |
This lets the sender get arbitrarily far ahead of the receiever, which we want for parquet metadata footers (probably).
…fetch-metadata-actors-2
|
There were two things:
|
|
Sorry for pushing you in circles a bit here. I think I’m finally able to describe my main concern more concretely, but feel free to push back if you just want to get something in asap and move on to data pre-fetching... This is looking better, but I’m still worried that the metadata protocol is aligned to each A Would this be relatively clean if the scan actor owned a small For example:
Then |
|
|
||
|
|
||
| class ParquetMetadataCache: | ||
| """ |
There was a problem hiding this comment.
Each actor gets a reference to this object.
The .get method is where all the complexity is, but the goal is to ensure that things are only prefetched once. The first time we see .paths, we
- Create a Future whose result will be the
list[CachedParquetInfo]for.paths. Its tracked in the local cache. - Create and await a coroutine (offloaded to a worker thread) that fetches the metadata for
.paths. - Trampoline the result of that coro to the
Futurefor.paths
If a second actor comes along and asks for .paths while the coroutine is still fetching the metadata, it'll be pointed to the same Future. Upon future.set_result() that coro will wake up and be given the same parquet metadata.
If the result is already done, it's just handed the cached parquet metadata.
I'm a bit concerned about the overhead of doing anything at the individual That said, I'm not sure why we'd want any batching other than what's on the I do want to get some additional metrics/traces for when these Actors complete there work (and maybe even log things like individual messages, depending on the config). But my guess now is that the metadata actor gets quite far ahead of the Scan actor. |
Yes, this aligns with my thinking as well.
If we are reading 10 large files, we will have >10 Even if each |
|
This is currently slower than our metadata prefetching on 
There's no reason for that read_parquet_footers to take 3 seconds, which makes me think there's a problem with this lock mechanism. I'll do some more investigation. |
3 participants
Description
This moves the parquet metadata prefetching that (optionally) happens between IR lowering and streaming network execution into the streaming network. This ensures that any given
Scannode can begin executing as soon as its metadata is ready, rather than all the metadata in a query.At a high level, this PR adds new Actors to the streaming network that are responsible solely for fetching parquet metadata for some
paths. We make exactly one new Prefetching actor per Scan actor. The prefetching andThe actors send that metadata to the Actors running the
Scannodes.Checklist