Fix use-after-destroy and stream ordering in Parquet IO utils

[mhaseeb123]

Description

This PR fixes the use-after-destroy and stream ordering (with PTDS input) issue (with host buffer source) in the fetch_byte_ranges_to_device_async IO utility used by parquet and hybrid scan.

See follow up PR #22550 that reduces the locked region size by moving all host_read_async outside it.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[coderabbitai]

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Batches host reads for parquet byte-range transfers: adds CUDA memcpy and tuple includes, reserves merged-range vectors, collects host-read futures returning owning buffers, resolves them to extract source pointers, and enqueues a single cudf::detail::memcpy_batch_async to copy to device destinations. Deferred return waits only on device-read futures.

Changes

Host-to-Device Batched Transfer

Layer / File(s)	Summary
Dependencies and vector pre-reservation cpp/src/io/parquet/io_utils/parquet_io_utils.cpp	CUDA memcpy include and <tuple> are added; capacity is reserved for merged I/O metadata vectors (io_offsets, io_sizes, destinations).
Host-read batch transfer data structures cpp/src/io/parquet/io_utils/parquet_io_utils.cpp	Replaces prior host-read worker task with std::future<host_read_buffer> futures (owning std::unique_ptr buffers) and adds vectors to accumulate batched-copy destinations, sizes, and resolved host-buffer pointers.
Batched host-to-device memcpy implementation cpp/src/io/parquet/io_utils/parquet_io_utils.cpp	Queues datasource.host_read_async futures, resolves them to obtain host data pointers, schedules a single cudf::detail::memcpy_batch_async for all host->device copies while keeping host buffers alive until enqueue, and returns a future that waits only on device-read futures.

---

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

---

Suggested labels

cuIO, 4 - Needs Review

---

Suggested reviewers

• ttnghia
• vuule
• pmattione-nvidia

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately describes the main fix: addressing use-after-destroy and stream ordering issues in Parquet IO utilities, which aligns with the code changes.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	The PR description clearly relates to the changeset, explaining the fix for use-after-destroy and stream ordering issues in fetch_byte_ranges_to_device_async, which aligns with the code modifications.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

cpp/src/io/parquet/io_utils/parquet_io_utils.cpp (1)

136-184: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Gate the blocking synchronizations to the paths that need them.

Line 137 now blocks even when every range uses host_read_async, and Lines 174-175 block even when every range uses device_read_async, so this *_async helper can end up fully synchronous before it returns. Those waits also happen while the static mutex from Line 142 is held, which serializes concurrent callers for the entire host-read/H2D phase.

Please restrict the pre-sync to the device-read path, keep the post-copy sync inside the host-read branch, and release the mutex once the async reads have been enqueued. As per coding guidelines, "Avoid unnecessary host-device synchronization that blocks the GPU pipeline".

🤖 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 `@cpp/src/io/parquet/io_utils/parquet_io_utils.cpp` around lines 136 - 184, The
current implementation calls stream.synchronize() unconditionally and holds
mutex across async enqueues and waits, causing unnecessary blocking; change the
logic in the function that builds device_read_tasks/host_read_tasks so that (1)
the initial stream.synchronize() is only invoked when at least one device read
is scheduled (i.e. when datasource.supports_device_read() &&
is_device_read_preferred(...) leads to any device_read_tasks), (2) release the
std::lock_guard<std::mutex> mutex immediately after enqueuing both
device_read_tasks and host_read_tasks so the critical section does not cover
subsequent waits or host-device copies, and (3) move the second
stream.synchronize() so it only runs inside the branch that performed host reads
(immediately after calling cudf::detail::memcpy_batch_async and before
host_read_buffers/other temporaries go out of scope); keep the deferred
sync_function that waits on device_read_tasks unchanged but ensure it is created
after the mutex is released.

🤖 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 `@cpp/src/io/parquet/io_utils/parquet_io_utils.cpp`:
- Around line 136-184: The current implementation calls stream.synchronize()
unconditionally and holds mutex across async enqueues and waits, causing
unnecessary blocking; change the logic in the function that builds
device_read_tasks/host_read_tasks so that (1) the initial stream.synchronize()
is only invoked when at least one device read is scheduled (i.e. when
datasource.supports_device_read() && is_device_read_preferred(...) leads to any
device_read_tasks), (2) release the std::lock_guard<std::mutex> mutex
immediately after enqueuing both device_read_tasks and host_read_tasks so the
critical section does not cover subsequent waits or host-device copies, and (3)
move the second stream.synchronize() so it only runs inside the branch that
performed host reads (immediately after calling cudf::detail::memcpy_batch_async
and before host_read_buffers/other temporaries go out of scope); keep the
deferred sync_function that waits on device_read_tasks unchanged but ensure it
is created after the mutex is released.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 3ecc1eea-9a86-4ce3-aed7-7b7a5e03514c

📥 Commits

Reviewing files that changed from the base of the PR and between 18af3a4 and 3d54fcf.

📒 Files selected for processing (1)

• cpp/src/io/parquet/io_utils/parquet_io_utils.cpp

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[mhaseeb123]

@coderabbitai full review

[coderabbitai]

✅ Actions performed

Full review triggered.

[aminaramoon]

LGTM!

[vuule]

Having a hard time convincing myself that there's no performance impact here in any case, but the changes look good.
Do you have any performance numbers?

[mhaseeb123]

Do you have any performance numbers?

Let me gather and paste some here

[mhaseeb123]

Performance Impact

This PR shows a rough ~12% end-to-end improvement when reading 16 x (1GB disk size) Parquet files (clickbench) using the hybrid_scan_multifile_single_step example with HOST_BUFFER datasource.

Note: Perf for FILEPATH datasource remains unchanged.

CC: @vuule

[mhaseeb123]

/merge