From 49709761545bcc4758f863d4d8ba7c08752afba3 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 17 Aug 2026 14:38:17 +0000 Subject: [PATCH 1/2] Adopt hybrid scan reader in cudf-polars for split scans --- .../cudf_polars/cudf_polars/dsl/utils/io.py | 39 +++- .../cudf_polars/cudf_polars/streaming/io.py | 192 +++++++++++++++++- .../cudf_polars/cudf_polars/utils/config.py | 25 +++ .../cudf_polars/tests/streaming/test_scan.py | 36 ++++ python/cudf_polars/tests/test_config.py | 5 + 5 files changed, 293 insertions(+), 4 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index 3ac9db651b01..f52b5da2ad9e 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -6,7 +6,7 @@ import concurrent.futures import contextlib -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING import pylibcudf as plc @@ -46,6 +46,27 @@ class CachedParquetInfo: path: str size: int | None file_metadata: plc.io.parquet_metadata.FileMetaData + # Pre-created during footer prefetch; shared across all splits and scans of this file. + # HybridScanReader is not cached: it holds mutable per-read state so each worker + # creates its own from the shared metadata. + _hybrid_scan_metadata: list[plc.io.experimental.HybridScanMetadata] = field( + default_factory=list, compare=False, repr=False + ) + + def hybrid_scan_reader( # pragma: no cover; only called from thread pool workers where coverage.py does not trace + self, + options: plc.io.parquet.ParquetReaderOptions, + ) -> plc.io.experimental.HybridScanReader: + """Return a fresh HybridScanReader backed by shared pre-parsed file metadata.""" + if not self._hybrid_scan_metadata: + self._hybrid_scan_metadata.append( + plc.io.experimental.HybridScanMetadata.from_parquet_metadata( + self.file_metadata, options + ) + ) + return plc.io.experimental.HybridScanReader.from_metadata( + self._hybrid_scan_metadata[0] + ) @nvtx_annotate_cudf_polars(message="fetch_parquet_footers_for_paths") @@ -99,10 +120,24 @@ def _prefetch_parquet_footers_for_paths(paths: list[str]) -> list[CachedParquetI ) ) - return [ + infos = [ CachedParquetInfo(path, size, file_metadata) for path, size, file_metadata in zip(paths, sizes, metadata, strict=True) ] + for info in infos: + options = ( + plc.io.parquet.ParquetReaderOptions.builder( + plc.io.SourceInfo([plc.io.types.FilepathSource(info.path, info.size)]) + ) + .decimal_width(plc.TypeId.DECIMAL128) + .build() + ) + info._hybrid_scan_metadata.append( + plc.io.experimental.HybridScanMetadata.from_parquet_metadata( + info.file_metadata, options + ) + ) + return infos @nvtx_annotate_cudf_polars(message="prefetch_parquet_file_metadata_for_ir") diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 4e061e6a3cb2..efeb4622181f 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -17,6 +17,7 @@ import pylibcudf as plc +from cudf_polars.containers import Column, DataFrame from cudf_polars.dsl.ir import ( IR, DataFrameScan, @@ -24,7 +25,9 @@ PythonScan, Scan, Sink, + _prepare_parquet_predicate, ) +from cudf_polars.dsl.to_ast import to_parquet_filter from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars from cudf_polars.streaming.base import ( IOPartitionFlavor, @@ -40,7 +43,10 @@ if TYPE_CHECKING: from collections.abc import Hashable, MutableMapping, Sequence - from cudf_polars.containers import DataFrame, DataType + import pylibcudf.expressions as plc_expr + from rmm.pylibrmm.stream import Stream + + from cudf_polars.containers import DataType from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import CachedParquetInfo, IRExecutionContext from cudf_polars.streaming.base import ( @@ -81,6 +87,13 @@ def scan_partition_plan( """Extract the partitioning plan of a Scan operation.""" if ir.typ == "parquet": blocksize: int = config_options.executor.target_partition_size + single_file = len(ir.paths) == 1 + # A single file always uses SplitScan when hybrid scan is enabled, so the + # hybrid reader can be used on it even when it would otherwise not split. + # The split factor is still size-based, so a large file is split into many. + hybrid_single_file = ( + single_file and config_options.parquet_options.use_hybrid_scan + ) if source := stats.scan_stats.get(ir): column_sizes = [ sz @@ -97,12 +110,18 @@ def scan_partition_plan( <= abs(file_size / k_hi - blocksize) else k_hi ) - if factor >= 2: + if factor >= 2 or hybrid_single_file: return IOPartitionPlan( factor, IOPartitionFlavor.SPLIT_FILES, estimated_chunk_bytes=file_size // factor, ) + elif hybrid_single_file: + return IOPartitionPlan( + 1, + IOPartitionFlavor.SPLIT_FILES, + estimated_chunk_bytes=file_size, + ) else: k_lo = min(blocksize // int(file_size), len(ir.paths)) k_hi = k_lo + 1 @@ -119,6 +138,9 @@ def scan_partition_plan( estimated_chunk_bytes=file_size * factor, ) + if hybrid_single_file: + return IOPartitionPlan(1, IOPartitionFlavor.SPLIT_FILES) + # TODO: Use file sizes for csv and json return IOPartitionPlan(1, IOPartitionFlavor.SINGLE_FILE) @@ -181,6 +203,136 @@ def expand_scan_for_rank( ) +def _fetch_byte_ranges( + source_info: plc.io.SourceInfo, + byte_ranges: list[plc.io.text.ByteRangeInfo], + stream: Stream, +) -> list[plc.gpumemoryview]: + return plc.io.parquet_io_utils.fetch_byte_ranges_to_device( + source_info, byte_ranges, stream=stream + ) + + +def _read_with_hybrid_scan( + schema: Schema, + paths: list[str], + with_columns: list[str] | None, + plc_filter: plc_expr.Expression, + row_group_indices: list[int], + stream: Stream, + cached_info: CachedParquetInfo, + *, + split_index: int = 0, + total_splits: int = 1, + stats_pruning: bool = True, +) -> DataFrame: + """Two-pass parquet read via HybridScanReader for a row-group-aligned split.""" + assert plc_filter is not None + assert len(paths) == 1, ( + "hybrid scan only supported for SplitScan; one physical file" + ) + with nvtx_annotate_cudf_polars( + message="HybridScan", payload=(split_index + 1, total_splits) + ): + source_info = plc.io.SourceInfo( + [plc.io.types.FilepathSource(cached_info.path, cached_info.size)] + ) + options = ( + plc.io.parquet.ParquetReaderOptions.builder(source_info) + .decimal_width(plc.TypeId.DECIMAL128) + .build() + ) + if with_columns is not None: + options.set_column_names(with_columns) + options.set_filter(plc_filter) + + reader = cached_info.hybrid_scan_reader(options) + + if stats_pruning: + row_group_indices = reader.filter_row_groups_with_stats( + row_group_indices, options, stream=stream + ) + + if row_group_indices: + bloom_ranges, _ = reader.secondary_filters_byte_ranges( + row_group_indices, options + ) + if bloom_ranges: + bloom_chunks = _fetch_byte_ranges(source_info, bloom_ranges, stream) + row_group_indices = reader.filter_row_groups_with_bloom_filters( + bloom_chunks, row_group_indices, options, stream=stream + ) + + if not row_group_indices: + col_names = with_columns if with_columns is not None else list(schema) + return DataFrame( + [ + Column( + plc.column_factories.make_empty_column( + schema[name].plc_type, stream=stream + ), + dtype=schema[name], + name=name, + ) + for name in col_names + ], + stream=stream, + ) + + # TODO: Consider implementing page-index stats pruning. For SplitScans, we can + # reuse the same page index for all splits of the same file, so the overhead of + # reading the page index can be amortized. For FusedScans, we would need to read + # the page index for all files, which may be too expensive. + row_mask = reader.build_all_true_row_mask(row_group_indices, stream=stream) + + filter_chunks = _fetch_byte_ranges( + source_info, + reader.filter_column_chunks_byte_ranges(row_group_indices, options), + stream, + ) + filter_tbl_w_meta = reader.materialize_filter_columns( + row_group_indices, + filter_chunks, + row_mask, + plc.io.experimental.UseDataPageMask.YES, + options, + stream=stream, + ) + + payload_chunks = _fetch_byte_ranges( + source_info, + reader.payload_column_chunks_byte_ranges(row_group_indices, options), + stream, + ) + payload_tbl_w_meta = reader.materialize_payload_columns( + row_group_indices, + payload_chunks, + row_mask, + plc.io.experimental.UseDataPageMask.YES, + options, + stream=stream, + ) + + filter_names = filter_tbl_w_meta.column_names(include_children=False) + payload_names = payload_tbl_w_meta.column_names(include_children=False) + filter_df = DataFrame.from_table( + filter_tbl_w_meta.tbl, + filter_names, + [schema[n] for n in filter_names], + stream=stream, + ) + payload_df = DataFrame.from_table( + payload_tbl_w_meta.tbl, + payload_names, + [schema[n] for n in payload_names], + stream=stream, + ) + stream.synchronize() + return DataFrame( + [*filter_df.columns, *payload_df.columns], stream=stream + ).select(list(schema.keys())) + + class SplitScan(IR): """ Input from a split file. @@ -336,6 +488,42 @@ def do_evaluate( skip_rgs = rg_stride * split_index skip_rows = sum(row_group_num_rows[:skip_rgs]) n_rows = sum(row_group_num_rows[skip_rgs : skip_rgs + rg_stride]) + # Hybrid scan reads through the prefetched, shared file metadata, so + # it is only used when footer prefetching is enabled. + # TODO: Investigate re-enabling for some of the excluded paths + # (row_index / include_file_paths). Needs performance investigation. + if ( + parquet_options.use_hybrid_scan + and cached_parquet_info is not None + and row_index is None + and include_file_paths is None + and predicate is not None + ): + stream = context.get_cuda_stream() + plc_filter = to_parquet_filter( + _prepare_parquet_predicate( + predicate.value, paths, schema, with_columns + ), + stream=stream, + ) + if plc_filter is not None: + end_rg = ( + total_row_groups + if split_index == total_splits - 1 + else skip_rgs + rg_stride + ) + return _read_with_hybrid_scan( + schema, + paths, + with_columns, + plc_filter, + list(range(skip_rgs, end_rg)), + stream, + cached_parquet_info[0], + split_index=split_index, + total_splits=total_splits, + stats_pruning=parquet_options._hybrid_scan_stats_pruning, + ) else: # There are not enough row-groups to align # all "total_splits" of our reads with row-group diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index ec144fc2fc73..517135d31b12 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -267,6 +267,10 @@ class ParquetOptions: When enabled, filter predicates are JIT-compiled to CUDA kernels for improved performance on large datasets with complex filters. Default is False. + use_hybrid_scan + Whether to use the two-pass ``HybridScanReader`` for ``SplitScan`` + tasks when a predicate can be pushed down to a parquet filter. + Default is False. """ _env_prefix = "CUDF_POLARS__PARQUET_OPTIONS" @@ -308,6 +312,23 @@ class ParquetOptions: default=UNSPECIFIED, ) ) + use_hybrid_scan: bool = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__USE_HYBRID_SCAN", + _bool_converter, + default=False, + ) + ) + # Internal benchmarking flag. When False, skips stats and bloom-filter pruning + # before the first pass of a hybrid scan so you can measure two-pass read + # overhead in isolation. No reason to set this to False in production. + _hybrid_scan_stats_pruning: bool = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__HYBRID_SCAN_STATS_PRUNING", + _bool_converter, + default=True, + ) + ) use_jit_filter: bool = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__USE_JIT_FILTER", @@ -331,6 +352,10 @@ def __post_init__(self) -> None: # noqa: D105 raise TypeError("max_row_group_samples must be an int") if not isinstance(self.prefetch_file_metadata, (bool, Unspecified)): raise TypeError("prefetch_file_metadata must be a bool when specified") + if not isinstance(self.use_hybrid_scan, bool): + raise TypeError("use_hybrid_scan must be a bool") + if not isinstance(self._hybrid_scan_stats_pruning, bool): + raise TypeError("_hybrid_scan_stats_pruning must be a bool") if not isinstance(self.use_jit_filter, bool): raise TypeError("use_jit_filter must be a bool") diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index a44af628ea1c..81c0ef009d07 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -370,6 +370,42 @@ def test_streaming_scan_raises() -> None: StreamingScan.do_evaluate([fused], scan, context=ctx) +@pytest.mark.parametrize( + "predicate,use_columns", + [ + # uses hybrid scan reader + (pl.col("x") < 1_000, None), + (pl.col("x") < 1_000, ["x", "z"]), + # fallsback to default parquet reader + (pl.col("y").str.contains("cat"), None), + (None, None), + ], +) +def test_split_scan_hybrid( + tmp_path: Path, + df: pl.DataFrame, + predicate: pl.Expr | None, + use_columns: list[str] | None, + streaming_engine_factory: Callable[..., StreamingEngine], +) -> None: + streaming_engine = streaming_engine_factory( + StreamingOptions( + target_partition_size=1_000, + parquet_options={ + "use_hybrid_scan": True, + "prefetch_file_metadata": True, + }, + ), + ) + make_partitioned_source(df, tmp_path, "parquet", n_files=1, row_group_size=100) + q = pl.scan_parquet(tmp_path) + if use_columns is not None: + q = q.select(use_columns) + if predicate is not None: + q = q.filter(predicate) + assert_gpu_result_equal(q, engine=streaming_engine) + + def test_scan_path_mismatch_raises() -> None: # This isn't reachable by polars' public API, so we test it directly. scan = _make_parquet_scan( diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 883d5cd06374..0cd616e9789f 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -361,6 +361,8 @@ def test_parquet_options_from_env(monkeypatch: pytest.MonkeyPatch) -> None: m.setenv("CUDF_POLARS__PARQUET_OPTIONS__PASS_READ_LIMIT", "200") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__MAX_FOOTER_SAMPLES", "0") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__MAX_ROW_GROUP_SAMPLES", "0") + m.setenv("CUDF_POLARS__PARQUET_OPTIONS__USE_HYBRID_SCAN", "0") + m.setenv("CUDF_POLARS__PARQUET_OPTIONS__HYBRID_SCAN_STATS_PRUNING", "0") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__PREFETCH_FILE_METADATA", "1") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__USE_JIT_FILTER", "1") @@ -373,6 +375,8 @@ def test_parquet_options_from_env(monkeypatch: pytest.MonkeyPatch) -> None: assert config.parquet_options.pass_read_limit == 200 assert config.parquet_options.max_footer_samples == 0 assert config.parquet_options.max_row_group_samples == 0 + assert config.parquet_options.use_hybrid_scan is False + assert config.parquet_options._hybrid_scan_stats_pruning is False assert config.parquet_options.prefetch_file_metadata is True assert config.parquet_options.use_jit_filter is True @@ -488,6 +492,7 @@ def test_fallback_mode_default(monkeypatch: pytest.MonkeyPatch) -> None: "max_footer_samples", "max_row_group_samples", "prefetch_file_metadata", + "_hybrid_scan_stats_pruning", "use_jit_filter", ], ) From 0b7b5d9dc5429f484e5f759bae4ed881a468a1b7 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 17 Aug 2026 14:51:04 +0000 Subject: [PATCH 2/2] Prefetch parquet byte ranges with kvikio for hybrid scan splits --- .../cudf_polars/cudf_polars/dsl/utils/io.py | 37 +- .../cudf_polars/streaming/actor_graph/io.py | 235 ++++++++----- .../cudf_polars/cudf_polars/streaming/io.py | 170 ++++++++- .../cudf_polars/streaming/prefetch.py | 330 ++++++++++++++++++ .../cudf_polars/cudf_polars/utils/config.py | 12 + .../cudf_polars/tests/streaming/test_scan.py | 1 + python/cudf_polars/tests/test_config.py | 1 + 7 files changed, 685 insertions(+), 101 deletions(-) create mode 100644 python/cudf_polars/cudf_polars/streaming/prefetch.py diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index f52b5da2ad9e..f6f6636eb17e 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -7,7 +7,12 @@ import concurrent.futures import contextlib from dataclasses import dataclass, field -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any + +try: # pragma: no cover; kvikio is optional + import kvikio +except ImportError: + kvikio = None import pylibcudf as plc @@ -52,6 +57,7 @@ class CachedParquetInfo: _hybrid_scan_metadata: list[plc.io.experimental.HybridScanMetadata] = field( default_factory=list, compare=False, repr=False ) + _remote_handle: list[Any] = field(default_factory=list, compare=False, repr=False) def hybrid_scan_reader( # pragma: no cover; only called from thread pool workers where coverage.py does not trace self, @@ -68,6 +74,19 @@ def hybrid_scan_reader( # pragma: no cover; only called from thread pool worker self._hybrid_scan_metadata[0] ) + def remote_handle(self) -> Any: # pragma: no cover; requires kvikio + """Return the kvikio handle for this file.""" + if not self._remote_handle: + if kvikio is None: + raise ImportError("kvikio is required for hybrid scan prefetching") + if plc.io.SourceInfo._is_remote_uri(self.path): + self._remote_handle.append( + kvikio.RemoteFile.open(self.path, nbytes=self.size) + ) + else: + self._remote_handle.append(kvikio.CuFile(self.path)) + return self._remote_handle[0] + @nvtx_annotate_cudf_polars(message="fetch_parquet_footers_for_paths") def _prefetch_parquet_footers_for_paths(paths: list[str]) -> list[CachedParquetInfo]: @@ -93,15 +112,10 @@ def _prefetch_parquet_footers_for_paths(paths: list[str]) -> list[CachedParquetI # For now, we'll just use kvikio to explicitly get the size. sizes: list[int | None] = [] - try: # pragma: no cover; kvikio is optional - import kvikio - except ImportError: - kvikio = None - for path in paths: - if ( - paths and kvikio is not None and plc.io.SourceInfo._is_remote_uri(path) - ): # pragma: no cover; kvikio is optional + if kvikio is not None and plc.io.SourceInfo._is_remote_uri( + path + ): # pragma: no cover # We're OK to use `kvikio.RemoteFile.open` here. It does make an HTTP HEAD # request for S3/HTTP endpoints, but that's the entire reason we're running # this code. So long as it makes just *one* HTTP request, there's no advantage @@ -137,6 +151,11 @@ def _prefetch_parquet_footers_for_paths(paths: list[str]) -> list[CachedParquetI info.file_metadata, options ) ) + if kvikio is not None: # pragma: no cover; requires kvikio + # Open kvikio handles eagerly on the main thread before any prefetch workers + # start, so all splits sharing a file get the same handle without races. + for info in infos: + info.remote_handle() return infos diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py index e89035e12e2a..610be3bd1e63 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio +import contextlib import functools import io import math @@ -20,10 +21,20 @@ from rapidsmpf.streaming.core.message import Message from cudf_polars.containers import DataFrame -from cudf_polars.dsl.ir import IR, DataFrameScan, PythonScan, Sink -from cudf_polars.dsl.tracing import Scope, log -from cudf_polars.streaming.actor_graph.dispatch import generate_ir_sub_network -from cudf_polars.streaming.actor_graph.nodes import define_actor, shutdown_on_error +from cudf_polars.dsl.ir import ( + IR, + DataFrameScan, + PythonScan, + Sink, +) +from cudf_polars.dsl.tracing import Scope, log, nvtx_annotate_cudf_polars +from cudf_polars.streaming.actor_graph.dispatch import ( + generate_ir_sub_network, +) +from cudf_polars.streaming.actor_graph.nodes import ( + define_actor, + shutdown_on_error, +) from cudf_polars.streaming.actor_graph.tracing import send_chunk from cudf_polars.streaming.actor_graph.utils import ( ChannelManager, @@ -35,11 +46,14 @@ send_metadata, ) from cudf_polars.streaming.io import ( + SplitScan, StreamingScan, StreamingSink, + _evaluate_with_prefetch, _prepare_sink_directory, _sink_to_file, ) +from cudf_polars.streaming.prefetch import HybridScanPrefetchExecutor from cudf_polars.streaming.rank_aware_source import RankAwareSource if TYPE_CHECKING: @@ -56,7 +70,7 @@ IOPartitionPlan, PartitionInfo, ) - from cudf_polars.streaming.io import FusedScan, SplitScan + from cudf_polars.streaming.io import FusedScan, PrefetchedByteRanges class Lineariser: @@ -497,6 +511,21 @@ def _( return nodes, channels +def evaluate_with_prefetch( + scan: SplitScan, + prefetcher: HybridScanPrefetchExecutor, + task_idx: int, + *, + context: IRExecutionContext, +) -> DataFrame: + """Evaluate a scan using parquet byte ranges prefetched into pinned host memory.""" + prefetched: PrefetchedByteRanges | None = prefetcher.result(task_idx) + if prefetched is None: + # The predicate could not be expressed as a parquet filter for this split. + return scan.do_evaluate(*scan._non_child_args, context=context) + return _evaluate_with_prefetch(scan, prefetched, context=context) + + async def read_chunk( context: Context, scan: IR, @@ -505,6 +534,8 @@ async def read_chunk( ir_context: IRExecutionContext, estimated_chunk_bytes: int, tracer: ActorTracer | None = None, + *, + prefetcher: HybridScanPrefetchExecutor | None = None, ) -> None: """ Read a chunk from disk and send it to the output channel. @@ -526,6 +557,9 @@ async def read_chunk( for admission before launching the read. tracer The actor tracer for collecting runtime statistics. + prefetcher + Optional prefetch pipeline. When set, retrieves the prefetched + I/O result before evaluating. """ reservation_bytes = ( estimated_chunk_bytes @@ -540,17 +574,27 @@ async def read_chunk( ) admitted = time.monotonic_ns() with opaque_memory_usage(reservation): - df = await ir_context.to_thread( - scan.do_evaluate, - *scan._non_child_args, - context=ir_context, - ) - chunk = TableChunk.from_pylibcudf_table( - df.table, - df.stream, - exclusive_view=True, - br=context.br(), - ) + if prefetcher is not None: + df = await ir_context.to_thread( + evaluate_with_prefetch, + scan, # type: ignore[arg-type] + prefetcher, + seq_num, + context=ir_context, + ) + else: + df = await ir_context.to_thread( + scan.do_evaluate, + *scan._non_child_args, + context=ir_context, + ) + with nvtx_annotate_cudf_polars(message="TableChunk.from_pylibcudf_table"): + chunk = TableChunk.from_pylibcudf_table( + df.table, + df.stream, + exclusive_view=True, + br=context.br(), + ) stop = time.monotonic_ns() log( "IO Task", @@ -575,6 +619,7 @@ async def scan_node( ch_out: Channel[TableChunk], *, num_producers: int, + num_prefetch_workers: int | None, estimated_chunk_bytes: int, ) -> None: """ @@ -592,79 +637,109 @@ async def scan_node( The output Channel[TableChunk]. num_producers The number of producers to use for the scan node. + num_prefetch_workers + The number of prefetch workers for the hybrid scan prefetch pipeline. + When ``None``, uses one worker per split. estimated_chunk_bytes Estimated retained output size of each chunk in bytes. Used to estimate peak memory for admission before launching each read. """ scans: Sequence[SplitScan] | Sequence[FusedScan] = ir.scans - async with shutdown_on_error( - context, ch_out, trace_ir=ir, ir_context=ir_context - ) as tracer: - # Send basic metadata - await send_metadata( - ch_out, - context, - ChannelMetadata(local_count=len(scans)), + first = scans[0] if scans else None + use_prefetch = ( + first is not None + and ir.scan_type == "split" + and first.parquet_options.use_hybrid_scan + and first.parquet_options.prefetch_file_metadata + and first.cached_parquet_info is not None + and first.base_scan.predicate is not None + and isinstance(first, SplitScan) + and first.total_splits + <= len(first.cached_parquet_info[0].file_metadata.row_group_num_rows) + and context.br().pinned_mr is not None + ) + prefetcher: HybridScanPrefetchExecutor | None = ( + HybridScanPrefetchExecutor.from_scans( + list(scans), # type: ignore[arg-type] + num_workers=num_prefetch_workers + if num_prefetch_workers is not None + else len(scans), + context=context, ) + if use_prefetch + else None + ) + with prefetcher or contextlib.nullcontext(): + async with shutdown_on_error( + context, ch_out, trace_ir=ir, ir_context=ir_context + ) as tracer: + # Send basic metadata + await send_metadata( + ch_out, + context, + ChannelMetadata(local_count=len(scans)), + ) - # If there is nothing to scan, drain the channel and return - if len(scans) == 0: - await ch_out.drain(context) - return - - # If there is only one scan or one producer, we can - # skip the lineariser and read the chunks directly - if len(scans) == 1 or num_producers == 1: - for seq_num, scan in enumerate(scans): - await read_chunk( - context, - scan, - seq_num, - ch_out, - ir_context, - estimated_chunk_bytes, - tracer=tracer, - ) - await ch_out.drain(context) - return - - # Use Lineariser to ensure ordered delivery - num_producers = min(num_producers, len(scans)) - lineariser = Lineariser(context, ch_out, num_producers) - - # Assign tasks to producers using round-robin - producer_tasks: list[list[tuple[int, SplitScan | FusedScan]]] = [ - [] for _ in range(num_producers) - ] - for task_idx, scan in enumerate(scans): - producer_id = task_idx % num_producers - # mypy resolves __iter__ on union-of-sequences to the common base (IR) - producer_tasks[producer_id].append((task_idx, scan)) # type: ignore[arg-type] - - async def _producer(producer_id: int, ch_out: Channel) -> None: - for task_idx, scan in producer_tasks[producer_id]: - await read_chunk( - context, - scan, - task_idx, - ch_out, - ir_context, - estimated_chunk_bytes, - tracer=tracer, + # If there is nothing to scan, drain the channel and return + if len(scans) == 0: + await ch_out.drain(context) + return + + # If there is only one scan or one producer, we can + # skip the lineariser and read the chunks directly + if len(scans) == 1 or num_producers == 1: + for seq_num, scan in enumerate(scans): + await read_chunk( + context, + scan, + seq_num, + ch_out, + ir_context, + estimated_chunk_bytes, + tracer=tracer, + prefetcher=prefetcher, + ) + await ch_out.drain(context) + return + + # Use Lineariser to ensure ordered delivery + num_producers = min(num_producers, len(scans)) + lineariser = Lineariser(context, ch_out, num_producers) + + # Assign tasks to producers using round-robin + producer_tasks: list[list[tuple[int, SplitScan | FusedScan]]] = [ + [] for _ in range(num_producers) + ] + for task_idx, scan in enumerate(scans): + producer_id = task_idx % num_producers + # mypy resolves __iter__ on union-of-sequences to the common base (IR) + producer_tasks[producer_id].append((task_idx, scan)) # type: ignore[arg-type] + + async def _producer(producer_id: int, ch_out: Channel) -> None: + for task_idx, scan in producer_tasks[producer_id]: + await read_chunk( + context, + scan, + task_idx, + ch_out, + ir_context, + estimated_chunk_bytes, + tracer=tracer, + prefetcher=prefetcher, + ) + await ch_out.drain(context) + + async with ( + shutdown_on_error(context, *lineariser.input_channels, trace_ir=ir), + ): + await gather_in_task_group( + lineariser.drain(), + *( + _producer(i, ch_in) + for i, ch_in in enumerate(lineariser.input_channels) + ), ) - await ch_out.drain(context) - - async with ( - shutdown_on_error(context, *lineariser.input_channels, trace_ir=ir), - ): - await gather_in_task_group( - lineariser.drain(), - *( - _producer(i, ch_in) - for i, ch_in in enumerate(lineariser.input_channels) - ), - ) @generate_ir_sub_network.register(StreamingScan) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index efeb4622181f..b6d9ed2f902a 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -13,9 +13,12 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, Self, overload +import nvtx + import polars as pl import pylibcudf as plc +from rmm import DeviceBuffer from cudf_polars.containers import Column, DataFrame from cudf_polars.dsl.ir import ( @@ -28,7 +31,7 @@ _prepare_parquet_predicate, ) from cudf_polars.dsl.to_ast import to_parquet_filter -from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars +from cudf_polars.dsl.tracing import CUDF_POLARS_NVTX_DOMAIN, nvtx_annotate_cudf_polars from cudf_polars.streaming.base import ( IOPartitionFlavor, IOPartitionPlan, @@ -43,6 +46,8 @@ if TYPE_CHECKING: from collections.abc import Hashable, MutableMapping, Sequence + from kvikio.cufile import IOFuture + import pylibcudf.expressions as plc_expr from rmm.pylibrmm.stream import Stream @@ -55,6 +60,7 @@ StatsCollector, ) from cudf_polars.streaming.dispatch import LowerIRTransformer + from cudf_polars.streaming.prefetch import PinnedBuffer from cudf_polars.typing import Schema from cudf_polars.utils.config import ( ConfigOptions, @@ -203,6 +209,45 @@ def expand_scan_for_rank( ) +@dataclasses.dataclass +class PrefetchedByteRanges: + """Prefetched byte ranges and pinned host buffers for a single scan task.""" + + row_group_indices: list[int] + filter_ranges: list[plc.io.text.ByteRangeInfo] + payload_ranges: list[plc.io.text.ByteRangeInfo] + filter_host: memoryview | None + payload_host: memoryview | None + filter_futures: list[IOFuture] = dataclasses.field( + default_factory=list, compare=False, repr=False + ) + payload_futures: list[IOFuture] = dataclasses.field( + default_factory=list, compare=False, repr=False + ) + filter_buf: PinnedBuffer | None = dataclasses.field( + default=None, compare=False, repr=False + ) + payload_buf: PinnedBuffer | None = dataclasses.field( + default=None, compare=False, repr=False + ) + + @classmethod + def empty(cls) -> PrefetchedByteRanges: + """Return a fully-pruned split with no rows to read.""" + return cls( + row_group_indices=[], + filter_ranges=[], + payload_ranges=[], + filter_host=None, + payload_host=None, + ) + + def release(self) -> None: + """Release pinned host memory reservations after the H2D copy completes.""" + self.filter_buf = None + self.payload_buf = None + + def _fetch_byte_ranges( source_info: plc.io.SourceInfo, byte_ranges: list[plc.io.text.ByteRangeInfo], @@ -213,6 +258,41 @@ def _fetch_byte_ranges( ) +def copy_host_ranges_to_device( + host: memoryview, + ranges: list[plc.io.text.ByteRangeInfo], + futures: list[IOFuture], + stream: Stream, + *, + base_scan_id: int = 0, + split_index: int = 0, + total_splits: int = 1, + label: str = "", +) -> list[plc.gpumemoryview]: + """Wait for in-flight S3 reads then copy pinned host ranges to device.""" + total = sum(r.size for r in ranges) + if not total: + return [] + rng = nvtx.start_range("copy_host_ranges_to_device", domain=CUDF_POLARS_NVTX_DOMAIN) + with nvtx_annotate_cudf_polars( + message=f"pread_ranges:wait:{label}" if label else "pread_ranges:wait", + payload=(base_scan_id, split_index + 1, total_splits, total), + ): + for f in futures: + f.get() + # TODO: Reserve device memory via rapidsmpf before allocating. + buf = DeviceBuffer(size=total) + buf.copy_from_host(host[:total], stream=stream) + gv = plc.gpumemoryview(buf) + result = [] + offset = 0 + for r in ranges: + result.append(gv.byte_slice(slice(offset, offset + r.size))) + offset += r.size + nvtx.end_range(rng) + return result + + def _read_with_hybrid_scan( schema: Schema, paths: list[str], @@ -221,10 +301,12 @@ def _read_with_hybrid_scan( row_group_indices: list[int], stream: Stream, cached_info: CachedParquetInfo, + base_scan_id: int, *, split_index: int = 0, total_splits: int = 1, stats_pruning: bool = True, + prefetched: PrefetchedByteRanges | None = None, ) -> DataFrame: """Two-pass parquet read via HybridScanReader for a row-group-aligned split.""" assert plc_filter is not None @@ -285,11 +367,23 @@ def _read_with_hybrid_scan( # the page index for all files, which may be too expensive. row_mask = reader.build_all_true_row_mask(row_group_indices, stream=stream) - filter_chunks = _fetch_byte_ranges( - source_info, - reader.filter_column_chunks_byte_ranges(row_group_indices, options), - stream, - ) + if prefetched is not None and prefetched.filter_host is not None: + filter_chunks = copy_host_ranges_to_device( + prefetched.filter_host, + prefetched.filter_ranges, + prefetched.filter_futures, + stream, + base_scan_id=base_scan_id, + split_index=split_index, + total_splits=total_splits, + label="filter", + ) + else: + filter_chunks = _fetch_byte_ranges( + source_info, + reader.filter_column_chunks_byte_ranges(row_group_indices, options), + stream, + ) filter_tbl_w_meta = reader.materialize_filter_columns( row_group_indices, filter_chunks, @@ -299,11 +393,23 @@ def _read_with_hybrid_scan( stream=stream, ) - payload_chunks = _fetch_byte_ranges( - source_info, - reader.payload_column_chunks_byte_ranges(row_group_indices, options), - stream, - ) + if prefetched is not None and prefetched.payload_host is not None: + payload_chunks = copy_host_ranges_to_device( + prefetched.payload_host, + prefetched.payload_ranges, + prefetched.payload_futures, + stream, + base_scan_id=base_scan_id, + split_index=split_index, + total_splits=total_splits, + label="payload", + ) + else: + payload_chunks = _fetch_byte_ranges( + source_info, + reader.payload_column_chunks_byte_ranges(row_group_indices, options), + stream, + ) payload_tbl_w_meta = reader.materialize_payload_columns( row_group_indices, payload_chunks, @@ -328,11 +434,47 @@ def _read_with_hybrid_scan( stream=stream, ) stream.synchronize() + if prefetched is not None: + prefetched.release() return DataFrame( [*filter_df.columns, *payload_df.columns], stream=stream ).select(list(schema.keys())) +def _evaluate_with_prefetch( + scan: SplitScan, + prefetched: PrefetchedByteRanges, + *, + context: IRExecutionContext, +) -> DataFrame: + """Evaluate a SplitScan using already-prefetched I/O results.""" + stream = context.get_cuda_stream() + predicate = scan.base_scan.predicate + assert predicate is not None + plc_filter = to_parquet_filter( + _prepare_parquet_predicate( + predicate.value, scan.paths, scan.schema, scan.base_scan.with_columns + ), + stream=stream, + ) + assert plc_filter is not None + assert scan.cached_parquet_info is not None + return _read_with_hybrid_scan( + scan.schema, + scan.paths, + scan.base_scan.with_columns, + plc_filter, + prefetched.row_group_indices, + stream, + scan.cached_parquet_info[0], + id(scan.base_scan), + split_index=scan.split_index, + total_splits=scan.total_splits, + stats_pruning=False, + prefetched=prefetched, + ) + + class SplitScan(IR): """ Input from a split file. @@ -360,7 +502,7 @@ class SplitScan(IR): "total_splits", "parquet_options", ) - _n_non_child_args = 13 + _n_non_child_args = 15 base_scan: Scan """Scan operation this node is based on.""" paths: list[str] @@ -402,6 +544,7 @@ def __init__( base_scan.include_file_paths, base_scan.predicate, parquet_options, + id(base_scan), cached_parquet_info, ) self.parquet_options = parquet_options @@ -440,6 +583,7 @@ def do_evaluate( include_file_paths: str | None, predicate: NamedExpr | None, parquet_options: ParquetOptions, + base_scan_id: int, cached_parquet_info: list[CachedParquetInfo] | None, *, context: IRExecutionContext, @@ -520,10 +664,12 @@ def do_evaluate( list(range(skip_rgs, end_rg)), stream, cached_parquet_info[0], + base_scan_id, split_index=split_index, total_splits=total_splits, stats_pruning=parquet_options._hybrid_scan_stats_pruning, ) + else: # There are not enough row-groups to align # all "total_splits" of our reads with row-group diff --git a/python/cudf_polars/cudf_polars/streaming/prefetch.py b/python/cudf_polars/cudf_polars/streaming/prefetch.py new file mode 100644 index 000000000000..9a524bebf6c7 --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/prefetch.py @@ -0,0 +1,330 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hybrid scan prefetch pipeline.""" + +from __future__ import annotations + +import asyncio +import ctypes +from concurrent.futures import ThreadPoolExecutor +from typing import TYPE_CHECKING, Any, Self + +import pylibcudf as plc +from rapidsmpf.memory.buffer import MemoryType +from rapidsmpf.streaming.core.memory_reserve_or_wait import reserve_memory + +from cudf_polars.dsl.ir import _prepare_parquet_predicate +from cudf_polars.dsl.to_ast import to_parquet_filter +from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars +from cudf_polars.streaming.io import PrefetchedByteRanges, _fetch_byte_ranges + +if TYPE_CHECKING: + from concurrent.futures import Future + + from kvikio.cufile import CuFile, IOFuture + from kvikio.remote_file import RemoteFile + + from rapidsmpf.memory.memory_reservation import MemoryReservation + from rapidsmpf.memory.pinned_memory_resource import PinnedMemoryResource + from rapidsmpf.streaming.core.context import Context + from rmm.pylibrmm.stream import Stream + + from cudf_polars.streaming.io import SplitScan + + +class PinnedBuffer: + """Pinned host buffer backed by a rapidsmpf PinnedMemoryResource pool.""" + + __slots__ = ("array", "mr", "nbytes", "ptr", "reservation", "stream") + + def __init__( + self, + mr: PinnedMemoryResource, + nbytes: int, + stream: Stream, + reservation: MemoryReservation, + ) -> None: + self.mr = mr + self.nbytes = nbytes + self.stream = stream + self.reservation = reservation + self.ptr = mr.allocate(nbytes, stream) + self.array = memoryview((ctypes.c_uint8 * nbytes).from_address(self.ptr)) + + def __del__(self) -> None: # noqa: D105 + # Guard against partial init. + if hasattr(self, "reservation"): + self.reservation.clear() + if hasattr(self, "ptr"): + self.mr.deallocate(self.ptr, self.nbytes, self.stream) + + +def pread_ranges( + handle: CuFile | RemoteFile, + ranges: list[plc.io.text.ByteRangeInfo], + pinned_mr: PinnedMemoryResource, + stream: Stream, + context: Context, + loop: asyncio.AbstractEventLoop, +) -> tuple[memoryview | None, list[IOFuture], PinnedBuffer | None]: + """Issue concurrent async reads for each range into a single pinned host buffer.""" + total = sum(r.size for r in ranges) + if not total: + return None, [], None + # Blocks this worker thread, not the event loop. The loop stays free to + # run other coroutines while we wait for the reservation. + with nvtx_annotate_cudf_polars(message="reserve_pinned_memory", payload=total): + reservation = asyncio.run_coroutine_threadsafe( + reserve_memory( + context, + size=total, + net_memory_delta=total, + mem_type=MemoryType.PINNED_HOST, + ), + loop, + ).result() + buf = PinnedBuffer(pinned_mr, total, stream, reservation) + futures = [] + offset = 0 + with nvtx_annotate_cudf_polars(message="read_ranges:submit", payload=total): + for r in ranges: + futures.append( + handle.pread( + buf.array[offset : offset + r.size], + size=r.size, + file_offset=r.offset, + ) + ) + offset += r.size + return buf.array, futures, buf + + +def prefetch_scan_byte_ranges( + scan: SplitScan, + stream: Stream, + pinned_mr: PinnedMemoryResource, + context: Context, + loop: asyncio.AbstractEventLoop, +) -> PrefetchedByteRanges | None: + """ + Run stats and bloom pruning for one SplitScan and issue async reads. + + Parameters + ---------- + scan + The split scan task to prefetch. + stream + CUDA stream used for filter expression compilation. + pinned_mr + Pinned memory resource to allocate host buffers from. + context + rapidsmpf context used for pinned memory reservation. + loop + Event loop used to submit the reservation coroutine from a worker thread. + + Returns + ------- + PrefetchedByteRanges | None + ``None`` when the predicate cannot be expressed as a parquet filter, + in which case the producer falls back to ``SplitScan.do_evaluate``. + :meth:`PrefetchedByteRanges.empty` when all row groups are pruned away. + """ + cached_info = scan.cached_parquet_info + assert cached_info is not None + + row_group_num_rows = cached_info[0].file_metadata.row_group_num_rows + total_row_groups = len(row_group_num_rows) + + rg_stride = total_row_groups // scan.total_splits + skip_rgs = rg_stride * scan.split_index + end_rg = ( + total_row_groups + if scan.split_index == scan.total_splits - 1 + else skip_rgs + rg_stride + ) + row_group_indices = list(range(skip_rgs, end_rg)) + + predicate = scan.base_scan.predicate + assert predicate is not None + + with nvtx_annotate_cudf_polars(message="to_parquet_filter"): + plc_filter = to_parquet_filter( + _prepare_parquet_predicate( + predicate.value, scan.paths, scan.schema, scan.base_scan.with_columns + ), + stream=stream, + ) + if plc_filter is None: + return None + + with nvtx_annotate_cudf_polars(message="build_reader_options"): + options = ( + plc.io.parquet.ParquetReaderOptions.builder( + plc.io.SourceInfo( + [ + plc.io.types.FilepathSource( + cached_info[0].path, cached_info[0].size + ) + ] + ) + ) + .decimal_width(plc.TypeId.DECIMAL128) + .build() + ) + if scan.base_scan.with_columns is not None: + options.set_column_names(scan.base_scan.with_columns) + options.set_filter(plc_filter) + + with nvtx_annotate_cudf_polars(message="hybrid_scan_reader"): + reader = cached_info[0].hybrid_scan_reader(options) + + if scan.parquet_options._hybrid_scan_stats_pruning: + with nvtx_annotate_cudf_polars(message="filter_row_groups_with_stats"): + row_group_indices = reader.filter_row_groups_with_stats( + row_group_indices, options, stream=stream + ) + + if row_group_indices: + bloom_ranges, _ = reader.secondary_filters_byte_ranges( + row_group_indices, options + ) + if bloom_ranges: + with nvtx_annotate_cudf_polars( + message="filter_row_groups_with_bloom_filters" + ): + bloom_chunks = _fetch_byte_ranges( + plc.io.SourceInfo( + [ + plc.io.types.FilepathSource( + cached_info[0].path, cached_info[0].size + ) + ] + ), + bloom_ranges, + stream, + ) + row_group_indices = reader.filter_row_groups_with_bloom_filters( + bloom_chunks, row_group_indices, options, stream=stream + ) + + if not row_group_indices: + return PrefetchedByteRanges.empty() + + with nvtx_annotate_cudf_polars(message="byte_range_computation"): + filter_ranges = reader.filter_column_chunks_byte_ranges( + row_group_indices, options + ) + payload_ranges = reader.payload_column_chunks_byte_ranges( + row_group_indices, options + ) + + handle = cached_info[0].remote_handle() + filter_bytes = sum(r.size for r in filter_ranges) + payload_bytes = sum(r.size for r in payload_ranges) + # TODO: coalesce nearby ranges before issuing pread calls. + # https://github.com/rapidsai/cudf/pull/23317#discussion_r3668809937 + with nvtx_annotate_cudf_polars( + message="pread_filter_and_payload", + payload=( + id(scan.base_scan), + scan.split_index + 1, + scan.total_splits, + filter_bytes, + payload_bytes, + ), + ): + filter_host, filter_futures, filter_buf = pread_ranges( + handle, filter_ranges, pinned_mr, stream, context, loop + ) + payload_host, payload_futures, payload_buf = pread_ranges( + handle, payload_ranges, pinned_mr, stream, context, loop + ) + + return PrefetchedByteRanges( + row_group_indices=row_group_indices, + filter_ranges=filter_ranges, + payload_ranges=payload_ranges, + filter_host=filter_host, + payload_host=payload_host, + filter_futures=filter_futures, + payload_futures=payload_futures, + filter_buf=filter_buf, + payload_buf=payload_buf, + ) + + +# TODO: Replace with a cucascade::io::datasource that accepts fadvise() hints +# issued before evaluation, so pre-reading is driven by the datasource layer +# rather than a separate host-pinned executor. +class HybridScanPrefetchExecutor: + """Prefetch executor for SplitScan tasks.""" + + def __init__( + self, + futures: list[Future[PrefetchedByteRanges | None]], + executor: ThreadPoolExecutor, + ): + self.futures = futures + self._executor = executor + + @classmethod + def from_scans( + cls, + scans: list[SplitScan], + num_workers: int, + context: Context, + ) -> Self: + """ + Submit prefetch tasks for all scans. + + Parameters + ---------- + scans + Tasks to prefetch. + num_workers + Number of background worker threads. + context + rapidsmpf context. Pinned memory must be enabled. + + Returns + ------- + HybridScanPrefetchExecutor + + Raises + ------ + ValueError + If ``context.br().pinned_mr`` is ``None``. + """ + pinned_mr = context.br().pinned_mr + if pinned_mr is None: + raise ValueError( + "HybridScanPrefetchExecutor requires a PinnedMemoryResource; " + "enable pinned memory via --pinned-memory." + ) + loop = asyncio.get_running_loop() + stream_pool = context.br().stream_pool + # TODO: Consider reusing ir_context.py_executor instead of a dedicated pool. + executor = ThreadPoolExecutor( + max_workers=num_workers, + thread_name_prefix="hybrid-prefetch", + ) + + def _task(s: SplitScan) -> PrefetchedByteRanges | None: + return prefetch_scan_byte_ranges( + s, stream_pool.get_stream(), pinned_mr, context, loop + ) + + futures = [executor.submit(_task, scan) for scan in scans] + return cls(futures, executor) + + def __enter__(self) -> Self: + """Enter the context manager.""" + return self + + def __exit__(self, *args: Any) -> None: + """Shut down the thread pool, cancelling pending futures.""" + self._executor.shutdown(cancel_futures=True, wait=False) + + def result(self, task_idx: int) -> PrefetchedByteRanges | None: + """Block until the tasks' prefetch result is ready and return it.""" + return self.futures[task_idx].result() diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 517135d31b12..eb9e6a03691e 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -747,6 +747,9 @@ class StreamingExecutor: - ``executor_options`` passed to ``polars.GPUEngine`` - the ``CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS`` environment variable + num_prefetch_workers + Number of prefetch worker threads for the hybrid scan prefetch pipeline. + Default is 2. Set to ``None`` to use one worker per split. num_py_executors Maximum number of workers for the Python ThreadPoolExecutor. Default is 8. @@ -816,6 +819,11 @@ class StreamingExecutor: f"{_env_prefix}__MAX_CONCURRENT_IO_TASKS", int, default=2 ) ) + num_prefetch_workers: int | None = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__NUM_PREFETCH_WORKERS", int, default=2 + ) + ) num_py_executors: int = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__NUM_PY_EXECUTORS", int, default=8 @@ -903,6 +911,10 @@ def __post_init__(self) -> None: # noqa: D105 raise TypeError("client_device_threshold must be a float") if not isinstance(self.max_concurrent_io_tasks, int): raise TypeError("max_concurrent_io_tasks must be an int") + if self.num_prefetch_workers is not None and not isinstance( + self.num_prefetch_workers, int + ): + raise TypeError("num_prefetch_workers must be an int or None") if not isinstance(self.num_py_executors, int): raise TypeError("num_py_executors must be an int") diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 81c0ef009d07..dabaf2a90a86 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -470,6 +470,7 @@ def test_split_scan_do_evaluate_missing_prefetch_metadata() -> None: None, None, parquet_options, + 0, [], context=context, ) diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 0cd616e9789f..17668bee6061 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -325,6 +325,7 @@ def test_validate_cluster() -> None: "broadcast_limit", "sink_to_directory", "client_device_threshold", + "num_prefetch_workers", "max_concurrent_io_tasks", "num_py_executors", ],