Prefetch byte ranges using kvikio

[Matt711]

Description

Hybrid scan lets us have more control over what happens when executing a Scan. By default, the streaming engine uses cudf::io::read_parquet to compute each Scan. Using cudf::io::parquet::experimental::HybridScanReader in cudf_polars essentially lets us break up the I/O and compute that happens during a call to cudf::io::read_parquet. In hybrid-scan speak, we split the read into two passes: first we read only the filter columns and compute a row mask, then we read only the payload columns for rows that survive the filter and combine. The typical benefit is for a selective filter. We only transfer the payload columns for rows that pass, rather than reading everything upfront and discarding filtered rows afterward.

For remote IO, HybridScanReader exposes APIs to compute the exact byte ranges needed for the filter and payload columns separately. Once we have the byte ranges for a split, we prefetch the data using kvikio into pinned host memory and transfer to device once the producer is ready.

We only dispatch to HybridScanReader for SplitScans in the streaming engine. SplitScan is used by the streaming engine to parallelize reads of a single large file. FusedScan fuses multiple smaller files into one (think logically one larger file) and AFAICT would be a good candidate for the multi-file hybrid scan reader (see #22583). Finally for the in-memory engine, there isn't a good justification for using hybrid scan so it's only supported for the streaming engine.

Dispatch is also conditional on having a filter predicate (we'll fallback to default libcudf parquet reader otherwise).

So for this first PR: Enabling CUDF_POLARS__PARQUET_OPTIONS__USE_HYBRID_SCAN=1 means dispatching SplitScans to HybridScanReader whenever there's a filter predicate (that we support translating to libcudf).

• Contributes to [PERF]: Fully adopt hybrid scan reader in cudf-polars #23539 by adopting hybrid scan for SplitScan tasks. FusedScan tasks are left for a follow-up.
• Contributes to [PERF]: kvikio byte range prefetch backend for hybrid scan #23535 by using kvikio's EASY_THREADPOOL backend to drive prefetching. We're leaving investigation into kvikio's MULTI_POLL backend to a follow-up.

Checklist

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

[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.

[Matt711]

/ok to test 245827f

[wence-]

Perhaps I missed it, but for review purposes can you split the C++ hybrid scan changes, and the pylibcudf exposure, into separate PRs?

[wence-]

Can we use an existing vocab type. For example a rapidsmpf::Buffer should be fine? I think.

[Matt711]

I think so. We need two things first

1. make_buffer exposed on BufferResource
2. Implement __buffer__ to get a memoryview that we can hand to pread

[Matt711]

Perhaps I missed it, but for review purposes can you split the C++ hybrid scan changes, and the pylibcudf exposure, into separate PRs?

Yes, will open them shortly

[rjzamora]

I see a few places spread across this PR where the details of using HybridScan and prefetching data is very specific to SplitScan. Although this makes sense for prototyping, I have a strong feeling that we should start with infrastructure that is agnostic as possible to SplitScan-vs-FusedScan. You will obviously need the SplitScan variety to actually enable hybrid-scan (for now), but I'd consider the FusedScan variety a P0 follow-up to this work.

[rjzamora]

Maybe this check should be a dedicated helper/utility for now? I think you may do similar validation in other places, so it would possibly used there as well.

[rjzamora]

In the case that we are using hybrid scan (or maybe even whenever we have prefetched metadata), I think we should map our list of SplitScans or FusedScans onto a list of task objects that are row-group aligned and more friendly to hybrid scan. For example:

class AlignedParquetScan(IR):
    base_scan: Scan
    cached_parquet_info: list[CachedParquetInfo]
    row_groups_per_source: list[list[int]]
    parquet_options: ParquetOptions

Why?

• When hybrid scan is enabled and footer metadata has been prefetched, we can assign explicit row-group indices to each read task instead of relying on a uniform split_index / total_splits for every file. Hence “Aligned”: each task reads whole parquet row groups.
• This gives us a hybrid-scan-friendly execution node that is agnostic to how the original streaming task was planned (SplitScan vs FusedScan). That should make the follow-up multi-file path much easier: the first version may have exactly one cached_parquet_info item per task, but the same AlignedParquetScan design can later represent multi-source reads by filling cached_parquet_info and row_groups_per_source with one entry per source.
• This also centralizes hybrid-scan eligibility and fallback logic, instead of spreading SplitScan-specific checks across planning, prefetch, and evaluation.

[rjzamora]

I think you pointed this out to me offline. If I understand correctly, we are submitting prefetching tasks for all our data up-front, and we are relying on pinned-host memory reservations for back-pressure?

[rjzamora]

What happens when the projected columns are all filter columns?

For example, scan_parquet(...).select("x").filter(pl.col("x") < 10) seems like it could produce an empty payload-column set after materialize_filter_columns. Does this logic handle that case cleanly?