From 1615e0d389f6ab696354dc2ff280b8af1d0a6bd8 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 2 Sep 2026 11:02:29 -0700 Subject: [PATCH 01/31] introduce ParquetScanTask --- .../cudf_polars/cudf_polars/dsl/utils/io.py | 35 ++- .../cudf_polars/streaming/actor_graph/io.py | 9 +- .../cudf_polars/cudf_polars/streaming/io.py | 204 ++++++++++++++---- .../cudf_polars/tests/streaming/test_scan.py | 50 +++++ 4 files changed, 243 insertions(+), 55 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index 254f46ed65e2..d86db5d533d9 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -15,11 +15,17 @@ from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars from cudf_polars.dsl.traversal import traversal -from cudf_polars.streaming.io import Scan, StreamingScan +from cudf_polars.streaming.io import ( + ParquetScanTask, + ParquetSourceInfo, + Scan, + StreamingScan, +) if TYPE_CHECKING: from cudf_polars.dsl.ir import IR from cudf_polars.streaming.base import StatsCollector + from cudf_polars.streaming.io import StreamingScanTask @dataclass(frozen=True) @@ -181,8 +187,6 @@ def prefetch_parquet_file_metadata_for_ir( ------- A dictionary mapping each individual path to its cached parquet metadata. """ - from cudf_polars.streaming.io import ParquetSourceInfo, StreamingScan - all_paths: set[str] = set() for node in traversal([root]): @@ -254,10 +258,25 @@ def attach_cached_parquet_metadata( """ for node in traversal([root]): if isinstance(node, StreamingScan) and node.base_scan.typ == "parquet": + scans: list[StreamingScanTask] = [] + changed = False for scan in node.scans: - if not all(path in cached_parquet_info_map for path in scan.paths): + generic_scan = scan.scan if isinstance(scan, ParquetScanTask) else scan + if not all( + path in cached_parquet_info_map for path in generic_scan.paths + ): + scans.append(scan) continue - cached = [cached_parquet_info_map[path] for path in scan.paths] - Scan._validate_cached_parquet_info(scan.paths, cached) - scan.cached_parquet_info = cached - scan._non_child_args = (*scan._non_child_args[:-1], cached) + cached = [cached_parquet_info_map[path] for path in generic_scan.paths] + Scan._validate_cached_parquet_info(generic_scan.paths, cached) + generic_scan.cached_parquet_info = cached + generic_scan._non_child_args = ( + *generic_scan._non_child_args[:-1], + cached, + ) + task = ParquetScanTask.from_scan(generic_scan) + scans.append(task or generic_scan) + changed = changed or task is not None or scan is not generic_scan + if changed: + node.scans = scans + node._non_child_args = (node.scans, node.base_scan, node.scan_type) 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 9145c688da85..0299965db5c0 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -60,7 +60,7 @@ IOPartitionPlan, PartitionInfo, ) - from cudf_polars.streaming.io import FusedScan, SplitScan + from cudf_polars.streaming.io import StreamingScanTask from cudf_polars.utils.config import MaxConcurrentIOTasks @@ -613,7 +613,7 @@ async def scan_node( 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 + scans: Sequence[StreamingScanTask] = ir.scans async with shutdown_on_error( context, ch_out, trace_ir=ir, ir_context=ir_context @@ -651,13 +651,12 @@ async def scan_node( lineariser = Lineariser(context, ch_out, num_producers) # Assign tasks to producers using round-robin - producer_tasks: list[list[tuple[int, SplitScan | FusedScan]]] = [ + producer_tasks: list[list[tuple[int, StreamingScanTask]]] = [ [] 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] + producer_tasks[producer_id].append((task_idx, scan)) async def _producer(producer_id: int, ch_out: Channel) -> None: for task_idx, scan in producer_tasks[producer_id]: diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 3947b19eedff..0019bb78294c 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -11,7 +11,7 @@ import statistics from collections import defaultdict from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal, Self, overload +from typing import TYPE_CHECKING, Any, Literal, Self, TypeAlias, overload import polars as pl @@ -340,6 +340,25 @@ def _read_with_hybrid_scan( return DataFrame(columns, stream=stream).select(list(schema.keys())) +def _split_scan_row_groups(scan: SplitScan) -> list[list[int]] | None: + """Return the complete row groups read by a row-group-aligned split.""" + if scan.cached_parquet_info is None or len(scan.cached_parquet_info) != 1: + return None + + total_row_groups = len(scan.cached_parquet_info[0].file_metadata.row_group_num_rows) + if scan.total_splits > total_row_groups: + return None + + row_group_stride = total_row_groups // scan.total_splits + start = row_group_stride * scan.split_index + stop = ( + total_row_groups + if scan.split_index == scan.total_splits - 1 + else start + row_group_stride + ) + return [list(range(start, stop))] if start < stop else None + + class SplitScan(IR): """ Input from a split file. @@ -495,44 +514,6 @@ 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 hybrid_scan_eligible( - parquet_options, - cached_parquet_info=cached_parquet_info, - row_index=row_index, - include_file_paths=include_file_paths, - predicate=predicate, - ): - assert predicate is not None - assert cached_parquet_info is not None - stream = context.get_cuda_stream() - plc_filter, residual = to_parquet_filter( - _prepare_parquet_predicate( - predicate.value, paths, schema, with_columns - ), - stream=stream, - ) - if plc_filter is not None and residual is 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 @@ -674,6 +655,145 @@ def do_evaluate( ) +class ParquetScanTask(IR): + """Parquet scan task that reads complete row groups.""" + + __slots__ = ( + "paths", + "row_groups", + "scan", + "schema", + ) + _non_child = ("scan", "row_groups") + _n_non_child_args = 2 + + scan: FusedScan | SplitScan + """Underlying generic scan task.""" + paths: list[str] + """File paths assigned to this task.""" + row_groups: list[list[int]] + """Row-group indices to read from each input path.""" + + def __init__(self, scan: FusedScan | SplitScan, row_groups: list[list[int]]): + if scan.base_scan.typ != "parquet": # pragma: no cover + raise ValueError(f"Expected a parquet scan task, got: {scan.base_scan.typ}") + if len(scan.paths) != len(row_groups): # pragma: no cover + raise ValueError("Expected one row-group list for each input path.") + + self.scan = scan + self.paths = scan.paths + self.row_groups = row_groups + self.schema = scan.schema + self._non_child_args = (scan, row_groups) + self.children = () + + @property + def base_scan(self) -> Scan: + """Scan operation this node is based on.""" + return self.scan.base_scan + + @property + def parquet_options(self) -> ParquetOptions: + """Parquet-specific options.""" + return self.scan.parquet_options + + @property + def cached_parquet_info(self) -> list[CachedParquetInfo] | None: + """Cached parquet metadata.""" + return self.scan.cached_parquet_info + + @cached_parquet_info.setter + def cached_parquet_info(self, value: list[CachedParquetInfo] | None) -> None: + self.scan.cached_parquet_info = value + self.scan._non_child_args = (*self.scan._non_child_args[:-1], value) + + @classmethod + def from_scan(cls, scan: FusedScan | SplitScan) -> Self | None: + """Create a parquet row-group task for scans that are fully row-group based.""" + if scan.base_scan.typ != "parquet" or scan.cached_parquet_info is None: + return None + + if isinstance(scan, SplitScan): + row_groups = _split_scan_row_groups(scan) + if row_groups is None: + return None + else: + row_groups = [ + list(range(len(info.file_metadata.row_group_num_rows))) + for info in scan.cached_parquet_info + ] + + return cls(scan, row_groups) + + def get_hashable(self) -> Hashable: + """Hashable representation of the node.""" + return ( + type(self), + self.scan.get_hashable(), + tuple(tuple(row_groups) for row_groups in self.row_groups), + ) + + @classmethod + def do_evaluate( + cls, + scan: FusedScan | SplitScan, + row_groups: list[list[int]], + *, + context: IRExecutionContext, + ) -> DataFrame: + """Evaluate a parquet row-group task.""" + base_scan = scan.base_scan + cached_parquet_info = scan.cached_parquet_info + + # 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 ( + len(scan.paths) == 1 + and len(row_groups) == 1 + and hybrid_scan_eligible( + scan.parquet_options, + cached_parquet_info=cached_parquet_info, + row_index=base_scan.row_index, + include_file_paths=base_scan.include_file_paths, + predicate=base_scan.predicate, + ) + ): + assert base_scan.predicate is not None + assert cached_parquet_info is not None + stream = context.get_cuda_stream() + plc_filter, residual = to_parquet_filter( + _prepare_parquet_predicate( + base_scan.predicate.value, + scan.paths, + scan.schema, + base_scan.with_columns, + ), + stream=stream, + ) + if plc_filter is not None and residual is None: + split_index = scan.split_index if isinstance(scan, SplitScan) else 0 + total_splits = scan.total_splits if isinstance(scan, SplitScan) else 1 + return _read_with_hybrid_scan( + scan.schema, + scan.paths, + base_scan.with_columns, + plc_filter, + row_groups[0], + stream, + cached_parquet_info[0], + split_index=split_index, + total_splits=total_splits, + stats_pruning=scan.parquet_options._hybrid_scan_stats_pruning, + ) + + return scan.do_evaluate(*scan._non_child_args, context=context) + + +StreamingScanTask: TypeAlias = SplitScan | FusedScan | ParquetScanTask + + @lower_ir_node.register(Empty) def _( ir: Empty, rec: LowerIRTransformer @@ -753,12 +873,12 @@ class StreamingScan(IR): "scan_type", ) _n_non_child_args = 3 - scans: Sequence[SplitScan] | Sequence[FusedScan] + scans: Sequence[StreamingScanTask] base_scan: Scan def __init__( self, - scans: Sequence[SplitScan] | Sequence[FusedScan], + scans: Sequence[StreamingScanTask], base_scan: Scan, scan_type: Literal["split", "fused"], ): @@ -842,7 +962,7 @@ def get_hashable(self) -> Hashable: @classmethod def do_evaluate( cls, - scans: Sequence[SplitScan] | Sequence[FusedScan], + scans: Sequence[StreamingScanTask], base_scan: Scan, *, context: IRExecutionContext, diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 72389d58e9fd..d0512fa35923 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -20,6 +20,7 @@ from cudf_polars.dsl.utils.io import ( CachedParquetInfo, _prefetch_parquet_footers_for_paths, + attach_cached_parquet_metadata, prefetch_parquet_file_metadata_for_ir, ) from cudf_polars.engine.options import StreamingOptions @@ -32,6 +33,7 @@ ) from cudf_polars.streaming.io import ( FusedScan, + ParquetScanTask, SplitScan, StreamingScan, expand_scan_for_rank, @@ -496,6 +498,54 @@ def test_expand_scan_for_rank_split_files( assert scan.paths == ["file.parquet"] +def test_attach_cached_parquet_metadata_creates_parquet_scan_tasks( + tmp_path: Path, +) -> None: + source = tmp_path / "data.parquet" + pl.DataFrame({"x": range(4)}).write_parquet(source, row_group_size=2) + + base = _make_parquet_scan([str(source)]) + streaming_scan = StreamingScan.for_split_files( + base, + IOPartitionPlan(2, IOPartitionFlavor.SPLIT_FILES), + partition_count=2, + rank=0, + nranks=1, + parquet_options=base.parquet_options, + ) + + cached = prefetch_parquet_file_metadata_for_ir(streaming_scan, None) + attach_cached_parquet_metadata(streaming_scan, cached) + + parquet_tasks = [ + scan for scan in streaming_scan.scans if isinstance(scan, ParquetScanTask) + ] + assert len(parquet_tasks) == len(streaming_scan.scans) + assert [scan.row_groups for scan in parquet_tasks] == [[[0]], [[1]]] + + +def test_attach_cached_parquet_metadata_keeps_sub_row_group_split_scan( + tmp_path: Path, +) -> None: + source = tmp_path / "data.parquet" + pl.DataFrame({"x": range(4)}).write_parquet(source, row_group_size=2) + + base = _make_parquet_scan([str(source)]) + streaming_scan = StreamingScan.for_split_files( + base, + IOPartitionPlan(4, IOPartitionFlavor.SPLIT_FILES), + partition_count=4, + rank=0, + nranks=1, + parquet_options=base.parquet_options, + ) + + cached = prefetch_parquet_file_metadata_for_ir(streaming_scan, None) + attach_cached_parquet_metadata(streaming_scan, cached) + + assert all(isinstance(scan, SplitScan) for scan in streaming_scan.scans) + + def test_streaming_scan_raises() -> None: # This isn't reachable by normal cudf-polars usage. scan = _make_parquet_scan(["file.parquet"]) From d3699efa4b35f7c154a48bb9ab04374cce985e7c Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 2 Sep 2026 11:55:44 -0700 Subject: [PATCH 02/31] variation 2 --- python/cudf_polars/cudf_polars/dsl/ir.py | 90 +++++++++-- .../cudf_polars/cudf_polars/dsl/utils/io.py | 50 +++--- .../cudf_polars/cudf_polars/streaming/io.py | 146 +++++++++++------- 3 files changed, 188 insertions(+), 98 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 024102296210..1d9f89e1986c 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -921,26 +921,36 @@ def _get_parquet_row_count_from_metadata( n_rows: int, parquet_options: ParquetOptions, cached_parquet_info: list[CachedParquetInfo] | None, + row_groups: list[list[int]] | None = None, ) -> int: # Zero-width parquet files lose their row count when read through # pylibcudf. See https://github.com/NVIDIA/cudf/issues/21428 if cached_parquet_info is not None: Scan._validate_cached_parquet_info(paths, cached_parquet_info) - parquet_metadatas = [ - info.file_metadata for info in cached_parquet_info - ] # pragma: no cover - num_rows = sum( - metadata.num_rows for metadata in parquet_metadatas - ) # pragma: no cover + if row_groups is None: + num_rows = sum( + info.file_metadata.num_rows for info in cached_parquet_info + ) # pragma: no cover + else: + num_rows = sum( + info.file_metadata.row_group_num_rows[row_group] + for info, file_row_groups in zip( + cached_parquet_info, row_groups, strict=True + ) + for row_group in file_row_groups + ) else: + if row_groups is not None: # pragma: no cover + raise ValueError("Explicit parquet row groups require cached metadata.") meta = plc.io.parquet_metadata.read_parquet_metadata( plc.io.SourceInfo(paths) ) num_rows = meta.num_rows() - num_rows -= skip_rows - if n_rows != -1: - num_rows = min(num_rows, n_rows) + if row_groups is None: + num_rows -= skip_rows + if n_rows != -1: + num_rows = min(num_rows, n_rows) return max(num_rows, 0) @staticmethod @@ -957,9 +967,7 @@ def _apply_parquet_projection( return plc.Table([columns[name] for name in with_columns]), with_columns @classmethod - @log_do_evaluate - @nvtx_annotate_cudf_polars(message="Scan") - def do_evaluate( + def _do_evaluate( cls, schema: Schema, typ: str, @@ -973,6 +981,7 @@ def do_evaluate( predicate: expr.NamedExpr | None, parquet_options: ParquetOptions, cached_parquet_info: list[CachedParquetInfo] | None, + row_groups: list[list[int]] | None = None, *, context: IRExecutionContext, ) -> DataFrame: @@ -1089,6 +1098,13 @@ def read_csv_header( df, ) elif typ == "parquet": + if row_groups is not None and ( + skip_rows != 0 or n_rows != -1 or row_index is not None + ): # pragma: no cover + raise ValueError( + "Explicit parquet row groups cannot be combined with row " + "slicing or row indexes." + ) if cached_parquet_info is not None: Scan._validate_cached_parquet_info(paths, cached_parquet_info) filepath_sources = [] @@ -1129,10 +1145,13 @@ def read_csv_header( parquet_reader_options.set_column_names(with_columns) if filters is not None: parquet_reader_options.set_filter(filters) - if n_rows != -1: - parquet_reader_options.set_num_rows(n_rows) - if skip_rows != 0: - parquet_reader_options.set_skip_rows(skip_rows) + if row_groups is not None: + parquet_reader_options.set_row_groups(row_groups) + else: + if n_rows != -1: + parquet_reader_options.set_num_rows(n_rows) + if skip_rows != 0: + parquet_reader_options.set_skip_rows(skip_rows) if parquet_options.chunked: reader = plc.io.parquet.ChunkedParquetReader( parquet_reader_options, @@ -1165,6 +1184,7 @@ def read_csv_header( n_rows, parquet_options, cached_parquet_info, + row_groups, ), ) df = DataFrame.from_table( @@ -1197,6 +1217,7 @@ def read_csv_header( n_rows, parquet_options, cached_parquet_info, + row_groups, ), ) df = DataFrame.from_table( @@ -1262,6 +1283,43 @@ def read_csv_header( ) return apply_predicate(df, effective_predicate) + @classmethod + @log_do_evaluate + @nvtx_annotate_cudf_polars(message="Scan") + def do_evaluate( + cls, + schema: Schema, + typ: str, + reader_options: dict[str, Any], + paths: list[str], + with_columns: list[str] | None, + skip_rows: int, + n_rows: int, + row_index: tuple[str, int] | None, + include_file_paths: str | None, + predicate: expr.NamedExpr | None, + parquet_options: ParquetOptions, + cached_parquet_info: list[CachedParquetInfo] | None, + *, + context: IRExecutionContext, + ) -> DataFrame: + """Evaluate and return a dataframe.""" + return cls._do_evaluate( + schema, + typ, + reader_options, + paths, + with_columns, + skip_rows, + n_rows, + row_index, + include_file_paths, + predicate, + parquet_options, + cached_parquet_info, + context=context, + ) + class Sink(IR): """Sink a dataframe to a file.""" diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index d86db5d533d9..a68c3fa6eab5 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -240,6 +240,23 @@ def prefetch_parquet_file_metadata_for_ir( return cached_parquet_info +def _attach_cached_metadata_to_scan_task( + scan: StreamingScanTask, + cached_parquet_info_map: dict[str, CachedParquetInfo], +) -> StreamingScanTask: + """Attach cached parquet metadata and specialize row-group-aligned tasks.""" + if isinstance(scan, ParquetScanTask): + return scan + if not all(path in cached_parquet_info_map for path in scan.paths): + return scan + + cached = [cached_parquet_info_map[path] for path in scan.paths] + Scan._validate_cached_parquet_info(scan.paths, cached) + scan.cached_parquet_info = cached + scan._non_child_args = (*scan._non_child_args[:-1], cached) + return ParquetScanTask.from_scan(scan) or scan + + def attach_cached_parquet_metadata( root: IR, cached_parquet_info_map: dict[str, CachedParquetInfo], @@ -257,26 +274,13 @@ def attach_cached_parquet_metadata( Mapping from file paths to cached parquet metadata. """ for node in traversal([root]): - if isinstance(node, StreamingScan) and node.base_scan.typ == "parquet": - scans: list[StreamingScanTask] = [] - changed = False - for scan in node.scans: - generic_scan = scan.scan if isinstance(scan, ParquetScanTask) else scan - if not all( - path in cached_parquet_info_map for path in generic_scan.paths - ): - scans.append(scan) - continue - cached = [cached_parquet_info_map[path] for path in generic_scan.paths] - Scan._validate_cached_parquet_info(generic_scan.paths, cached) - generic_scan.cached_parquet_info = cached - generic_scan._non_child_args = ( - *generic_scan._non_child_args[:-1], - cached, - ) - task = ParquetScanTask.from_scan(generic_scan) - scans.append(task or generic_scan) - changed = changed or task is not None or scan is not generic_scan - if changed: - node.scans = scans - node._non_child_args = (node.scans, node.base_scan, node.scan_type) + if not isinstance(node, StreamingScan) or node.base_scan.typ != "parquet": + continue + + scans = [ + _attach_cached_metadata_to_scan_task(scan, cached_parquet_info_map) + for scan in node.scans + ] + if any(scan is not old for scan, old in zip(scans, node.scans, strict=True)): + node.scans = scans + node._non_child_args = (node.scans, node.base_scan, node.scan_type) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 0019bb78294c..56b75d8c88ad 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -234,10 +234,8 @@ def _read_with_hybrid_scan( total_splits: int = 1, stats_pruning: bool = True, ) -> DataFrame: - """Two-pass parquet read via HybridScanReader for a row-group-aligned split.""" - assert len(paths) == 1, ( - "hybrid scan only supported for SplitScan; one physical file" - ) + """Two-pass parquet read via HybridScanReader for a row-group-aligned task.""" + assert len(paths) == 1, "hybrid scan only supports one physical file" with nvtx_annotate_cudf_polars( message="HybridScan", payload=(split_index + 1, total_splits) ): @@ -342,10 +340,11 @@ def _read_with_hybrid_scan( def _split_scan_row_groups(scan: SplitScan) -> list[list[int]] | None: """Return the complete row groups read by a row-group-aligned split.""" - if scan.cached_parquet_info is None or len(scan.cached_parquet_info) != 1: + cached_parquet_info = scan.cached_parquet_info + if cached_parquet_info is None or len(cached_parquet_info) != 1: return None - total_row_groups = len(scan.cached_parquet_info[0].file_metadata.row_group_num_rows) + total_row_groups = len(cached_parquet_info[0].file_metadata.row_group_num_rows) if scan.total_splits > total_row_groups: return None @@ -659,58 +658,67 @@ class ParquetScanTask(IR): """Parquet scan task that reads complete row groups.""" __slots__ = ( + "base_scan", + "cached_parquet_info", + "parquet_options", "paths", "row_groups", - "scan", "schema", ) - _non_child = ("scan", "row_groups") - _n_non_child_args = 2 + _non_child = ( + "base_scan", + "paths", + "row_groups", + "parquet_options", + "cached_parquet_info", + ) + _n_non_child_args = 5 - scan: FusedScan | SplitScan - """Underlying generic scan task.""" + base_scan: Scan + """Scan operation this task is based on.""" paths: list[str] """File paths assigned to this task.""" row_groups: list[list[int]] """Row-group indices to read from each input path.""" + parquet_options: ParquetOptions + """Parquet-specific options.""" + cached_parquet_info: list[CachedParquetInfo] + """Cached parquet metadata.""" - def __init__(self, scan: FusedScan | SplitScan, row_groups: list[list[int]]): - if scan.base_scan.typ != "parquet": # pragma: no cover - raise ValueError(f"Expected a parquet scan task, got: {scan.base_scan.typ}") - if len(scan.paths) != len(row_groups): # pragma: no cover + def __init__( + self, + base_scan: Scan, + paths: list[str], + row_groups: list[list[int]], + parquet_options: ParquetOptions, + cached_parquet_info: list[CachedParquetInfo], + ): + if base_scan.typ != "parquet": # pragma: no cover + raise ValueError(f"Expected a parquet scan task, got: {base_scan.typ}") + if len(paths) != len(row_groups): # pragma: no cover raise ValueError("Expected one row-group list for each input path.") - self.scan = scan - self.paths = scan.paths + Scan._validate_cached_parquet_info(paths, cached_parquet_info) + self.base_scan = base_scan + self.paths = paths self.row_groups = row_groups - self.schema = scan.schema - self._non_child_args = (scan, row_groups) + self.parquet_options = parquet_options + self.cached_parquet_info = cached_parquet_info + self.schema = base_scan.schema + self._non_child_args = ( + base_scan, + paths, + row_groups, + parquet_options, + cached_parquet_info, + ) self.children = () - @property - def base_scan(self) -> Scan: - """Scan operation this node is based on.""" - return self.scan.base_scan - - @property - def parquet_options(self) -> ParquetOptions: - """Parquet-specific options.""" - return self.scan.parquet_options - - @property - def cached_parquet_info(self) -> list[CachedParquetInfo] | None: - """Cached parquet metadata.""" - return self.scan.cached_parquet_info - - @cached_parquet_info.setter - def cached_parquet_info(self, value: list[CachedParquetInfo] | None) -> None: - self.scan.cached_parquet_info = value - self.scan._non_child_args = (*self.scan._non_child_args[:-1], value) - @classmethod def from_scan(cls, scan: FusedScan | SplitScan) -> Self | None: """Create a parquet row-group task for scans that are fully row-group based.""" - if scan.base_scan.typ != "parquet" or scan.cached_parquet_info is None: + cached_parquet_info = scan.cached_parquet_info + if scan.base_scan.typ != "parquet" or cached_parquet_info is None: return None if isinstance(scan, SplitScan): @@ -720,40 +728,48 @@ def from_scan(cls, scan: FusedScan | SplitScan) -> Self | None: else: row_groups = [ list(range(len(info.file_metadata.row_group_num_rows))) - for info in scan.cached_parquet_info + for info in cached_parquet_info ] - return cls(scan, row_groups) + return cls( + scan.base_scan, + scan.paths, + row_groups, + scan.parquet_options, + cached_parquet_info, + ) def get_hashable(self) -> Hashable: """Hashable representation of the node.""" return ( type(self), - self.scan.get_hashable(), + self.base_scan.get_hashable(), + tuple(self.paths), + self.parquet_options, tuple(tuple(row_groups) for row_groups in self.row_groups), ) @classmethod def do_evaluate( cls, - scan: FusedScan | SplitScan, + base_scan: Scan, + paths: list[str], row_groups: list[list[int]], + parquet_options: ParquetOptions, + cached_parquet_info: list[CachedParquetInfo], *, context: IRExecutionContext, ) -> DataFrame: """Evaluate a parquet row-group task.""" - base_scan = scan.base_scan - cached_parquet_info = scan.cached_parquet_info - # 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 ( - len(scan.paths) == 1 + len(paths) == 1 and len(row_groups) == 1 and hybrid_scan_eligible( - scan.parquet_options, + parquet_options, cached_parquet_info=cached_parquet_info, row_index=base_scan.row_index, include_file_paths=base_scan.include_file_paths, @@ -766,29 +782,41 @@ def do_evaluate( plc_filter, residual = to_parquet_filter( _prepare_parquet_predicate( base_scan.predicate.value, - scan.paths, - scan.schema, + paths, + base_scan.schema, base_scan.with_columns, ), stream=stream, ) if plc_filter is not None and residual is None: - split_index = scan.split_index if isinstance(scan, SplitScan) else 0 - total_splits = scan.total_splits if isinstance(scan, SplitScan) else 1 return _read_with_hybrid_scan( - scan.schema, - scan.paths, + base_scan.schema, + paths, base_scan.with_columns, plc_filter, row_groups[0], stream, cached_parquet_info[0], - split_index=split_index, - total_splits=total_splits, - stats_pruning=scan.parquet_options._hybrid_scan_stats_pruning, + stats_pruning=parquet_options._hybrid_scan_stats_pruning, ) - return scan.do_evaluate(*scan._non_child_args, context=context) + with nvtx_annotate_cudf_polars(message=f"ParquetScan: {', '.join(paths)}"): + return Scan._do_evaluate( + base_scan.schema, + base_scan.typ, + base_scan.reader_options, + paths, + base_scan.with_columns, + base_scan.skip_rows, + base_scan.n_rows, + base_scan.row_index, + base_scan.include_file_paths, + base_scan.predicate, + parquet_options, + cached_parquet_info, + row_groups, + context=context, + ) StreamingScanTask: TypeAlias = SplitScan | FusedScan | ParquetScanTask From e3c2d087a53362e71ac81201e22cb9292c87f35d Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 2 Sep 2026 12:19:07 -0700 Subject: [PATCH 03/31] cleanup --- python/cudf_polars/cudf_polars/dsl/ir.py | 90 ++++--------------- .../cudf_polars/cudf_polars/streaming/io.py | 53 ++++++++++- 2 files changed, 65 insertions(+), 78 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 1d9f89e1986c..024102296210 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -921,36 +921,26 @@ def _get_parquet_row_count_from_metadata( n_rows: int, parquet_options: ParquetOptions, cached_parquet_info: list[CachedParquetInfo] | None, - row_groups: list[list[int]] | None = None, ) -> int: # Zero-width parquet files lose their row count when read through # pylibcudf. See https://github.com/NVIDIA/cudf/issues/21428 if cached_parquet_info is not None: Scan._validate_cached_parquet_info(paths, cached_parquet_info) - if row_groups is None: - num_rows = sum( - info.file_metadata.num_rows for info in cached_parquet_info - ) # pragma: no cover - else: - num_rows = sum( - info.file_metadata.row_group_num_rows[row_group] - for info, file_row_groups in zip( - cached_parquet_info, row_groups, strict=True - ) - for row_group in file_row_groups - ) + parquet_metadatas = [ + info.file_metadata for info in cached_parquet_info + ] # pragma: no cover + num_rows = sum( + metadata.num_rows for metadata in parquet_metadatas + ) # pragma: no cover else: - if row_groups is not None: # pragma: no cover - raise ValueError("Explicit parquet row groups require cached metadata.") meta = plc.io.parquet_metadata.read_parquet_metadata( plc.io.SourceInfo(paths) ) num_rows = meta.num_rows() - if row_groups is None: - num_rows -= skip_rows - if n_rows != -1: - num_rows = min(num_rows, n_rows) + num_rows -= skip_rows + if n_rows != -1: + num_rows = min(num_rows, n_rows) return max(num_rows, 0) @staticmethod @@ -967,7 +957,9 @@ def _apply_parquet_projection( return plc.Table([columns[name] for name in with_columns]), with_columns @classmethod - def _do_evaluate( + @log_do_evaluate + @nvtx_annotate_cudf_polars(message="Scan") + def do_evaluate( cls, schema: Schema, typ: str, @@ -981,7 +973,6 @@ def _do_evaluate( predicate: expr.NamedExpr | None, parquet_options: ParquetOptions, cached_parquet_info: list[CachedParquetInfo] | None, - row_groups: list[list[int]] | None = None, *, context: IRExecutionContext, ) -> DataFrame: @@ -1098,13 +1089,6 @@ def read_csv_header( df, ) elif typ == "parquet": - if row_groups is not None and ( - skip_rows != 0 or n_rows != -1 or row_index is not None - ): # pragma: no cover - raise ValueError( - "Explicit parquet row groups cannot be combined with row " - "slicing or row indexes." - ) if cached_parquet_info is not None: Scan._validate_cached_parquet_info(paths, cached_parquet_info) filepath_sources = [] @@ -1145,13 +1129,10 @@ def read_csv_header( parquet_reader_options.set_column_names(with_columns) if filters is not None: parquet_reader_options.set_filter(filters) - if row_groups is not None: - parquet_reader_options.set_row_groups(row_groups) - else: - if n_rows != -1: - parquet_reader_options.set_num_rows(n_rows) - if skip_rows != 0: - parquet_reader_options.set_skip_rows(skip_rows) + if n_rows != -1: + parquet_reader_options.set_num_rows(n_rows) + if skip_rows != 0: + parquet_reader_options.set_skip_rows(skip_rows) if parquet_options.chunked: reader = plc.io.parquet.ChunkedParquetReader( parquet_reader_options, @@ -1184,7 +1165,6 @@ def read_csv_header( n_rows, parquet_options, cached_parquet_info, - row_groups, ), ) df = DataFrame.from_table( @@ -1217,7 +1197,6 @@ def read_csv_header( n_rows, parquet_options, cached_parquet_info, - row_groups, ), ) df = DataFrame.from_table( @@ -1283,43 +1262,6 @@ def read_csv_header( ) return apply_predicate(df, effective_predicate) - @classmethod - @log_do_evaluate - @nvtx_annotate_cudf_polars(message="Scan") - def do_evaluate( - cls, - schema: Schema, - typ: str, - reader_options: dict[str, Any], - paths: list[str], - with_columns: list[str] | None, - skip_rows: int, - n_rows: int, - row_index: tuple[str, int] | None, - include_file_paths: str | None, - predicate: expr.NamedExpr | None, - parquet_options: ParquetOptions, - cached_parquet_info: list[CachedParquetInfo] | None, - *, - context: IRExecutionContext, - ) -> DataFrame: - """Evaluate and return a dataframe.""" - return cls._do_evaluate( - schema, - typ, - reader_options, - paths, - with_columns, - skip_rows, - n_rows, - row_index, - include_file_paths, - predicate, - parquet_options, - cached_parquet_info, - context=context, - ) - class Sink(IR): """Sink a dataframe to a file.""" diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 56b75d8c88ad..4ae29e16dbf5 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -338,6 +338,40 @@ def _read_with_hybrid_scan( return DataFrame(columns, stream=stream).select(list(schema.keys())) +def _row_slice_from_row_groups( + cached_parquet_info: list[CachedParquetInfo], row_groups: list[list[int]] +) -> tuple[int, int]: + """Translate a contiguous row-group selection to skip_rows and n_rows.""" + skip_rows = n_rows = total_rows = 0 + selected_started = selected_stopped = False + + for info, file_row_groups in zip(cached_parquet_info, row_groups, strict=True): + file_row_counts = info.file_metadata.row_group_num_rows + if file_row_groups != sorted(set(file_row_groups)): # pragma: no cover + raise ValueError("Expected row groups in on-disk order without duplicates.") + if file_row_groups and ( + file_row_groups[0] < 0 or file_row_groups[-1] >= len(file_row_counts) + ): # pragma: no cover + raise ValueError("Row group index out of range.") + + selected = set(file_row_groups) + for row_group, rows in enumerate(file_row_counts): + total_rows += rows + if row_group in selected: + if selected_stopped: # pragma: no cover + raise ValueError("Expected one contiguous row-group selection.") + selected_started = True + n_rows += rows + elif selected_started: + selected_stopped = True + else: + skip_rows += rows + + if selected_started and skip_rows + n_rows == total_rows: + return skip_rows, -1 + return skip_rows, n_rows + + def _split_scan_row_groups(scan: SplitScan) -> list[list[int]] | None: """Return the complete row groups read by a row-group-aligned split.""" cached_parquet_info = scan.cached_parquet_info @@ -801,20 +835,31 @@ def do_evaluate( ) with nvtx_annotate_cudf_polars(message=f"ParquetScan: {', '.join(paths)}"): - return Scan._do_evaluate( + if ( + base_scan.skip_rows != 0 + or base_scan.n_rows != -1 + or base_scan.row_index is not None + ): # pragma: no cover + raise ValueError( + "Parquet row-group tasks cannot be combined with row slicing " + "or row indexes." + ) + skip_rows, n_rows = _row_slice_from_row_groups( + cached_parquet_info, row_groups + ) + return Scan.do_evaluate( base_scan.schema, base_scan.typ, base_scan.reader_options, paths, base_scan.with_columns, - base_scan.skip_rows, - base_scan.n_rows, + skip_rows, + n_rows, base_scan.row_index, base_scan.include_file_paths, base_scan.predicate, parquet_options, cached_parquet_info, - row_groups, context=context, ) From 0ebcc5c5b3b29dbf6cfd034016b38efe7bddb54a Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 2 Sep 2026 13:30:30 -0700 Subject: [PATCH 04/31] minor cleanup --- .../cudf_polars/cudf_polars/dsl/utils/io.py | 49 ++++----- .../cudf_polars/cudf_polars/streaming/io.py | 101 +++++++++--------- 2 files changed, 69 insertions(+), 81 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index a68c3fa6eab5..75049daf3b4d 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -25,7 +25,6 @@ if TYPE_CHECKING: from cudf_polars.dsl.ir import IR from cudf_polars.streaming.base import StatsCollector - from cudf_polars.streaming.io import StreamingScanTask @dataclass(frozen=True) @@ -240,23 +239,6 @@ def prefetch_parquet_file_metadata_for_ir( return cached_parquet_info -def _attach_cached_metadata_to_scan_task( - scan: StreamingScanTask, - cached_parquet_info_map: dict[str, CachedParquetInfo], -) -> StreamingScanTask: - """Attach cached parquet metadata and specialize row-group-aligned tasks.""" - if isinstance(scan, ParquetScanTask): - return scan - if not all(path in cached_parquet_info_map for path in scan.paths): - return scan - - cached = [cached_parquet_info_map[path] for path in scan.paths] - Scan._validate_cached_parquet_info(scan.paths, cached) - scan.cached_parquet_info = cached - scan._non_child_args = (*scan._non_child_args[:-1], cached) - return ParquetScanTask.from_scan(scan) or scan - - def attach_cached_parquet_metadata( root: IR, cached_parquet_info_map: dict[str, CachedParquetInfo], @@ -274,13 +256,24 @@ def attach_cached_parquet_metadata( Mapping from file paths to cached parquet metadata. """ for node in traversal([root]): - if not isinstance(node, StreamingScan) or node.base_scan.typ != "parquet": - continue - - scans = [ - _attach_cached_metadata_to_scan_task(scan, cached_parquet_info_map) - for scan in node.scans - ] - if any(scan is not old for scan, old in zip(scans, node.scans, strict=True)): - node.scans = scans - node._non_child_args = (node.scans, node.base_scan, node.scan_type) + if isinstance(node, StreamingScan) and node.base_scan.typ == "parquet": + scans = list(node.scans) + converted = False + for i, scan in enumerate(scans): + if isinstance(scan, ParquetScanTask) or not all( + path in cached_parquet_info_map for path in scan.paths + ): + continue + + cached = [cached_parquet_info_map[path] for path in scan.paths] + Scan._validate_cached_parquet_info(scan.paths, cached) + scan.cached_parquet_info = cached + scan._non_child_args = (*scan._non_child_args[:-1], cached) + + if (parquet_scan := ParquetScanTask.from_scan(scan)) is not None: + scans[i] = parquet_scan + converted = True + + if converted: + node.scans = scans + node._non_child_args = (node.scans, node.base_scan, node.scan_type) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 4ae29e16dbf5..3e97ab3bb87e 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -338,47 +338,16 @@ def _read_with_hybrid_scan( return DataFrame(columns, stream=stream).select(list(schema.keys())) -def _row_slice_from_row_groups( - cached_parquet_info: list[CachedParquetInfo], row_groups: list[list[int]] -) -> tuple[int, int]: - """Translate a contiguous row-group selection to skip_rows and n_rows.""" - skip_rows = n_rows = total_rows = 0 - selected_started = selected_stopped = False - - for info, file_row_groups in zip(cached_parquet_info, row_groups, strict=True): - file_row_counts = info.file_metadata.row_group_num_rows - if file_row_groups != sorted(set(file_row_groups)): # pragma: no cover - raise ValueError("Expected row groups in on-disk order without duplicates.") - if file_row_groups and ( - file_row_groups[0] < 0 or file_row_groups[-1] >= len(file_row_counts) - ): # pragma: no cover - raise ValueError("Row group index out of range.") - - selected = set(file_row_groups) - for row_group, rows in enumerate(file_row_counts): - total_rows += rows - if row_group in selected: - if selected_stopped: # pragma: no cover - raise ValueError("Expected one contiguous row-group selection.") - selected_started = True - n_rows += rows - elif selected_started: - selected_stopped = True - else: - skip_rows += rows - - if selected_started and skip_rows + n_rows == total_rows: - return skip_rows, -1 - return skip_rows, n_rows - - -def _split_scan_row_groups(scan: SplitScan) -> list[list[int]] | None: - """Return the complete row groups read by a row-group-aligned split.""" +def _split_scan_row_groups_and_slice( + scan: SplitScan, +) -> tuple[list[list[int]], int, int] | None: + """Return the row groups and row slice for a row-group-aligned split.""" cached_parquet_info = scan.cached_parquet_info if cached_parquet_info is None or len(cached_parquet_info) != 1: return None - total_row_groups = len(cached_parquet_info[0].file_metadata.row_group_num_rows) + row_group_num_rows = cached_parquet_info[0].file_metadata.row_group_num_rows + total_row_groups = len(row_group_num_rows) if scan.total_splits > total_row_groups: return None @@ -389,7 +358,13 @@ def _split_scan_row_groups(scan: SplitScan) -> list[list[int]] | None: if scan.split_index == scan.total_splits - 1 else start + row_group_stride ) - return [list(range(start, stop))] if start < stop else None + row_groups = list(range(start, stop)) + if not row_groups: + return None + + skip_rows = sum(row_group_num_rows[:start]) + n_rows = -1 if stop == total_row_groups else sum(row_group_num_rows[start:stop]) + return [row_groups], skip_rows, n_rows class SplitScan(IR): @@ -694,19 +669,23 @@ class ParquetScanTask(IR): __slots__ = ( "base_scan", "cached_parquet_info", + "n_rows", "parquet_options", "paths", "row_groups", "schema", + "skip_rows", ) _non_child = ( "base_scan", "paths", "row_groups", + "skip_rows", + "n_rows", "parquet_options", "cached_parquet_info", ) - _n_non_child_args = 5 + _n_non_child_args = 7 base_scan: Scan """Scan operation this task is based on.""" @@ -714,6 +693,10 @@ class ParquetScanTask(IR): """File paths assigned to this task.""" row_groups: list[list[int]] """Row-group indices to read from each input path.""" + skip_rows: int + """Number of rows to skip for the fallback parquet read.""" + n_rows: int + """Number of rows to read for the fallback parquet read.""" parquet_options: ParquetOptions """Parquet-specific options.""" cached_parquet_info: list[CachedParquetInfo] @@ -724,6 +707,8 @@ def __init__( base_scan: Scan, paths: list[str], row_groups: list[list[int]], + skip_rows: int, + n_rows: int, parquet_options: ParquetOptions, cached_parquet_info: list[CachedParquetInfo], ): @@ -731,11 +716,22 @@ def __init__( raise ValueError(f"Expected a parquet scan task, got: {base_scan.typ}") if len(paths) != len(row_groups): # pragma: no cover raise ValueError("Expected one row-group list for each input path.") + if ( + base_scan.skip_rows != 0 + or base_scan.n_rows != -1 + or base_scan.row_index is not None + ): # pragma: no cover + raise ValueError( + "Parquet row-group tasks cannot be combined with row slicing " + "or row indexes." + ) Scan._validate_cached_parquet_info(paths, cached_parquet_info) self.base_scan = base_scan self.paths = paths self.row_groups = row_groups + self.skip_rows = skip_rows + self.n_rows = n_rows self.parquet_options = parquet_options self.cached_parquet_info = cached_parquet_info self.schema = base_scan.schema @@ -743,6 +739,8 @@ def __init__( base_scan, paths, row_groups, + skip_rows, + n_rows, parquet_options, cached_parquet_info, ) @@ -756,19 +754,24 @@ def from_scan(cls, scan: FusedScan | SplitScan) -> Self | None: return None if isinstance(scan, SplitScan): - row_groups = _split_scan_row_groups(scan) - if row_groups is None: + result = _split_scan_row_groups_and_slice(scan) + if result is None: return None + row_groups, skip_rows, n_rows = result else: row_groups = [ list(range(len(info.file_metadata.row_group_num_rows))) for info in cached_parquet_info ] + skip_rows = 0 + n_rows = -1 return cls( scan.base_scan, scan.paths, row_groups, + skip_rows, + n_rows, scan.parquet_options, cached_parquet_info, ) @@ -781,6 +784,8 @@ def get_hashable(self) -> Hashable: tuple(self.paths), self.parquet_options, tuple(tuple(row_groups) for row_groups in self.row_groups), + self.skip_rows, + self.n_rows, ) @classmethod @@ -789,6 +794,8 @@ def do_evaluate( base_scan: Scan, paths: list[str], row_groups: list[list[int]], + skip_rows: int, + n_rows: int, parquet_options: ParquetOptions, cached_parquet_info: list[CachedParquetInfo], *, @@ -835,18 +842,6 @@ def do_evaluate( ) with nvtx_annotate_cudf_polars(message=f"ParquetScan: {', '.join(paths)}"): - if ( - base_scan.skip_rows != 0 - or base_scan.n_rows != -1 - or base_scan.row_index is not None - ): # pragma: no cover - raise ValueError( - "Parquet row-group tasks cannot be combined with row slicing " - "or row indexes." - ) - skip_rows, n_rows = _row_slice_from_row_groups( - cached_parquet_info, row_groups - ) return Scan.do_evaluate( base_scan.schema, base_scan.typ, From 18eb8afc589e3a4773bd6a73051cd9d325d78502 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 2 Sep 2026 13:39:02 -0700 Subject: [PATCH 05/31] update docstring --- python/cudf_polars/cudf_polars/dsl/utils/io.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index 75049daf3b4d..bf6c40d20b2b 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -244,9 +244,12 @@ def attach_cached_parquet_metadata( cached_parquet_info_map: dict[str, CachedParquetInfo], ) -> None: """ - Attach prefetched metadata to scan nodes. + Attach prefetched metadata to scan nodes and specialize parquet scan tasks. - This is an optimization only and does not affect IR identity. + This is an optimization only and does not affect IR identity. When cached + metadata is available, row-group-aligned parquet scan tasks are converted to + ``ParquetScanTask`` so downstream parquet-specific code can use the resolved + path and row-group assignment directly. Parameters ---------- From 4ce5e77f0d5af19749994a08a33afdb678c0a814 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 2 Sep 2026 14:53:49 -0700 Subject: [PATCH 06/31] address comment --- .../cudf_polars/cudf_polars/streaming/io.py | 11 ++++- .../cudf_polars/tests/streaming/test_scan.py | 44 +++++++++++++++++-- 2 files changed, 49 insertions(+), 6 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 3e97ab3bb87e..6f370324abb8 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -750,7 +750,14 @@ def __init__( def from_scan(cls, scan: FusedScan | SplitScan) -> Self | None: """Create a parquet row-group task for scans that are fully row-group based.""" cached_parquet_info = scan.cached_parquet_info - if scan.base_scan.typ != "parquet" or cached_parquet_info is None: + base_scan = scan.base_scan + if base_scan.typ != "parquet" or cached_parquet_info is None: + return None + if ( + base_scan.skip_rows != 0 + or base_scan.n_rows != -1 + or base_scan.row_index is not None + ): return None if isinstance(scan, SplitScan): @@ -767,7 +774,7 @@ def from_scan(cls, scan: FusedScan | SplitScan) -> Self | None: n_rows = -1 return cls( - scan.base_scan, + base_scan, scan.paths, row_groups, skip_rows, diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index d0512fa35923..20a452230bcb 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -396,7 +396,12 @@ def test_scan_union(engine: pl.GPUEngine, tmp_path: Path) -> None: def _make_parquet_scan( - paths: list[str], parquet_options: ParquetOptions | None = None + paths: list[str], + parquet_options: ParquetOptions | None = None, + *, + skip_rows: int = 0, + n_rows: int = -1, + row_index: tuple[str, int] | None = None, ) -> Scan: parquet_options = parquet_options or ParquetOptions() return Scan( @@ -406,9 +411,9 @@ def _make_parquet_scan( None, paths, None, - 0, - -1, - None, + skip_rows, + n_rows, + row_index, None, None, parquet_options, @@ -546,6 +551,37 @@ def test_attach_cached_parquet_metadata_keeps_sub_row_group_split_scan( assert all(isinstance(scan, SplitScan) for scan in streaming_scan.scans) +@pytest.mark.parametrize( + "skip_rows,n_rows,row_index", + [(1, -1, None), (0, 2, None), (0, -1, ("index", 0))], +) +def test_attach_cached_parquet_metadata_keeps_sliced_fused_scan( + tmp_path: Path, + skip_rows: int, + n_rows: int, + row_index: tuple[str, int] | None, +) -> None: + source = tmp_path / "data.parquet" + pl.DataFrame({"x": range(4)}).write_parquet(source, row_group_size=2) + + base = _make_parquet_scan( + [str(source)], skip_rows=skip_rows, n_rows=n_rows, row_index=row_index + ) + streaming_scan = StreamingScan.for_fused_files( + base, + IOPartitionPlan(1, IOPartitionFlavor.SINGLE_READ), + partition_count=1, + rank=0, + nranks=1, + parquet_options=base.parquet_options, + ) + + cached = prefetch_parquet_file_metadata_for_ir(streaming_scan, None) + attach_cached_parquet_metadata(streaming_scan, cached) + + assert all(isinstance(scan, FusedScan) for scan in streaming_scan.scans) + + def test_streaming_scan_raises() -> None: # This isn't reachable by normal cudf-polars usage. scan = _make_parquet_scan(["file.parquet"]) From ae2c72856f7b3f25fef03aeca6fee6f588e5f4ef Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 3 Sep 2026 06:13:39 -0700 Subject: [PATCH 07/31] rename class --- python/cudf_polars/cudf_polars/dsl/utils/io.py | 8 ++++---- python/cudf_polars/cudf_polars/streaming/io.py | 12 ++++++------ python/cudf_polars/tests/streaming/test_scan.py | 4 ++-- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index bf6c40d20b2b..739ee12afc16 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -16,7 +16,7 @@ from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars from cudf_polars.dsl.traversal import traversal from cudf_polars.streaming.io import ( - ParquetScanTask, + AlignedParquetScan, ParquetSourceInfo, Scan, StreamingScan, @@ -248,7 +248,7 @@ def attach_cached_parquet_metadata( This is an optimization only and does not affect IR identity. When cached metadata is available, row-group-aligned parquet scan tasks are converted to - ``ParquetScanTask`` so downstream parquet-specific code can use the resolved + ``AlignedParquetScan`` so downstream parquet-specific code can use the resolved path and row-group assignment directly. Parameters @@ -263,7 +263,7 @@ def attach_cached_parquet_metadata( scans = list(node.scans) converted = False for i, scan in enumerate(scans): - if isinstance(scan, ParquetScanTask) or not all( + if isinstance(scan, AlignedParquetScan) or not all( path in cached_parquet_info_map for path in scan.paths ): continue @@ -273,7 +273,7 @@ def attach_cached_parquet_metadata( scan.cached_parquet_info = cached scan._non_child_args = (*scan._non_child_args[:-1], cached) - if (parquet_scan := ParquetScanTask.from_scan(scan)) is not None: + if (parquet_scan := AlignedParquetScan.from_scan(scan)) is not None: scans[i] = parquet_scan converted = True diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index b7b055dd4117..c31c6ba3f1a2 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -668,8 +668,8 @@ def do_evaluate( ) -class ParquetScanTask(IR): - """Parquet scan task that reads complete row groups.""" +class AlignedParquetScan(IR): + """Parquet scan task aligned to complete row groups.""" __slots__ = ( "base_scan", @@ -718,7 +718,7 @@ def __init__( cached_parquet_info: list[CachedParquetInfo], ): if base_scan.typ != "parquet": # pragma: no cover - raise ValueError(f"Expected a parquet scan task, got: {base_scan.typ}") + raise ValueError(f"Expected a parquet scan, got: {base_scan.typ}") if len(paths) != len(row_groups): # pragma: no cover raise ValueError("Expected one row-group list for each input path.") if ( @@ -753,7 +753,7 @@ def __init__( @classmethod def from_scan(cls, scan: FusedScan | SplitScan) -> Self | None: - """Create a parquet row-group task for scans that are fully row-group based.""" + """Create an aligned parquet scan from a row-group-aligned scan task.""" cached_parquet_info = scan.cached_parquet_info base_scan = scan.base_scan if base_scan.typ != "parquet" or cached_parquet_info is None: @@ -813,7 +813,7 @@ def do_evaluate( *, context: IRExecutionContext, ) -> DataFrame: - """Evaluate a parquet row-group task.""" + """Evaluate a row-group-aligned parquet scan.""" # 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 @@ -871,7 +871,7 @@ def do_evaluate( ) -StreamingScanTask: TypeAlias = SplitScan | FusedScan | ParquetScanTask +StreamingScanTask: TypeAlias = SplitScan | FusedScan | AlignedParquetScan @lower_ir_node.register(Empty) diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 20a452230bcb..d7b3adfe0a85 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -32,8 +32,8 @@ StatsCollector, ) from cudf_polars.streaming.io import ( + AlignedParquetScan, FusedScan, - ParquetScanTask, SplitScan, StreamingScan, expand_scan_for_rank, @@ -523,7 +523,7 @@ def test_attach_cached_parquet_metadata_creates_parquet_scan_tasks( attach_cached_parquet_metadata(streaming_scan, cached) parquet_tasks = [ - scan for scan in streaming_scan.scans if isinstance(scan, ParquetScanTask) + scan for scan in streaming_scan.scans if isinstance(scan, AlignedParquetScan) ] assert len(parquet_tasks) == len(streaming_scan.scans) assert [scan.row_groups for scan in parquet_tasks] == [[[0]], [[1]]] From dcd5a0873394a9bc0ee658371b789daee924bfec Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 3 Sep 2026 07:35:28 -0700 Subject: [PATCH 08/31] experiment with new design --- .../cudf_polars/cudf_polars/dsl/utils/io.py | 27 +- .../cudf_polars/cudf_polars/streaming/io.py | 329 ++++++++++-------- .../cudf_polars/tests/streaming/test_scan.py | 40 ++- 3 files changed, 213 insertions(+), 183 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index 739ee12afc16..bde00c72cc63 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -16,7 +16,7 @@ from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars from cudf_polars.dsl.traversal import traversal from cudf_polars.streaming.io import ( - AlignedParquetScan, + ParquetScanTask, ParquetSourceInfo, Scan, StreamingScan, @@ -244,12 +244,9 @@ def attach_cached_parquet_metadata( cached_parquet_info_map: dict[str, CachedParquetInfo], ) -> None: """ - Attach prefetched metadata to scan nodes and specialize parquet scan tasks. + Attach prefetched metadata to parquet scan tasks. - This is an optimization only and does not affect IR identity. When cached - metadata is available, row-group-aligned parquet scan tasks are converted to - ``AlignedParquetScan`` so downstream parquet-specific code can use the resolved - path and row-group assignment directly. + This is an optimization only and does not affect IR identity. Parameters ---------- @@ -260,23 +257,11 @@ def attach_cached_parquet_metadata( """ for node in traversal([root]): if isinstance(node, StreamingScan) and node.base_scan.typ == "parquet": - scans = list(node.scans) - converted = False - for i, scan in enumerate(scans): - if isinstance(scan, AlignedParquetScan) or not all( + for scan in node.scans: + if not isinstance(scan, ParquetScanTask) or not all( path in cached_parquet_info_map for path in scan.paths ): continue cached = [cached_parquet_info_map[path] for path in scan.paths] - Scan._validate_cached_parquet_info(scan.paths, cached) - scan.cached_parquet_info = cached - scan._non_child_args = (*scan._non_child_args[:-1], cached) - - if (parquet_scan := AlignedParquetScan.from_scan(scan)) is not None: - scans[i] = parquet_scan - converted = True - - if converted: - node.scans = scans - node._non_child_args = (node.scans, node.base_scan, node.scan_type) + scan.attach_cached_parquet_info(cached) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index c31c6ba3f1a2..9974ae160e9f 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -345,10 +345,10 @@ def _read_with_hybrid_scan( def _split_scan_row_groups_and_slice( scan: SplitScan, + cached_parquet_info: list[CachedParquetInfo], ) -> tuple[list[list[int]], int, int] | None: """Return the row groups and row slice for a row-group-aligned split.""" - cached_parquet_info = scan.cached_parquet_info - if cached_parquet_info is None or len(cached_parquet_info) != 1: + if len(cached_parquet_info) != 1: return None row_group_num_rows = cached_parquet_info[0].file_metadata.row_group_num_rows @@ -668,95 +668,130 @@ def do_evaluate( ) -class AlignedParquetScan(IR): - """Parquet scan task aligned to complete row groups.""" +def _evaluate_scan_task( + scan: SplitScan | FusedScan, + cached_parquet_info: list[CachedParquetInfo] | None, + *, + context: IRExecutionContext, +) -> DataFrame: + """Evaluate a generic streaming scan task with optional parquet metadata.""" + base_scan = scan.base_scan + if isinstance(scan, SplitScan): + return SplitScan.do_evaluate( + scan.split_index, + scan.total_splits, + base_scan.schema, + base_scan.typ, + base_scan.reader_options, + scan.paths, + base_scan.with_columns, + base_scan.skip_rows, + base_scan.n_rows, + base_scan.row_index, + base_scan.include_file_paths, + base_scan.predicate, + scan.parquet_options, + cached_parquet_info, + context=context, + ) + return FusedScan.do_evaluate( + base_scan.schema, + base_scan.typ, + base_scan.reader_options, + scan.paths, + base_scan.with_columns, + base_scan.skip_rows, + base_scan.n_rows, + base_scan.row_index, + base_scan.include_file_paths, + base_scan.predicate, + scan.parquet_options, + cached_parquet_info, + context=context, + ) + + +def _set_scan_task_cached_parquet_info( + scan: SplitScan | FusedScan, + cached_parquet_info: list[CachedParquetInfo], +) -> None: + """Attach cached parquet metadata to a generic streaming scan task.""" + scan.cached_parquet_info = cached_parquet_info + scan._non_child_args = (*scan._non_child_args[:-1], cached_parquet_info) + + +class ParquetScanTask(IR): + """Parquet scan task wrapping a generic streaming scan task.""" __slots__ = ( "base_scan", + "base_task", "cached_parquet_info", - "n_rows", "parquet_options", "paths", - "row_groups", "schema", - "skip_rows", ) - _non_child = ( - "base_scan", - "paths", - "row_groups", - "skip_rows", - "n_rows", - "parquet_options", - "cached_parquet_info", - ) - _n_non_child_args = 7 + _non_child = ("base_task",) + _n_non_child_args = 2 + base_task: SplitScan | FusedScan + """Generic streaming scan task being specialized for parquet.""" base_scan: Scan """Scan operation this task is based on.""" paths: list[str] """File paths assigned to this task.""" - row_groups: list[list[int]] - """Row-group indices to read from each input path.""" - skip_rows: int - """Number of rows to skip for the fallback parquet read.""" - n_rows: int - """Number of rows to read for the fallback parquet read.""" parquet_options: ParquetOptions """Parquet-specific options.""" - cached_parquet_info: list[CachedParquetInfo] + cached_parquet_info: list[CachedParquetInfo] | None """Cached parquet metadata.""" def __init__( self, - base_scan: Scan, - paths: list[str], - row_groups: list[list[int]], - skip_rows: int, - n_rows: int, - parquet_options: ParquetOptions, - cached_parquet_info: list[CachedParquetInfo], + base_task: SplitScan | FusedScan, + cached_parquet_info: list[CachedParquetInfo] | None = None, ): + base_scan = base_task.base_scan if base_scan.typ != "parquet": # pragma: no cover raise ValueError(f"Expected a parquet scan, got: {base_scan.typ}") - if len(paths) != len(row_groups): # pragma: no cover - raise ValueError("Expected one row-group list for each input path.") - if ( - base_scan.skip_rows != 0 - or base_scan.n_rows != -1 - or base_scan.row_index is not None - ): # pragma: no cover - raise ValueError( - "Parquet row-group tasks cannot be combined with row slicing " - "or row indexes." - ) - - Scan._validate_cached_parquet_info(paths, cached_parquet_info) + if cached_parquet_info is None: + cached_parquet_info = base_task.cached_parquet_info + if cached_parquet_info == [] and base_task.paths: + cached_parquet_info = None + Scan._validate_cached_parquet_info(base_task.paths, cached_parquet_info) + self.base_task = base_task self.base_scan = base_scan - self.paths = paths - self.row_groups = row_groups - self.skip_rows = skip_rows - self.n_rows = n_rows - self.parquet_options = parquet_options + self.paths = base_task.paths + self.parquet_options = base_task.parquet_options self.cached_parquet_info = cached_parquet_info self.schema = base_scan.schema self._non_child_args = ( - base_scan, - paths, - row_groups, - skip_rows, - n_rows, - parquet_options, + base_task, cached_parquet_info, ) self.children = () + if cached_parquet_info is not None: + _set_scan_task_cached_parquet_info(base_task, cached_parquet_info) - @classmethod - def from_scan(cls, scan: FusedScan | SplitScan) -> Self | None: - """Create an aligned parquet scan from a row-group-aligned scan task.""" - cached_parquet_info = scan.cached_parquet_info - base_scan = scan.base_scan - if base_scan.typ != "parquet" or cached_parquet_info is None: + def attach_cached_parquet_info( + self, cached_parquet_info: list[CachedParquetInfo] + ) -> None: + """Attach cached metadata to this parquet scan task.""" + Scan._validate_cached_parquet_info(self.paths, cached_parquet_info) + self.cached_parquet_info = cached_parquet_info + _set_scan_task_cached_parquet_info(self.base_task, cached_parquet_info) + self._non_child_args = (self.base_task, cached_parquet_info) + + def row_groups_and_slice(self) -> tuple[list[list[int]], int, int] | None: + """Return row groups and row slice when this task is row-group aligned.""" + return self._row_groups_and_slice(self.base_task, self.cached_parquet_info) + + @staticmethod + def _row_groups_and_slice( + base_task: SplitScan | FusedScan, + cached_parquet_info: list[CachedParquetInfo] | None, + ) -> tuple[list[list[int]], int, int] | None: + base_scan = base_task.base_scan + if cached_parquet_info is None: return None if ( base_scan.skip_rows != 0 @@ -765,113 +800,111 @@ def from_scan(cls, scan: FusedScan | SplitScan) -> Self | None: ): return None - if isinstance(scan, SplitScan): - result = _split_scan_row_groups_and_slice(scan) - if result is None: - return None - row_groups, skip_rows, n_rows = result - else: - row_groups = [ - list(range(len(info.file_metadata.row_group_num_rows))) - for info in cached_parquet_info - ] - skip_rows = 0 - n_rows = -1 + if isinstance(base_task, SplitScan): + return _split_scan_row_groups_and_slice(base_task, cached_parquet_info) - return cls( - base_scan, - scan.paths, - row_groups, - skip_rows, - n_rows, - scan.parquet_options, - cached_parquet_info, - ) + row_groups = [ + list(range(len(info.file_metadata.row_group_num_rows))) + for info in cached_parquet_info + ] + return row_groups, 0, -1 def get_hashable(self) -> Hashable: """Hashable representation of the node.""" return ( type(self), - self.base_scan.get_hashable(), - tuple(self.paths), - self.parquet_options, - tuple(tuple(row_groups) for row_groups in self.row_groups), - self.skip_rows, - self.n_rows, + self.base_task.get_hashable(), ) @classmethod def do_evaluate( cls, - base_scan: Scan, - paths: list[str], - row_groups: list[list[int]], - skip_rows: int, - n_rows: int, - parquet_options: ParquetOptions, - cached_parquet_info: list[CachedParquetInfo], + base_task: SplitScan | FusedScan, + cached_parquet_info: list[CachedParquetInfo] | None, *, context: IRExecutionContext, ) -> DataFrame: - """Evaluate a row-group-aligned parquet scan.""" - # 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 ( - len(paths) == 1 - and len(row_groups) == 1 - and hybrid_scan_eligible( - parquet_options, - cached_parquet_info=cached_parquet_info, - row_index=base_scan.row_index, - include_file_paths=base_scan.include_file_paths, - predicate=base_scan.predicate, - ) - ): - assert base_scan.predicate is not None - assert cached_parquet_info is not None - stream = context.get_cuda_stream() - plc_filter, residual = to_parquet_filter( - _prepare_parquet_predicate( - base_scan.predicate.value, - paths, - base_scan.schema, - base_scan.with_columns, - ), - stream=stream, - ) - if plc_filter is not None and residual is None: - return _read_with_hybrid_scan( + """Evaluate a parquet scan task.""" + base_scan = base_task.base_scan + paths = base_task.paths + parquet_options = base_task.parquet_options + row_group_info = cls._row_groups_and_slice(base_task, cached_parquet_info) + + if row_group_info is not None: + row_groups, skip_rows, n_rows = row_group_info + # 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 ( + len(paths) == 1 + and len(row_groups) == 1 + and hybrid_scan_eligible( + parquet_options, + cached_parquet_info=cached_parquet_info, + row_index=base_scan.row_index, + include_file_paths=base_scan.include_file_paths, + predicate=base_scan.predicate, + ) + ): + assert base_scan.predicate is not None + assert cached_parquet_info is not None + stream = context.get_cuda_stream() + plc_filter, residual = to_parquet_filter( + _prepare_parquet_predicate( + base_scan.predicate.value, + paths, + base_scan.schema, + base_scan.with_columns, + ), + stream=stream, + ) + if plc_filter is not None and residual is None: + return _read_with_hybrid_scan( + base_scan.schema, + paths, + base_scan.with_columns, + plc_filter, + row_groups[0], + stream, + cached_parquet_info[0], + stats_pruning=parquet_options._hybrid_scan_stats_pruning, + ) + + with nvtx_annotate_cudf_polars(message=f"ParquetScan: {', '.join(paths)}"): + return Scan.do_evaluate( base_scan.schema, + base_scan.typ, + base_scan.reader_options, paths, base_scan.with_columns, - plc_filter, - row_groups[0], - stream, - cached_parquet_info[0], - stats_pruning=parquet_options._hybrid_scan_stats_pruning, + skip_rows, + n_rows, + base_scan.row_index, + base_scan.include_file_paths, + base_scan.predicate, + parquet_options, + cached_parquet_info, + context=context, ) - with nvtx_annotate_cudf_polars(message=f"ParquetScan: {', '.join(paths)}"): - return Scan.do_evaluate( - base_scan.schema, - base_scan.typ, - base_scan.reader_options, - paths, - base_scan.with_columns, - skip_rows, - n_rows, - base_scan.row_index, - base_scan.include_file_paths, - base_scan.predicate, - parquet_options, - cached_parquet_info, - context=context, - ) + return _evaluate_scan_task( + base_task, + cached_parquet_info, + context=context, + ) + + +StreamingScanTask: TypeAlias = SplitScan | FusedScan | ParquetScanTask -StreamingScanTask: TypeAlias = SplitScan | FusedScan | AlignedParquetScan +def _streaming_scan_task(scan: StreamingScanTask) -> StreamingScanTask: + """Wrap parquet scans in a parquet-specific streaming task.""" + if isinstance(scan, ParquetScanTask): + return scan + if scan.base_scan.typ == "parquet": + return ParquetScanTask(scan) + return scan @lower_ir_node.register(Empty) @@ -962,6 +995,8 @@ def __init__( base_scan: Scan, scan_type: Literal["split", "fused"], ): + if base_scan.typ == "parquet": + scans = [_streaming_scan_task(scan) for scan in scans] self.scans = scans self.base_scan = base_scan self.schema = base_scan.schema diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index d7b3adfe0a85..98e6b65311ab 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -32,8 +32,8 @@ StatsCollector, ) from cudf_polars.streaming.io import ( - AlignedParquetScan, FusedScan, + ParquetScanTask, SplitScan, StreamingScan, expand_scan_for_rank, @@ -467,7 +467,8 @@ def test_expand_scan_for_rank_fused_and_single_read( for scan, expected_paths in zip( streaming_scan.scans, expected_path_groups, strict=True ): - assert isinstance(scan, FusedScan) + assert isinstance(scan, ParquetScanTask) + assert isinstance(scan.base_task, FusedScan) assert scan.paths == expected_paths @@ -497,13 +498,14 @@ def test_expand_scan_for_rank_split_files( for scan, (split_index, total_splits) in zip( streaming_scan.scans, expected_splits, strict=True ): - assert isinstance(scan, SplitScan) - assert scan.split_index == split_index - assert scan.total_splits == total_splits + assert isinstance(scan, ParquetScanTask) + assert isinstance(scan.base_task, SplitScan) + assert scan.base_task.split_index == split_index + assert scan.base_task.total_splits == total_splits assert scan.paths == ["file.parquet"] -def test_attach_cached_parquet_metadata_creates_parquet_scan_tasks( +def test_attach_cached_parquet_metadata_resolves_row_groups( tmp_path: Path, ) -> None: source = tmp_path / "data.parquet" @@ -522,14 +524,16 @@ def test_attach_cached_parquet_metadata_creates_parquet_scan_tasks( cached = prefetch_parquet_file_metadata_for_ir(streaming_scan, None) attach_cached_parquet_metadata(streaming_scan, cached) - parquet_tasks = [ - scan for scan in streaming_scan.scans if isinstance(scan, AlignedParquetScan) - ] - assert len(parquet_tasks) == len(streaming_scan.scans) - assert [scan.row_groups for scan in parquet_tasks] == [[[0]], [[1]]] + row_groups = [] + for scan in streaming_scan.scans: + assert isinstance(scan, ParquetScanTask) + info = scan.row_groups_and_slice() + assert info is not None + row_groups.append(info[0]) + assert row_groups == [[[0]], [[1]]] -def test_attach_cached_parquet_metadata_keeps_sub_row_group_split_scan( +def test_attach_cached_parquet_metadata_leaves_sub_row_group_split_unaligned( tmp_path: Path, ) -> None: source = tmp_path / "data.parquet" @@ -548,14 +552,17 @@ def test_attach_cached_parquet_metadata_keeps_sub_row_group_split_scan( cached = prefetch_parquet_file_metadata_for_ir(streaming_scan, None) attach_cached_parquet_metadata(streaming_scan, cached) - assert all(isinstance(scan, SplitScan) for scan in streaming_scan.scans) + for scan in streaming_scan.scans: + assert isinstance(scan, ParquetScanTask) + assert isinstance(scan.base_task, SplitScan) + assert scan.row_groups_and_slice() is None @pytest.mark.parametrize( "skip_rows,n_rows,row_index", [(1, -1, None), (0, 2, None), (0, -1, ("index", 0))], ) -def test_attach_cached_parquet_metadata_keeps_sliced_fused_scan( +def test_attach_cached_parquet_metadata_leaves_sliced_fused_scan_unaligned( tmp_path: Path, skip_rows: int, n_rows: int, @@ -579,7 +586,10 @@ def test_attach_cached_parquet_metadata_keeps_sliced_fused_scan( cached = prefetch_parquet_file_metadata_for_ir(streaming_scan, None) attach_cached_parquet_metadata(streaming_scan, cached) - assert all(isinstance(scan, FusedScan) for scan in streaming_scan.scans) + for scan in streaming_scan.scans: + assert isinstance(scan, ParquetScanTask) + assert isinstance(scan.base_task, FusedScan) + assert scan.row_groups_and_slice() is None def test_streaming_scan_raises() -> None: From 57e2d42f158083e1560a70e5a91962e161691fe5 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 3 Sep 2026 08:14:22 -0700 Subject: [PATCH 09/31] avoid historical constraint --- .../cudf_polars/cudf_polars/dsl/utils/io.py | 18 ++-- .../cudf_polars/cudf_polars/streaming/io.py | 95 +++++++++++-------- 2 files changed, 66 insertions(+), 47 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index bde00c72cc63..595693407e89 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -16,10 +16,10 @@ from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars from cudf_polars.dsl.traversal import traversal from cudf_polars.streaming.io import ( - ParquetScanTask, ParquetSourceInfo, Scan, StreamingScan, + _set_scan_cached_parquet_info, ) if TYPE_CHECKING: @@ -244,7 +244,7 @@ def attach_cached_parquet_metadata( cached_parquet_info_map: dict[str, CachedParquetInfo], ) -> None: """ - Attach prefetched metadata to parquet scan tasks. + Attach prefetched metadata to parquet scan nodes. This is an optimization only and does not affect IR identity. @@ -257,11 +257,9 @@ def attach_cached_parquet_metadata( """ for node in traversal([root]): if isinstance(node, StreamingScan) and node.base_scan.typ == "parquet": - for scan in node.scans: - if not isinstance(scan, ParquetScanTask) or not all( - path in cached_parquet_info_map for path in scan.paths - ): - continue - - cached = [cached_parquet_info_map[path] for path in scan.paths] - scan.attach_cached_parquet_info(cached) + if not all( + path in cached_parquet_info_map for path in node.base_scan.paths + ): + continue + cached = [cached_parquet_info_map[path] for path in node.base_scan.paths] + _set_scan_cached_parquet_info(node.base_scan, cached) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 9974ae160e9f..bce03c05fb92 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -372,6 +372,56 @@ def _split_scan_row_groups_and_slice( return [row_groups], skip_rows, n_rows +def _cached_parquet_info_for_paths( + paths: list[str], + cached_parquet_info: list[CachedParquetInfo] | None, +) -> list[CachedParquetInfo] | None: + """Return cached parquet metadata matching ``paths``.""" + if cached_parquet_info is None or cached_parquet_info == []: + return None + if paths == [info.path for info in cached_parquet_info]: + return cached_parquet_info + + cached_by_path = {info.path: info for info in cached_parquet_info} + if not all(path in cached_by_path for path in paths): + return None + return [cached_by_path[path] for path in paths] + + +def _cached_parquet_info_for_task( + scan: SplitScan | FusedScan, +) -> list[CachedParquetInfo] | None: + """Return path-aligned cached parquet metadata for ``scan``.""" + return _cached_parquet_info_for_paths( + scan.paths, scan.base_scan.cached_parquet_info + ) or _cached_parquet_info_for_paths(scan.paths, scan.cached_parquet_info) + + +def _set_scan_cached_parquet_info( + scan: Scan, + cached_parquet_info: list[CachedParquetInfo], +) -> None: + """Attach cached parquet metadata to a scan.""" + Scan._validate_cached_parquet_info(scan.paths, cached_parquet_info) + scan.cached_parquet_info = cached_parquet_info + scan._non_child_args = (*scan._non_child_args[:-1], cached_parquet_info) + + +def _fetch_parquet_info_for_task( + scan: SplitScan | FusedScan, +) -> list[CachedParquetInfo]: + """Fetch parquet metadata for ``scan`` without requiring eager prefetch.""" + from cudf_polars.dsl.utils.io import _prefetch_parquet_footers_for_paths + + cached_parquet_info = _prefetch_parquet_footers_for_paths( + scan.paths, + parse_hybrid_metadata=scan.parquet_options.use_hybrid_scan, + ) + if scan.paths == scan.base_scan.paths: + _set_scan_cached_parquet_info(scan.base_scan, cached_parquet_info) + return cached_parquet_info + + class SplitScan(IR): """ Input from a split file. @@ -711,28 +761,18 @@ def _evaluate_scan_task( ) -def _set_scan_task_cached_parquet_info( - scan: SplitScan | FusedScan, - cached_parquet_info: list[CachedParquetInfo], -) -> None: - """Attach cached parquet metadata to a generic streaming scan task.""" - scan.cached_parquet_info = cached_parquet_info - scan._non_child_args = (*scan._non_child_args[:-1], cached_parquet_info) - - class ParquetScanTask(IR): """Parquet scan task wrapping a generic streaming scan task.""" __slots__ = ( "base_scan", "base_task", - "cached_parquet_info", "parquet_options", "paths", "schema", ) _non_child = ("base_task",) - _n_non_child_args = 2 + _n_non_child_args = 1 base_task: SplitScan | FusedScan """Generic streaming scan task being specialized for parquet.""" @@ -742,48 +782,27 @@ class ParquetScanTask(IR): """File paths assigned to this task.""" parquet_options: ParquetOptions """Parquet-specific options.""" - cached_parquet_info: list[CachedParquetInfo] | None - """Cached parquet metadata.""" def __init__( self, base_task: SplitScan | FusedScan, - cached_parquet_info: list[CachedParquetInfo] | None = None, ): base_scan = base_task.base_scan if base_scan.typ != "parquet": # pragma: no cover raise ValueError(f"Expected a parquet scan, got: {base_scan.typ}") - if cached_parquet_info is None: - cached_parquet_info = base_task.cached_parquet_info - if cached_parquet_info == [] and base_task.paths: - cached_parquet_info = None - Scan._validate_cached_parquet_info(base_task.paths, cached_parquet_info) self.base_task = base_task self.base_scan = base_scan self.paths = base_task.paths self.parquet_options = base_task.parquet_options - self.cached_parquet_info = cached_parquet_info self.schema = base_scan.schema - self._non_child_args = ( - base_task, - cached_parquet_info, - ) + self._non_child_args = (base_task,) self.children = () - if cached_parquet_info is not None: - _set_scan_task_cached_parquet_info(base_task, cached_parquet_info) - - def attach_cached_parquet_info( - self, cached_parquet_info: list[CachedParquetInfo] - ) -> None: - """Attach cached metadata to this parquet scan task.""" - Scan._validate_cached_parquet_info(self.paths, cached_parquet_info) - self.cached_parquet_info = cached_parquet_info - _set_scan_task_cached_parquet_info(self.base_task, cached_parquet_info) - self._non_child_args = (self.base_task, cached_parquet_info) def row_groups_and_slice(self) -> tuple[list[list[int]], int, int] | None: """Return row groups and row slice when this task is row-group aligned.""" - return self._row_groups_and_slice(self.base_task, self.cached_parquet_info) + return self._row_groups_and_slice( + self.base_task, _cached_parquet_info_for_task(self.base_task) + ) @staticmethod def _row_groups_and_slice( @@ -820,7 +839,6 @@ def get_hashable(self) -> Hashable: def do_evaluate( cls, base_task: SplitScan | FusedScan, - cached_parquet_info: list[CachedParquetInfo] | None, *, context: IRExecutionContext, ) -> DataFrame: @@ -828,6 +846,9 @@ def do_evaluate( base_scan = base_task.base_scan paths = base_task.paths parquet_options = base_task.parquet_options + cached_parquet_info = _cached_parquet_info_for_task(base_task) + if cached_parquet_info is None and isinstance(base_task, SplitScan): + cached_parquet_info = _fetch_parquet_info_for_task(base_task) row_group_info = cls._row_groups_and_slice(base_task, cached_parquet_info) if row_group_info is not None: From 780d8f5d6ca4a9ca6a1b81724309a597566ca185 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 3 Sep 2026 09:06:12 -0700 Subject: [PATCH 10/31] simplify again --- .../cudf_polars/cudf_polars/streaming/io.py | 165 +++++------------- .../cudf_polars/tests/streaming/test_scan.py | 59 ++----- 2 files changed, 56 insertions(+), 168 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index bce03c05fb92..a1e481884f72 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -394,7 +394,7 @@ def _cached_parquet_info_for_task( """Return path-aligned cached parquet metadata for ``scan``.""" return _cached_parquet_info_for_paths( scan.paths, scan.base_scan.cached_parquet_info - ) or _cached_parquet_info_for_paths(scan.paths, scan.cached_parquet_info) + ) def _set_scan_cached_parquet_info( @@ -407,21 +407,6 @@ def _set_scan_cached_parquet_info( scan._non_child_args = (*scan._non_child_args[:-1], cached_parquet_info) -def _fetch_parquet_info_for_task( - scan: SplitScan | FusedScan, -) -> list[CachedParquetInfo]: - """Fetch parquet metadata for ``scan`` without requiring eager prefetch.""" - from cudf_polars.dsl.utils.io import _prefetch_parquet_footers_for_paths - - cached_parquet_info = _prefetch_parquet_footers_for_paths( - scan.paths, - parse_hybrid_metadata=scan.parquet_options.use_hybrid_scan, - ) - if scan.paths == scan.base_scan.paths: - _set_scan_cached_parquet_info(scan.base_scan, cached_parquet_info) - return cached_parquet_info - - class SplitScan(IR): """ Input from a split file. @@ -434,7 +419,6 @@ class SplitScan(IR): __slots__ = ( "base_scan", - "cached_parquet_info", "parquet_options", "paths", "schema", @@ -442,14 +426,13 @@ class SplitScan(IR): "total_splits", ) _non_child = ( - "schema", "base_scan", "paths", "split_index", "total_splits", "parquet_options", ) - _n_non_child_args = 13 + _n_non_child_args = 5 base_scan: Scan """Scan operation this node is based on.""" paths: list[str] @@ -460,41 +443,28 @@ class SplitScan(IR): """Total number of splits.""" parquet_options: ParquetOptions """Parquet-specific options.""" - cached_parquet_info: list[CachedParquetInfo] | None def __init__( self, - schema: Schema, base_scan: Scan, paths: list[str], split_index: int, total_splits: int, parquet_options: ParquetOptions, - cached_parquet_info: list[CachedParquetInfo] | None = None, ): - self.schema = schema + self.schema = base_scan.schema self.base_scan = base_scan self.paths = paths self.split_index = split_index self.total_splits = total_splits self._non_child_args = ( + base_scan, + paths, split_index, total_splits, - base_scan.schema, - base_scan.typ, - base_scan.reader_options, - paths, - base_scan.with_columns, - base_scan.skip_rows, - base_scan.n_rows, - base_scan.row_index, - base_scan.include_file_paths, - base_scan.predicate, parquet_options, - cached_parquet_info, ) self.parquet_options = parquet_options - self.cached_parquet_info = cached_parquet_info self.children = () if base_scan.typ not in ("parquet",): # pragma: no cover raise NotImplementedError( @@ -516,26 +486,20 @@ def get_hashable(self) -> Hashable: @classmethod def do_evaluate( cls, + base_scan: Scan, + paths: list[str], split_index: int, total_splits: int, - schema: Schema, - typ: str, - reader_options: dict[str, Any], - paths: list[str], - with_columns: list[str] | None, - skip_rows: int, - n_rows: int, - row_index: tuple[str, int] | None, - include_file_paths: str | None, - predicate: NamedExpr | None, parquet_options: ParquetOptions, - cached_parquet_info: list[CachedParquetInfo] | None, *, context: IRExecutionContext, + cached_parquet_info: list[CachedParquetInfo] | None = None, ) -> DataFrame: """Evaluate and return a dataframe.""" - if typ not in ("parquet",): # pragma: no cover - raise NotImplementedError(f"Unhandled Scan type for file splitting: {typ}") + if base_scan.typ not in ("parquet",): # pragma: no cover + raise NotImplementedError( + f"Unhandled Scan type for file splitting: {base_scan.typ}" + ) if len(paths) > 1: # pragma: no cover raise ValueError(f"Expected a single path, got: {paths}") @@ -595,16 +559,16 @@ def do_evaluate( message=f"SplitScan: {paths[0]} [{split_index + 1}/{total_splits}]" ): return Scan.do_evaluate( - schema, - typ, - reader_options, + base_scan.schema, + base_scan.typ, + base_scan.reader_options, paths, - with_columns, + base_scan.with_columns, skip_rows, n_rows, - row_index, - include_file_paths, - predicate, + base_scan.row_index, + base_scan.include_file_paths, + base_scan.predicate, parquet_options, cached_parquet_info, context=context, @@ -621,53 +585,37 @@ class FusedScan(IR): __slots__ = ( "base_scan", - "cached_parquet_info", "parquet_options", "paths", "schema", ) _non_child = ( - "schema", "base_scan", "paths", "parquet_options", ) - _n_non_child_args = 11 + _n_non_child_args = 3 base_scan: Scan """Scan operation this node is based on.""" paths: list[str] """File paths assigned to this task.""" parquet_options: ParquetOptions """Parquet-specific options.""" - cached_parquet_info: list[CachedParquetInfo] | None - """Cached parquet metadata.""" def __init__( self, - schema: Schema, base_scan: Scan, paths: list[str], parquet_options: ParquetOptions, - cached_parquet_info: list[CachedParquetInfo] | None = None, ): - self.schema = schema + self.schema = base_scan.schema self.base_scan = base_scan self.paths = paths self.parquet_options = parquet_options - self.cached_parquet_info = cached_parquet_info self._non_child_args = ( - base_scan.schema, - base_scan.typ, - base_scan.reader_options, + base_scan, paths, - base_scan.with_columns, - base_scan.skip_rows, - base_scan.n_rows, - base_scan.row_index, - base_scan.include_file_paths, - base_scan.predicate, parquet_options, - cached_parquet_info, ) self.children = () @@ -684,34 +632,26 @@ def get_hashable(self) -> Hashable: @classmethod def do_evaluate( cls, - schema: Schema, - typ: str, - reader_options: dict[str, Any], + base_scan: Scan, paths: list[str], - with_columns: list[str] | None, - skip_rows: int, - n_rows: int, - row_index: tuple[str, int] | None, - include_file_paths: str | None, - predicate: NamedExpr | None, parquet_options: ParquetOptions, - cached_parquet_info: list[CachedParquetInfo] | None, *, context: IRExecutionContext, + cached_parquet_info: list[CachedParquetInfo] | None = None, ) -> DataFrame: """Evaluate and return a dataframe.""" with nvtx_annotate_cudf_polars(message=f"FusedScan: {', '.join(paths)}"): return Scan.do_evaluate( - schema, - typ, - reader_options, + base_scan.schema, + base_scan.typ, + base_scan.reader_options, paths, - with_columns, - skip_rows, - n_rows, - row_index, - include_file_paths, - predicate, + base_scan.with_columns, + base_scan.skip_rows, + base_scan.n_rows, + base_scan.row_index, + base_scan.include_file_paths, + base_scan.predicate, parquet_options, cached_parquet_info, context=context, @@ -728,36 +668,20 @@ def _evaluate_scan_task( base_scan = scan.base_scan if isinstance(scan, SplitScan): return SplitScan.do_evaluate( + base_scan, + scan.paths, scan.split_index, scan.total_splits, - base_scan.schema, - base_scan.typ, - base_scan.reader_options, - scan.paths, - base_scan.with_columns, - base_scan.skip_rows, - base_scan.n_rows, - base_scan.row_index, - base_scan.include_file_paths, - base_scan.predicate, scan.parquet_options, - cached_parquet_info, context=context, + cached_parquet_info=cached_parquet_info, ) return FusedScan.do_evaluate( - base_scan.schema, - base_scan.typ, - base_scan.reader_options, + base_scan, scan.paths, - base_scan.with_columns, - base_scan.skip_rows, - base_scan.n_rows, - base_scan.row_index, - base_scan.include_file_paths, - base_scan.predicate, scan.parquet_options, - cached_parquet_info, context=context, + cached_parquet_info=cached_parquet_info, ) @@ -801,7 +725,8 @@ def __init__( def row_groups_and_slice(self) -> tuple[list[list[int]], int, int] | None: """Return row groups and row slice when this task is row-group aligned.""" return self._row_groups_and_slice( - self.base_task, _cached_parquet_info_for_task(self.base_task) + self.base_task, + _cached_parquet_info_for_task(self.base_task), ) @staticmethod @@ -847,8 +772,6 @@ def do_evaluate( paths = base_task.paths parquet_options = base_task.parquet_options cached_parquet_info = _cached_parquet_info_for_task(base_task) - if cached_parquet_info is None and isinstance(base_task, SplitScan): - cached_parquet_info = _fetch_parquet_info_for_task(base_task) row_group_info = cls._row_groups_and_slice(base_task, cached_parquet_info) if row_group_info is not None: @@ -919,7 +842,9 @@ def do_evaluate( StreamingScanTask: TypeAlias = SplitScan | FusedScan | ParquetScanTask -def _streaming_scan_task(scan: StreamingScanTask) -> StreamingScanTask: +def _streaming_scan_task( + scan: StreamingScanTask, +) -> StreamingScanTask: """Wrap parquet scans in a parquet-specific streaming task.""" if isinstance(scan, ParquetScanTask): return scan @@ -1048,13 +973,11 @@ def for_split_files( while sindex < plan.factor and splits_created < local_count: scans.append( SplitScan( - base_scan.schema, base_scan, [path], sindex, plan.factor, parquet_options, - None, ) ) sindex += 1 @@ -1079,11 +1002,9 @@ def for_fused_files( paths_end = paths_start + plan.factor * local_count scans = [ FusedScan( - base_scan.schema, base_scan, base_scan.paths[offset : offset + plan.factor], parquet_options, - None, ) for offset in range(paths_start, paths_end, plan.factor) if base_scan.paths[offset : offset + plan.factor] diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 98e6b65311ab..b9e1fcc91a22 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -180,7 +180,7 @@ def recording_prefetch( ) scan = _make_parquet_scan(paths) - fused = FusedScan(scan.schema, scan, paths, scan.parquet_options, None) + fused = FusedScan(scan, paths, scan.parquet_options) streaming_scan = StreamingScan([fused], scan, "fused") result = prefetch_parquet_file_metadata_for_ir( @@ -196,7 +196,7 @@ def test_prefetch_parquet_file_metadata_remote_only(tmp_path, df) -> None: local_path = str(next(tmp_path.glob("*.parquet"))) scan = _make_parquet_scan([local_path]) - fused = FusedScan(scan.schema, scan, scan.paths, scan.parquet_options, []) + fused = FusedScan(scan, scan.paths, scan.parquet_options) streaming_scan = StreamingScan([fused], scan, "fused") # Local paths are skipped entirely when remote_only=True. @@ -595,7 +595,7 @@ def test_attach_cached_parquet_metadata_leaves_sliced_fused_scan_unaligned( def test_streaming_scan_raises() -> None: # This isn't reachable by normal cudf-polars usage. scan = _make_parquet_scan(["file.parquet"]) - fused = FusedScan(scan.schema, scan, scan.paths, scan.parquet_options, []) + fused = FusedScan(scan, scan.paths, scan.parquet_options) ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): StreamingScan.do_evaluate([fused], scan, context=ctx) @@ -672,7 +672,7 @@ def test_streaming_scan_missing_prefetch_metadata_raises() -> None: scan = _make_parquet_scan( ["file.parquet"], parquet_options=ParquetOptions(prefetch_file_metadata=True) ) - fused = FusedScan(scan.schema, scan, scan.paths, scan.parquet_options, []) + fused = FusedScan(scan, scan.paths, scan.parquet_options) ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): @@ -683,28 +683,19 @@ def test_split_scan_do_evaluate_missing_prefetch_metadata() -> None: paths = ["/some/missing/file.parquet"] parquet_options = ParquetOptions(prefetch_file_metadata=True) context = IRExecutionContext() - schema = {"x": DataType(pl.Int64())} with pytest.raises( AssertionError, match=(r"Paths do not match cached parquet info."), ): SplitScan.do_evaluate( - 0, - 4, - schema, - "parquet", - {}, + _make_parquet_scan(paths, parquet_options), paths, - None, 0, - -1, - None, - None, - None, + 4, parquet_options, - [], context=context, + cached_parquet_info=[], ) @@ -766,11 +757,10 @@ def test_prefetch_file_metadata_with_cached_scan_parent_nodes( def test_fused_scan_identity_equality() -> None: base = _make_parquet_scan(["a.parquet", "b.parquet"]) paths = ["a.parquet"] - info = _make_cached_parquet_info(paths) - a = FusedScan(base.schema, base, paths, base.parquet_options, info) - b = FusedScan(base.schema, base, paths, base.parquet_options, info.copy()) - c = FusedScan(base.schema, base, ["b.parquet"], base.parquet_options, info) + a = FusedScan(base, paths, base.parquet_options) + b = FusedScan(base, paths, base.parquet_options) + c = FusedScan(base, ["b.parquet"], base.parquet_options) assert a == b assert hash(a) == hash(b) @@ -779,13 +769,10 @@ def test_fused_scan_identity_equality() -> None: def test_split_scan_identity_equality() -> None: base = _make_parquet_scan(["a.parquet"]) - info = _make_cached_parquet_info(base.paths) - a = SplitScan(base.schema, base, base.paths, 0, 4, base.parquet_options, info) - b = SplitScan( - base.schema, base, base.paths, 0, 4, base.parquet_options, info.copy() - ) - c = SplitScan(base.schema, base, base.paths, 1, 4, base.parquet_options, info) + a = SplitScan(base, base.paths, 0, 4, base.parquet_options) + b = SplitScan(base, base.paths, 0, 4, base.parquet_options) + c = SplitScan(base, base.paths, 1, 4, base.parquet_options) assert a == b assert hash(a) == hash(b) @@ -795,31 +782,25 @@ def test_split_scan_identity_equality() -> None: def test_streaming_scan_identity_equality() -> None: base = _make_parquet_scan(["a.parquet"]) split = SplitScan( - base.schema, base, base.paths, 0, 2, base.parquet_options, - _make_cached_parquet_info(base.paths, size=10), ) split_same = SplitScan( - base.schema, base, base.paths, 0, 2, base.parquet_options, - _make_cached_parquet_info(base.paths, size=11), ) split_diff = SplitScan( - base.schema, base, base.paths, 1, 2, base.parquet_options, - _make_cached_parquet_info(base.paths, size=10), ) a = StreamingScan([split], base, "split") @@ -854,20 +835,6 @@ def test_cached_parquet_info_excluded_from_identity() -> None: assert scan_without == scan_with assert hash(scan_without) == hash(scan_with) - split_without = SplitScan( - base.schema, base, base.paths, 0, 4, base.parquet_options, None - ) - split_with = SplitScan( - base.schema, base, base.paths, 0, 4, base.parquet_options, info - ) - assert split_without == split_with - assert hash(split_without) == hash(split_with) - - fused_without = FusedScan(base.schema, base, base.paths, base.parquet_options, None) - fused_with = FusedScan(base.schema, base, base.paths, base.parquet_options, info) - assert fused_without == fused_with - assert hash(fused_without) == hash(fused_with) - class FooSource(DataSourceInfo): def __init__(self, size: int): From b327ec79a511e75a981e9ffbcffb0abbcd230116 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 3 Sep 2026 09:23:19 -0700 Subject: [PATCH 11/31] minor cleanup --- python/cudf_polars/cudf_polars/streaming/io.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index a1e481884f72..6b4a118ea121 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -388,15 +388,6 @@ def _cached_parquet_info_for_paths( return [cached_by_path[path] for path in paths] -def _cached_parquet_info_for_task( - scan: SplitScan | FusedScan, -) -> list[CachedParquetInfo] | None: - """Return path-aligned cached parquet metadata for ``scan``.""" - return _cached_parquet_info_for_paths( - scan.paths, scan.base_scan.cached_parquet_info - ) - - def _set_scan_cached_parquet_info( scan: Scan, cached_parquet_info: list[CachedParquetInfo], @@ -726,7 +717,10 @@ def row_groups_and_slice(self) -> tuple[list[list[int]], int, int] | None: """Return row groups and row slice when this task is row-group aligned.""" return self._row_groups_and_slice( self.base_task, - _cached_parquet_info_for_task(self.base_task), + _cached_parquet_info_for_paths( + self.base_task.paths, + self.base_task.base_scan.cached_parquet_info, + ), ) @staticmethod @@ -771,7 +765,9 @@ def do_evaluate( base_scan = base_task.base_scan paths = base_task.paths parquet_options = base_task.parquet_options - cached_parquet_info = _cached_parquet_info_for_task(base_task) + cached_parquet_info = _cached_parquet_info_for_paths( + base_task.paths, base_task.base_scan.cached_parquet_info + ) row_group_info = cls._row_groups_and_slice(base_task, cached_parquet_info) if row_group_info is not None: From 6426ead85719479fdd00de9723194346a8bfb0c4 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 3 Sep 2026 09:59:25 -0700 Subject: [PATCH 12/31] more cleanup --- .../cudf_polars/cudf_polars/dsl/utils/io.py | 12 +- .../cudf_polars/cudf_polars/streaming/io.py | 122 ++++++++---------- 2 files changed, 58 insertions(+), 76 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index 595693407e89..df568dc52fce 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -19,7 +19,6 @@ ParquetSourceInfo, Scan, StreamingScan, - _set_scan_cached_parquet_info, ) if TYPE_CHECKING: @@ -257,9 +256,10 @@ def attach_cached_parquet_metadata( """ for node in traversal([root]): if isinstance(node, StreamingScan) and node.base_scan.typ == "parquet": - if not all( - path in cached_parquet_info_map for path in node.base_scan.paths - ): + base_scan = node.base_scan + if not all(path in cached_parquet_info_map for path in base_scan.paths): continue - cached = [cached_parquet_info_map[path] for path in node.base_scan.paths] - _set_scan_cached_parquet_info(node.base_scan, cached) + cached = [cached_parquet_info_map[path] for path in base_scan.paths] + Scan._validate_cached_parquet_info(base_scan.paths, cached) + base_scan.cached_parquet_info = cached + base_scan._non_child_args = (*base_scan._non_child_args[:-1], cached) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 6b4a118ea121..d0bb91640f75 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -11,7 +11,7 @@ import statistics from collections import defaultdict from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal, Self, TypeAlias, overload +from typing import TYPE_CHECKING, Any, Literal, NamedTuple, Self, TypeAlias, overload import polars as pl @@ -356,27 +356,58 @@ def _split_scan_row_groups_and_slice( if scan.total_splits > total_row_groups: return None - row_group_stride = total_row_groups // scan.total_splits - start = row_group_stride * scan.split_index - stop = ( - total_row_groups - if scan.split_index == scan.total_splits - 1 - else start + row_group_stride - ) - row_groups = list(range(start, stop)) + bounds = _split_scan_bounds(row_group_num_rows, scan.split_index, scan.total_splits) + row_groups = list(range(bounds.row_group_start, bounds.row_group_stop)) if not row_groups: return None - skip_rows = sum(row_group_num_rows[:start]) - n_rows = -1 if stop == total_row_groups else sum(row_group_num_rows[start:stop]) - return [row_groups], skip_rows, n_rows + return [row_groups], bounds.skip_rows, bounds.n_rows -def _cached_parquet_info_for_paths( - paths: list[str], - cached_parquet_info: list[CachedParquetInfo] | None, +class _SplitScanBounds(NamedTuple): + """Row and row-group bounds for a split scan.""" + + row_group_start: int + row_group_stop: int + skip_rows: int + n_rows: int + + +def _split_scan_bounds( + row_group_num_rows: Sequence[int], + split_index: int, + total_splits: int, +) -> _SplitScanBounds: + """Return row and row-group bounds for a file split.""" + total_row_groups = len(row_group_num_rows) + if total_splits <= total_row_groups: + row_group_stride = total_row_groups // total_splits + row_group_start = row_group_stride * split_index + row_group_stop = ( + total_row_groups + if split_index == total_splits - 1 + else row_group_start + row_group_stride + ) + skip_rows = sum(row_group_num_rows[:row_group_start]) + n_rows = sum(row_group_num_rows[row_group_start:row_group_stop]) + else: + row_group_start = 0 + row_group_stop = 0 + total_rows = sum(row_group_num_rows) + n_rows = total_rows // total_splits + skip_rows = n_rows * split_index + + if split_index == total_splits - 1: + n_rows = -1 + return _SplitScanBounds(row_group_start, row_group_stop, skip_rows, n_rows) + + +def _cached_parquet_info_for_scan( + scan: SplitScan | FusedScan, ) -> list[CachedParquetInfo] | None: - """Return cached parquet metadata matching ``paths``.""" + """Return cached parquet metadata matching a scan task.""" + paths = scan.paths + cached_parquet_info = scan.base_scan.cached_parquet_info if cached_parquet_info is None or cached_parquet_info == []: return None if paths == [info.path for info in cached_parquet_info]: @@ -388,16 +419,6 @@ def _cached_parquet_info_for_paths( return [cached_by_path[path] for path in paths] -def _set_scan_cached_parquet_info( - scan: Scan, - cached_parquet_info: list[CachedParquetInfo], -) -> None: - """Attach cached parquet metadata to a scan.""" - Scan._validate_cached_parquet_info(scan.paths, cached_parquet_info) - scan.cached_parquet_info = cached_parquet_info - scan._non_child_args = (*scan._non_child_args[:-1], cached_parquet_info) - - class SplitScan(IR): """ Input from a split file. @@ -495,24 +516,13 @@ def do_evaluate( if len(paths) > 1: # pragma: no cover raise ValueError(f"Expected a single path, got: {paths}") - # Parquet logic: - # - We are one of "total_splits" SplitScan nodes - # assigned to the same file. - # - We know our index within this file ("split_index") - # - We can also use parquet metadata to query the - # total number of rows in each row-group of the file. - # - We can use all this information to calculate the - # "skip_rows" and "n_rows" options to use locally. - if cached_parquet_info is not None: parquet_metadatas = [info.file_metadata for info in cached_parquet_info] - row_group_num_rows = [ num_rows for metadata in parquet_metadatas for num_rows in metadata.row_group_num_rows ] - else: row_group_num_rows = [ rg["num_rows"] @@ -520,30 +530,7 @@ def do_evaluate( plc.io.SourceInfo(paths) ).rowgroup_metadata() ] - - total_row_groups = len(row_group_num_rows) - if total_splits <= total_row_groups: - # We have enough row-groups in the file to align - # all "total_splits" of our reads with row-group - # boundaries. Calculate which row-groups to include - # in the current read, and use metadata to translate - # the row-group indices to "skip_rows" and "n_rows". - rg_stride = total_row_groups // total_splits - 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]) - else: - # There are not enough row-groups to align - # all "total_splits" of our reads with row-group - # boundaries. Use metadata to directly calculate - # "skip_rows" and "n_rows" for the current read. - total_rows = sum(row_group_num_rows) - n_rows = total_rows // total_splits - skip_rows = n_rows * split_index - - # Last split should always read to end of file - if split_index == (total_splits - 1): - n_rows = -1 + bounds = _split_scan_bounds(row_group_num_rows, split_index, total_splits) # Perform the partial read with nvtx_annotate_cudf_polars( @@ -555,8 +542,8 @@ def do_evaluate( base_scan.reader_options, paths, base_scan.with_columns, - skip_rows, - n_rows, + bounds.skip_rows, + bounds.n_rows, base_scan.row_index, base_scan.include_file_paths, base_scan.predicate, @@ -717,10 +704,7 @@ def row_groups_and_slice(self) -> tuple[list[list[int]], int, int] | None: """Return row groups and row slice when this task is row-group aligned.""" return self._row_groups_and_slice( self.base_task, - _cached_parquet_info_for_paths( - self.base_task.paths, - self.base_task.base_scan.cached_parquet_info, - ), + _cached_parquet_info_for_scan(self.base_task), ) @staticmethod @@ -765,9 +749,7 @@ def do_evaluate( base_scan = base_task.base_scan paths = base_task.paths parquet_options = base_task.parquet_options - cached_parquet_info = _cached_parquet_info_for_paths( - base_task.paths, base_task.base_scan.cached_parquet_info - ) + cached_parquet_info = _cached_parquet_info_for_scan(base_task) row_group_info = cls._row_groups_and_slice(base_task, cached_parquet_info) if row_group_info is not None: From 018aa0b17d952fd0c8ca78ce45ea31c358ebf78e Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 3 Sep 2026 11:43:02 -0700 Subject: [PATCH 13/31] addresss CI --- python/cudf_polars/cudf_polars/dsl/utils/io.py | 9 ++++++++- .../cudf_polars/cudf_polars/streaming/actor_graph/io.py | 6 ++++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index df568dc52fce..f41ed2ada084 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -6,6 +6,7 @@ import concurrent.futures import contextlib +import inspect from dataclasses import dataclass, field from typing import TYPE_CHECKING @@ -21,6 +22,10 @@ StreamingScan, ) +_CACHED_PARQUET_INFO_ARG_INDEX = tuple( + inspect.signature(Scan.do_evaluate).parameters +).index("cached_parquet_info") + if TYPE_CHECKING: from cudf_polars.dsl.ir import IR from cudf_polars.streaming.base import StatsCollector @@ -262,4 +267,6 @@ def attach_cached_parquet_metadata( cached = [cached_parquet_info_map[path] for path in base_scan.paths] Scan._validate_cached_parquet_info(base_scan.paths, cached) base_scan.cached_parquet_info = cached - base_scan._non_child_args = (*base_scan._non_child_args[:-1], cached) + args = list(base_scan._non_child_args) + args[_CACHED_PARQUET_INFO_ARG_INDEX] = cached + base_scan._non_child_args = tuple(args) 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 0299965db5c0..bb5cb9fe42ce 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -39,6 +39,7 @@ send_metadata, ) from cudf_polars.streaming.io import ( + ParquetScanTask, StreamingScan, StreamingSink, _prepare_sink_directory, @@ -569,14 +570,15 @@ async def read_chunk( br=context.br(), ) stop = time.monotonic_ns() + trace_scan = scan.base_task if isinstance(scan, ParquetScanTask) else scan log( "IO Task", scope=Scope.IO_TASK.value, start=start, admitted=admitted, stop=stop, - ir_id=scan.get_stable_id(), - ir_type=type(scan).__name__, + ir_id=trace_scan.get_stable_id(), + ir_type=type(trace_scan).__name__, sequence_number=seq_num, estimated_output_bytes=estimated_chunk_bytes, reservation_bytes=reservation_bytes, From b3c11e295d064d7a7852c5ee9134bca29924ae6b Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 3 Sep 2026 11:46:50 -0700 Subject: [PATCH 14/31] remove silly fix --- .../cudf_polars/cudf_polars/dsl/utils/io.py | 34 +++++++++++++------ 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index f41ed2ada084..25b46cb8ceb0 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -6,7 +6,6 @@ import concurrent.futures import contextlib -import inspect from dataclasses import dataclass, field from typing import TYPE_CHECKING @@ -22,10 +21,6 @@ StreamingScan, ) -_CACHED_PARQUET_INFO_ARG_INDEX = tuple( - inspect.signature(Scan.do_evaluate).parameters -).index("cached_parquet_info") - if TYPE_CHECKING: from cudf_polars.dsl.ir import IR from cudf_polars.streaming.base import StatsCollector @@ -243,6 +238,29 @@ def prefetch_parquet_file_metadata_for_ir( return cached_parquet_info +def _set_scan_cached_parquet_info( + scan: Scan, + cached_parquet_info: list[CachedParquetInfo], +) -> None: + """Attach cached parquet metadata to a scan without changing its identity.""" + Scan._validate_cached_parquet_info(scan.paths, cached_parquet_info) + scan.cached_parquet_info = cached_parquet_info + scan._non_child_args = ( + scan.schema, + scan.typ, + scan.reader_options, + scan.paths, + scan.with_columns, + scan.skip_rows, + scan.n_rows, + scan.row_index, + scan.include_file_paths, + scan.predicate, + scan.parquet_options, + cached_parquet_info, + ) + + def attach_cached_parquet_metadata( root: IR, cached_parquet_info_map: dict[str, CachedParquetInfo], @@ -265,8 +283,4 @@ def attach_cached_parquet_metadata( if not all(path in cached_parquet_info_map for path in base_scan.paths): continue cached = [cached_parquet_info_map[path] for path in base_scan.paths] - Scan._validate_cached_parquet_info(base_scan.paths, cached) - base_scan.cached_parquet_info = cached - args = list(base_scan._non_child_args) - args[_CACHED_PARQUET_INFO_ARG_INDEX] = cached - base_scan._non_child_args = tuple(args) + _set_scan_cached_parquet_info(base_scan, cached) From b7169878750115ffba407b18491878a00fe45ffb Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 3 Sep 2026 11:49:29 -0700 Subject: [PATCH 15/31] roll back --- python/cudf_polars/cudf_polars/dsl/utils/io.py | 15 +-------------- 1 file changed, 1 insertion(+), 14 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index 25b46cb8ceb0..a7fea07d0d54 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -245,20 +245,7 @@ def _set_scan_cached_parquet_info( """Attach cached parquet metadata to a scan without changing its identity.""" Scan._validate_cached_parquet_info(scan.paths, cached_parquet_info) scan.cached_parquet_info = cached_parquet_info - scan._non_child_args = ( - scan.schema, - scan.typ, - scan.reader_options, - scan.paths, - scan.with_columns, - scan.skip_rows, - scan.n_rows, - scan.row_index, - scan.include_file_paths, - scan.predicate, - scan.parquet_options, - cached_parquet_info, - ) + scan._non_child_args = (*scan._non_child_args[:-1], cached_parquet_info) def attach_cached_parquet_metadata( From 6e1d5fc1b68a4d761d72d2f7ddae90e701f49514 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 3 Sep 2026 11:51:36 -0700 Subject: [PATCH 16/31] try again --- python/cudf_polars/cudf_polars/dsl/utils/io.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index a7fea07d0d54..df568dc52fce 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -238,16 +238,6 @@ def prefetch_parquet_file_metadata_for_ir( return cached_parquet_info -def _set_scan_cached_parquet_info( - scan: Scan, - cached_parquet_info: list[CachedParquetInfo], -) -> None: - """Attach cached parquet metadata to a scan without changing its identity.""" - Scan._validate_cached_parquet_info(scan.paths, cached_parquet_info) - scan.cached_parquet_info = cached_parquet_info - scan._non_child_args = (*scan._non_child_args[:-1], cached_parquet_info) - - def attach_cached_parquet_metadata( root: IR, cached_parquet_info_map: dict[str, CachedParquetInfo], @@ -270,4 +260,6 @@ def attach_cached_parquet_metadata( if not all(path in cached_parquet_info_map for path in base_scan.paths): continue cached = [cached_parquet_info_map[path] for path in base_scan.paths] - _set_scan_cached_parquet_info(base_scan, cached) + Scan._validate_cached_parquet_info(base_scan.paths, cached) + base_scan.cached_parquet_info = cached + base_scan._non_child_args = (*base_scan._non_child_args[:-1], cached) From 5e1278b08cd6c5b199a272a37558eea80d45e6c9 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 3 Sep 2026 12:16:25 -0700 Subject: [PATCH 17/31] more scan -> task cleanup --- .../cudf_polars/cudf_polars/dsl/utils/io.py | 4 +- .../cudf_polars/streaming/actor_graph/io.py | 42 +++++----- .../cudf_polars/streaming/explain.py | 2 +- .../cudf_polars/cudf_polars/streaming/io.py | 76 ++++++++----------- .../tests/streaming/test_explain.py | 2 +- .../cudf_polars/tests/streaming/test_scan.py | 42 +++------- 6 files changed, 66 insertions(+), 102 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index df568dc52fce..dd3f7fe480a8 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -189,8 +189,8 @@ def prefetch_parquet_file_metadata_for_ir( for node in traversal([root]): if isinstance(node, StreamingScan) and node.base_scan.typ == "parquet": - for scan in node.scans: - for path in scan.paths: + for task in node.tasks: + for path in task.paths: all_paths.add(path) elif isinstance(node, Scan) and node.typ == "parquet": # pragma: no cover raise RuntimeError("Unexpected parquet 'Scan' node in lowered IR graph.") 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 bb5cb9fe42ce..acd60d0dcd1e 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -517,7 +517,7 @@ def _( async def read_chunk( context: Context, - scan: IR, + task: IR, seq_num: int, ch_out: Channel[TableChunk], ir_context: IRExecutionContext, @@ -531,8 +531,8 @@ async def read_chunk( ---------- context The rapidsmpf context. - scan - The Scan or DataFrameScan node. + task + The scan task to evaluate. seq_num The sequence number. ch_out @@ -547,7 +547,7 @@ async def read_chunk( """ reservation_bytes = ( estimated_chunk_bytes - if isinstance(scan, DataFrameScan) + if isinstance(task, DataFrameScan) else 2 * estimated_chunk_bytes ) start = time.monotonic_ns() @@ -559,8 +559,8 @@ 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, + task.do_evaluate, + *task._non_child_args, context=ir_context, ) chunk = TableChunk.from_pylibcudf_table( @@ -570,15 +570,15 @@ async def read_chunk( br=context.br(), ) stop = time.monotonic_ns() - trace_scan = scan.base_task if isinstance(scan, ParquetScanTask) else scan + trace_task = task.base_task if isinstance(task, ParquetScanTask) else task log( "IO Task", scope=Scope.IO_TASK.value, start=start, admitted=admitted, stop=stop, - ir_id=trace_scan.get_stable_id(), - ir_type=type(trace_scan).__name__, + ir_id=trace_task.get_stable_id(), + ir_type=type(trace_task).__name__, sequence_number=seq_num, estimated_output_bytes=estimated_chunk_bytes, reservation_bytes=reservation_bytes, @@ -615,7 +615,7 @@ async def scan_node( Estimated retained output size of each chunk in bytes. Used to estimate peak memory for admission before launching each read. """ - scans: Sequence[StreamingScanTask] = ir.scans + tasks: Sequence[StreamingScanTask] = ir.tasks async with shutdown_on_error( context, ch_out, trace_ir=ir, ir_context=ir_context @@ -624,21 +624,21 @@ async def scan_node( await send_metadata( ch_out, context, - ChannelMetadata(local_count=len(scans)), + ChannelMetadata(local_count=len(tasks)), ) # If there is nothing to scan, drain the channel and return - if len(scans) == 0: + if len(tasks) == 0: await ch_out.drain(context) return - # If there is only one scan or one producer, we can + # If there is only one task 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): + if len(tasks) == 1 or num_producers == 1: + for seq_num, task in enumerate(tasks): await read_chunk( context, - scan, + task, seq_num, ch_out, ir_context, @@ -649,22 +649,22 @@ async def scan_node( return # Use Lineariser to ensure ordered delivery - num_producers = min(num_producers, len(scans)) + num_producers = min(num_producers, len(tasks)) lineariser = Lineariser(context, ch_out, num_producers) # Assign tasks to producers using round-robin producer_tasks: list[list[tuple[int, StreamingScanTask]]] = [ [] for _ in range(num_producers) ] - for task_idx, scan in enumerate(scans): + for task_idx, task in enumerate(tasks): producer_id = task_idx % num_producers - producer_tasks[producer_id].append((task_idx, scan)) + producer_tasks[producer_id].append((task_idx, task)) async def _producer(producer_id: int, ch_out: Channel) -> None: - for task_idx, scan in producer_tasks[producer_id]: + for task_idx, task in producer_tasks[producer_id]: await read_chunk( context, - scan, + task, task_idx, ch_out, ir_context, diff --git a/python/cudf_polars/cudf_polars/streaming/explain.py b/python/cudf_polars/cudf_polars/streaming/explain.py index 011e3be232a7..8d69a9292f90 100644 --- a/python/cudf_polars/cudf_polars/streaming/explain.py +++ b/python/cudf_polars/cudf_polars/streaming/explain.py @@ -563,7 +563,7 @@ def _(ir: Scan) -> dict[str, Serializable]: def _(ir: StreamingScan) -> dict[str, Serializable]: return { "typ": ir.base_scan.typ, - "scan_count": len(ir.scans), + "task_count": len(ir.tasks), "prefix": os.path.commonprefix(ir.base_scan.paths), "predicate": ( _serialize_expr(ir.base_scan.predicate) if ir.base_scan.predicate else None diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index d0bb91640f75..1a22d226a234 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -402,12 +402,12 @@ def _split_scan_bounds( return _SplitScanBounds(row_group_start, row_group_stop, skip_rows, n_rows) -def _cached_parquet_info_for_scan( - scan: SplitScan | FusedScan, +def _cached_parquet_info_for_paths( + base_scan: Scan, + paths: list[str], ) -> list[CachedParquetInfo] | None: - """Return cached parquet metadata matching a scan task.""" - paths = scan.paths - cached_parquet_info = scan.base_scan.cached_parquet_info + """Return cached parquet metadata matching a scan task's paths.""" + cached_parquet_info = base_scan.cached_parquet_info if cached_parquet_info is None or cached_parquet_info == []: return None if paths == [info.path for info in cached_parquet_info]: @@ -505,7 +505,6 @@ def do_evaluate( parquet_options: ParquetOptions, *, context: IRExecutionContext, - cached_parquet_info: list[CachedParquetInfo] | None = None, ) -> DataFrame: """Evaluate and return a dataframe.""" if base_scan.typ not in ("parquet",): # pragma: no cover @@ -516,6 +515,7 @@ def do_evaluate( if len(paths) > 1: # pragma: no cover raise ValueError(f"Expected a single path, got: {paths}") + cached_parquet_info = _cached_parquet_info_for_paths(base_scan, paths) if cached_parquet_info is not None: parquet_metadatas = [info.file_metadata for info in cached_parquet_info] row_group_num_rows = [ @@ -615,9 +615,9 @@ def do_evaluate( parquet_options: ParquetOptions, *, context: IRExecutionContext, - cached_parquet_info: list[CachedParquetInfo] | None = None, ) -> DataFrame: """Evaluate and return a dataframe.""" + cached_parquet_info = _cached_parquet_info_for_paths(base_scan, paths) with nvtx_annotate_cudf_polars(message=f"FusedScan: {', '.join(paths)}"): return Scan.do_evaluate( base_scan.schema, @@ -638,11 +638,10 @@ def do_evaluate( def _evaluate_scan_task( scan: SplitScan | FusedScan, - cached_parquet_info: list[CachedParquetInfo] | None, *, context: IRExecutionContext, ) -> DataFrame: - """Evaluate a generic streaming scan task with optional parquet metadata.""" + """Evaluate a generic streaming scan task.""" base_scan = scan.base_scan if isinstance(scan, SplitScan): return SplitScan.do_evaluate( @@ -652,14 +651,12 @@ def _evaluate_scan_task( scan.total_splits, scan.parquet_options, context=context, - cached_parquet_info=cached_parquet_info, ) return FusedScan.do_evaluate( base_scan, scan.paths, scan.parquet_options, context=context, - cached_parquet_info=cached_parquet_info, ) @@ -704,7 +701,7 @@ def row_groups_and_slice(self) -> tuple[list[list[int]], int, int] | None: """Return row groups and row slice when this task is row-group aligned.""" return self._row_groups_and_slice( self.base_task, - _cached_parquet_info_for_scan(self.base_task), + _cached_parquet_info_for_paths(self.base_scan, self.paths), ) @staticmethod @@ -749,7 +746,7 @@ def do_evaluate( base_scan = base_task.base_scan paths = base_task.paths parquet_options = base_task.parquet_options - cached_parquet_info = _cached_parquet_info_for_scan(base_task) + cached_parquet_info = _cached_parquet_info_for_paths(base_scan, paths) row_group_info = cls._row_groups_and_slice(base_task, cached_parquet_info) if row_group_info is not None: @@ -812,7 +809,6 @@ def do_evaluate( return _evaluate_scan_task( base_task, - cached_parquet_info, context=context, ) @@ -820,17 +816,6 @@ def do_evaluate( StreamingScanTask: TypeAlias = SplitScan | FusedScan | ParquetScanTask -def _streaming_scan_task( - scan: StreamingScanTask, -) -> StreamingScanTask: - """Wrap parquet scans in a parquet-specific streaming task.""" - if isinstance(scan, ParquetScanTask): - return scan - if scan.base_scan.typ == "parquet": - return ParquetScanTask(scan) - return scan - - @lower_ir_node.register(Empty) def _( ir: Empty, rec: LowerIRTransformer @@ -900,32 +885,31 @@ class StreamingScan(IR): __slots__ = ( "base_scan", - "scan_type", - "scans", "schema", + "tasks", ) _non_child = ( - "scans", + "tasks", "base_scan", - "scan_type", ) - _n_non_child_args = 3 - scans: Sequence[StreamingScanTask] + _n_non_child_args = 2 base_scan: Scan + tasks: Sequence[StreamingScanTask] def __init__( self, - scans: Sequence[StreamingScanTask], + tasks: Sequence[StreamingScanTask], base_scan: Scan, - scan_type: Literal["split", "fused"], ): if base_scan.typ == "parquet": - scans = [_streaming_scan_task(scan) for scan in scans] - self.scans = scans + tasks = [ + task if isinstance(task, ParquetScanTask) else ParquetScanTask(task) + for task in tasks + ] self.base_scan = base_scan self.schema = base_scan.schema - self.scan_type = scan_type - self._non_child_args = (scans, base_scan, scan_type) + self.tasks = tasks + self._non_child_args = (tasks, base_scan) self.children = () @classmethod @@ -945,11 +929,11 @@ def for_split_files( path_end = math.ceil((local_offset + local_count) / plan.factor) local_paths = base_scan.paths[path_offset:path_end] sindex = local_offset % plan.factor - scans: list[SplitScan] = [] + tasks: list[SplitScan] = [] splits_created = 0 for path in local_paths: while sindex < plan.factor and splits_created < local_count: - scans.append( + tasks.append( SplitScan( base_scan, [path], @@ -961,7 +945,7 @@ def for_split_files( sindex += 1 splits_created += 1 sindex = 0 - return cls(scans, base_scan, "split") + return cls(tasks, base_scan) @classmethod def for_fused_files( @@ -978,7 +962,7 @@ def for_fused_files( local_offset, local_count = _rank_slice(partition_count, rank, nranks) paths_start = local_offset * plan.factor paths_end = paths_start + plan.factor * local_count - scans = [ + tasks = [ FusedScan( base_scan, base_scan.paths[offset : offset + plan.factor], @@ -987,24 +971,24 @@ def for_fused_files( for offset in range(paths_start, paths_end, plan.factor) if base_scan.paths[offset : offset + plan.factor] ] - return cls(scans, base_scan, "fused") + return cls(tasks, base_scan) def get_hashable(self) -> Hashable: """Hashable representation of the node.""" - # We don't need to include base_scan / schema, since it's in all the scan nodes. - return (type(self), *tuple(x.get_hashable() for x in self.scans)) + # We don't need to include base_scan / schema, since it's in all the scan tasks. + return (type(self), *tuple(task.get_hashable() for task in self.tasks)) @classmethod def do_evaluate( cls, - scans: Sequence[StreamingScanTask], + tasks: Sequence[StreamingScanTask], base_scan: Scan, *, context: IRExecutionContext, ) -> DataFrame: """Raises NotImplementedError for StreamingScan nodes.""" raise NotImplementedError( - "StreamingScan.do_evaluate should not be called directly. Call Scan.do_evaluate on each scan node instead." + "StreamingScan.do_evaluate should not be called directly. Call Scan.do_evaluate on each scan task instead." ) diff --git a/python/cudf_polars/tests/streaming/test_explain.py b/python/cudf_polars/tests/streaming/test_explain.py index 5ea6be578ae4..707811cd5a71 100644 --- a/python/cudf_polars/tests/streaming/test_explain.py +++ b/python/cudf_polars/tests/streaming/test_explain.py @@ -558,7 +558,7 @@ def test_scan_properties(tmp_path: Path, predicate: pl.Expr | None): "prefix": f"{root}/", "typ": "parquet", "predicate": None, - "scan_count": 1, + "task_count": 1, } if predicate is not None: q = q.filter(predicate) diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index b9e1fcc91a22..33caef738b18 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -181,7 +181,7 @@ def recording_prefetch( scan = _make_parquet_scan(paths) fused = FusedScan(scan, paths, scan.parquet_options) - streaming_scan = StreamingScan([fused], scan, "fused") + streaming_scan = StreamingScan([fused], scan) result = prefetch_parquet_file_metadata_for_ir( streaming_scan, py_executor=None, stats=stats @@ -197,7 +197,7 @@ def test_prefetch_parquet_file_metadata_remote_only(tmp_path, df) -> None: scan = _make_parquet_scan([local_path]) fused = FusedScan(scan, scan.paths, scan.parquet_options) - streaming_scan = StreamingScan([fused], scan, "fused") + streaming_scan = StreamingScan([fused], scan) # Local paths are skipped entirely when remote_only=True. result = prefetch_parquet_file_metadata_for_ir( @@ -465,7 +465,7 @@ def test_expand_scan_for_rank_fused_and_single_read( parquet_options=ParquetOptions(), ) for scan, expected_paths in zip( - streaming_scan.scans, expected_path_groups, strict=True + streaming_scan.tasks, expected_path_groups, strict=True ): assert isinstance(scan, ParquetScanTask) assert isinstance(scan.base_task, FusedScan) @@ -494,9 +494,9 @@ def test_expand_scan_for_rank_split_files( nranks=2, parquet_options=ParquetOptions(), ) - assert len(streaming_scan.scans) == len(expected_splits) + assert len(streaming_scan.tasks) == len(expected_splits) for scan, (split_index, total_splits) in zip( - streaming_scan.scans, expected_splits, strict=True + streaming_scan.tasks, expected_splits, strict=True ): assert isinstance(scan, ParquetScanTask) assert isinstance(scan.base_task, SplitScan) @@ -525,7 +525,7 @@ def test_attach_cached_parquet_metadata_resolves_row_groups( attach_cached_parquet_metadata(streaming_scan, cached) row_groups = [] - for scan in streaming_scan.scans: + for scan in streaming_scan.tasks: assert isinstance(scan, ParquetScanTask) info = scan.row_groups_and_slice() assert info is not None @@ -552,7 +552,7 @@ def test_attach_cached_parquet_metadata_leaves_sub_row_group_split_unaligned( cached = prefetch_parquet_file_metadata_for_ir(streaming_scan, None) attach_cached_parquet_metadata(streaming_scan, cached) - for scan in streaming_scan.scans: + for scan in streaming_scan.tasks: assert isinstance(scan, ParquetScanTask) assert isinstance(scan.base_task, SplitScan) assert scan.row_groups_and_slice() is None @@ -586,7 +586,7 @@ def test_attach_cached_parquet_metadata_leaves_sliced_fused_scan_unaligned( cached = prefetch_parquet_file_metadata_for_ir(streaming_scan, None) attach_cached_parquet_metadata(streaming_scan, cached) - for scan in streaming_scan.scans: + for scan in streaming_scan.tasks: assert isinstance(scan, ParquetScanTask) assert isinstance(scan.base_task, FusedScan) assert scan.row_groups_and_slice() is None @@ -679,26 +679,6 @@ def test_streaming_scan_missing_prefetch_metadata_raises() -> None: StreamingScan.do_evaluate([fused], scan, context=ctx) -def test_split_scan_do_evaluate_missing_prefetch_metadata() -> None: - paths = ["/some/missing/file.parquet"] - parquet_options = ParquetOptions(prefetch_file_metadata=True) - context = IRExecutionContext() - - with pytest.raises( - AssertionError, - match=(r"Paths do not match cached parquet info."), - ): - SplitScan.do_evaluate( - _make_parquet_scan(paths, parquet_options), - paths, - 0, - 4, - parquet_options, - context=context, - cached_parquet_info=[], - ) - - def test_prefetch_file_metadata_join( tmp_path: Path, streaming_engine_factory: Callable[..., StreamingEngine] ) -> None: @@ -803,9 +783,9 @@ def test_streaming_scan_identity_equality() -> None: base.parquet_options, ) - a = StreamingScan([split], base, "split") - b = StreamingScan([split_same], base, "split") - c = StreamingScan([split_diff], base, "split") + a = StreamingScan([split], base) + b = StreamingScan([split_same], base) + c = StreamingScan([split_diff], base) assert a == b assert hash(a) == hash(b) From 5358c6f4f185256d6d2217277fa539816ece5fb2 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 3 Sep 2026 12:36:46 -0700 Subject: [PATCH 18/31] revise SplitScanBounds convention a bit --- .../cudf_polars/cudf_polars/streaming/io.py | 51 ++++++++----------- 1 file changed, 22 insertions(+), 29 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 1a22d226a234..a26ccce5d6f5 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -343,32 +343,10 @@ def _read_with_hybrid_scan( return DataFrame(columns, stream=stream).select(list(schema.keys())) -def _split_scan_row_groups_and_slice( - scan: SplitScan, - cached_parquet_info: list[CachedParquetInfo], -) -> tuple[list[list[int]], int, int] | None: - """Return the row groups and row slice for a row-group-aligned split.""" - if len(cached_parquet_info) != 1: - return None - - row_group_num_rows = cached_parquet_info[0].file_metadata.row_group_num_rows - total_row_groups = len(row_group_num_rows) - if scan.total_splits > total_row_groups: - return None - - bounds = _split_scan_bounds(row_group_num_rows, scan.split_index, scan.total_splits) - row_groups = list(range(bounds.row_group_start, bounds.row_group_stop)) - if not row_groups: - return None - - return [row_groups], bounds.skip_rows, bounds.n_rows - - -class _SplitScanBounds(NamedTuple): +class SplitScanBounds(NamedTuple): """Row and row-group bounds for a split scan.""" - row_group_start: int - row_group_stop: int + row_groups: list[list[int]] skip_rows: int n_rows: int @@ -377,7 +355,7 @@ def _split_scan_bounds( row_group_num_rows: Sequence[int], split_index: int, total_splits: int, -) -> _SplitScanBounds: +) -> SplitScanBounds: """Return row and row-group bounds for a file split.""" total_row_groups = len(row_group_num_rows) if total_splits <= total_row_groups: @@ -390,16 +368,16 @@ def _split_scan_bounds( ) skip_rows = sum(row_group_num_rows[:row_group_start]) n_rows = sum(row_group_num_rows[row_group_start:row_group_stop]) + row_groups = [list(range(row_group_start, row_group_stop))] else: - row_group_start = 0 - row_group_stop = 0 + row_groups = [] total_rows = sum(row_group_num_rows) n_rows = total_rows // total_splits skip_rows = n_rows * split_index if split_index == total_splits - 1: n_rows = -1 - return _SplitScanBounds(row_group_start, row_group_stop, skip_rows, n_rows) + return SplitScanBounds(row_groups, skip_rows, n_rows) def _cached_parquet_info_for_paths( @@ -720,7 +698,22 @@ def _row_groups_and_slice( return None if isinstance(base_task, SplitScan): - return _split_scan_row_groups_and_slice(base_task, cached_parquet_info) + if len(cached_parquet_info) != 1: + return None + + row_group_num_rows = cached_parquet_info[0].file_metadata.row_group_num_rows + if base_task.total_splits > len(row_group_num_rows): + return None + + bounds = _split_scan_bounds( + row_group_num_rows, + base_task.split_index, + base_task.total_splits, + ) + if not bounds.row_groups: + return None + + return bounds.row_groups, bounds.skip_rows, bounds.n_rows row_groups = [ list(range(len(info.file_metadata.row_group_num_rows))) From 30eda0d1f1c005e4027368a835fe54e593d1dc8e Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 3 Sep 2026 13:20:27 -0700 Subject: [PATCH 19/31] heavier cleanup --- .../cudf_polars/cudf_polars/streaming/io.py | 280 ++++++------------ .../cudf_polars/tests/streaming/test_scan.py | 14 +- 2 files changed, 105 insertions(+), 189 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index a26ccce5d6f5..a1fb8276e8dc 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -343,8 +343,8 @@ def _read_with_hybrid_scan( return DataFrame(columns, stream=stream).select(list(schema.keys())) -class SplitScanBounds(NamedTuple): - """Row and row-group bounds for a split scan.""" +class ParquetTaskBounds(NamedTuple): + """Read bounds for a parquet task.""" row_groups: list[list[int]] skip_rows: int @@ -355,7 +355,7 @@ def _split_scan_bounds( row_group_num_rows: Sequence[int], split_index: int, total_splits: int, -) -> SplitScanBounds: +) -> ParquetTaskBounds: """Return row and row-group bounds for a file split.""" total_row_groups = len(row_group_num_rows) if total_splits <= total_row_groups: @@ -377,7 +377,7 @@ def _split_scan_bounds( if split_index == total_splits - 1: n_rows = -1 - return SplitScanBounds(row_groups, skip_rows, n_rows) + return ParquetTaskBounds(row_groups, skip_rows, n_rows) def _cached_parquet_info_for_paths( @@ -398,14 +398,7 @@ def _cached_parquet_info_for_paths( class SplitScan(IR): - """ - Input from a split file. - - This class wraps a single-file ``Scan`` object. At - IO/evaluation time, this class will only perform - a partial read of the underlying file. The range - (skip_rows and n_rows) is calculated at IO time. - """ + """Streaming task describing one split of a parquet file.""" __slots__ = ( "base_scan", @@ -473,63 +466,6 @@ def get_hashable(self) -> Hashable: self.parquet_options, ) - @classmethod - def do_evaluate( - cls, - base_scan: Scan, - paths: list[str], - split_index: int, - total_splits: int, - parquet_options: ParquetOptions, - *, - context: IRExecutionContext, - ) -> DataFrame: - """Evaluate and return a dataframe.""" - if base_scan.typ not in ("parquet",): # pragma: no cover - raise NotImplementedError( - f"Unhandled Scan type for file splitting: {base_scan.typ}" - ) - - if len(paths) > 1: # pragma: no cover - raise ValueError(f"Expected a single path, got: {paths}") - - cached_parquet_info = _cached_parquet_info_for_paths(base_scan, paths) - if cached_parquet_info is not None: - parquet_metadatas = [info.file_metadata for info in cached_parquet_info] - row_group_num_rows = [ - num_rows - for metadata in parquet_metadatas - for num_rows in metadata.row_group_num_rows - ] - else: - row_group_num_rows = [ - rg["num_rows"] - for rg in plc.io.parquet_metadata.read_parquet_metadata( - plc.io.SourceInfo(paths) - ).rowgroup_metadata() - ] - bounds = _split_scan_bounds(row_group_num_rows, split_index, total_splits) - - # Perform the partial read - with nvtx_annotate_cudf_polars( - message=f"SplitScan: {paths[0]} [{split_index + 1}/{total_splits}]" - ): - return Scan.do_evaluate( - base_scan.schema, - base_scan.typ, - base_scan.reader_options, - paths, - base_scan.with_columns, - bounds.skip_rows, - bounds.n_rows, - base_scan.row_index, - base_scan.include_file_paths, - base_scan.predicate, - parquet_options, - cached_parquet_info, - context=context, - ) - class FusedScan(IR): """ @@ -614,37 +550,12 @@ def do_evaluate( ) -def _evaluate_scan_task( - scan: SplitScan | FusedScan, - *, - context: IRExecutionContext, -) -> DataFrame: - """Evaluate a generic streaming scan task.""" - base_scan = scan.base_scan - if isinstance(scan, SplitScan): - return SplitScan.do_evaluate( - base_scan, - scan.paths, - scan.split_index, - scan.total_splits, - scan.parquet_options, - context=context, - ) - return FusedScan.do_evaluate( - base_scan, - scan.paths, - scan.parquet_options, - context=context, - ) - - class ParquetScanTask(IR): """Parquet scan task wrapping a generic streaming scan task.""" __slots__ = ( "base_scan", "base_task", - "parquet_options", "paths", "schema", ) @@ -657,8 +568,6 @@ class ParquetScanTask(IR): """Scan operation this task is based on.""" paths: list[str] """File paths assigned to this task.""" - parquet_options: ParquetOptions - """Parquet-specific options.""" def __init__( self, @@ -670,56 +579,61 @@ def __init__( self.base_task = base_task self.base_scan = base_scan self.paths = base_task.paths - self.parquet_options = base_task.parquet_options self.schema = base_scan.schema self._non_child_args = (base_task,) self.children = () - def row_groups_and_slice(self) -> tuple[list[list[int]], int, int] | None: - """Return row groups and row slice when this task is row-group aligned.""" - return self._row_groups_and_slice( + def task_bounds(self) -> ParquetTaskBounds | None: + """Return parquet read bounds for this task.""" + return self._get_task_bounds( self.base_task, _cached_parquet_info_for_paths(self.base_scan, self.paths), + fetch_missing_metadata=False, ) @staticmethod - def _row_groups_and_slice( + def _get_task_bounds( base_task: SplitScan | FusedScan, cached_parquet_info: list[CachedParquetInfo] | None, - ) -> tuple[list[list[int]], int, int] | None: + *, + fetch_missing_metadata: bool, + ) -> ParquetTaskBounds | None: base_scan = base_task.base_scan - if cached_parquet_info is None: - return None - if ( - base_scan.skip_rows != 0 - or base_scan.n_rows != -1 - or base_scan.row_index is not None - ): - return None if isinstance(base_task, SplitScan): - if len(cached_parquet_info) != 1: - return None - - row_group_num_rows = cached_parquet_info[0].file_metadata.row_group_num_rows - if base_task.total_splits > len(row_group_num_rows): + if len(base_task.paths) > 1: # pragma: no cover + raise ValueError(f"Expected a single path, got: {base_task.paths}") + if cached_parquet_info is not None: + row_group_num_rows = cached_parquet_info[ + 0 + ].file_metadata.row_group_num_rows + elif fetch_missing_metadata: + row_group_num_rows = [ + rg["num_rows"] + for rg in plc.io.parquet_metadata.read_parquet_metadata( + plc.io.SourceInfo(base_task.paths) + ).rowgroup_metadata() + ] + else: return None - - bounds = _split_scan_bounds( + return _split_scan_bounds( row_group_num_rows, base_task.split_index, base_task.total_splits, ) - if not bounds.row_groups: - return None - - return bounds.row_groups, bounds.skip_rows, bounds.n_rows - row_groups = [ - list(range(len(info.file_metadata.row_group_num_rows))) - for info in cached_parquet_info - ] - return row_groups, 0, -1 + row_groups: list[list[int]] = [] + if ( + cached_parquet_info is not None + and base_scan.skip_rows == 0 + and base_scan.n_rows == -1 + and base_scan.row_index is None + ): + row_groups = [ + list(range(len(info.file_metadata.row_group_num_rows))) + for info in cached_parquet_info + ] + return ParquetTaskBounds(row_groups, base_scan.skip_rows, base_scan.n_rows) def get_hashable(self) -> Hashable: """Hashable representation of the node.""" @@ -740,70 +654,68 @@ def do_evaluate( paths = base_task.paths parquet_options = base_task.parquet_options cached_parquet_info = _cached_parquet_info_for_paths(base_scan, paths) - row_group_info = cls._row_groups_and_slice(base_task, cached_parquet_info) - - if row_group_info is not None: - row_groups, skip_rows, n_rows = row_group_info - # 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 ( - len(paths) == 1 - and len(row_groups) == 1 - and hybrid_scan_eligible( - parquet_options, - cached_parquet_info=cached_parquet_info, - row_index=base_scan.row_index, - include_file_paths=base_scan.include_file_paths, - predicate=base_scan.predicate, - ) - ): - assert base_scan.predicate is not None - assert cached_parquet_info is not None - stream = context.get_cuda_stream() - plc_filter, residual = to_parquet_filter( - _prepare_parquet_predicate( - base_scan.predicate.value, - paths, - base_scan.schema, - base_scan.with_columns, - ), - stream=stream, - ) - if plc_filter is not None and residual is None: - return _read_with_hybrid_scan( - base_scan.schema, - paths, - base_scan.with_columns, - plc_filter, - row_groups[0], - stream, - cached_parquet_info[0], - stats_pruning=parquet_options._hybrid_scan_stats_pruning, - ) + bounds = cls._get_task_bounds( + base_task, + cached_parquet_info, + fetch_missing_metadata=True, + ) - with nvtx_annotate_cudf_polars(message=f"ParquetScan: {', '.join(paths)}"): - return Scan.do_evaluate( + assert bounds is not None + # 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 ( + len(paths) == 1 + and len(bounds.row_groups) == 1 + and hybrid_scan_eligible( + parquet_options, + cached_parquet_info=cached_parquet_info, + row_index=base_scan.row_index, + include_file_paths=base_scan.include_file_paths, + predicate=base_scan.predicate, + ) + ): + assert base_scan.predicate is not None + assert cached_parquet_info is not None + stream = context.get_cuda_stream() + plc_filter, residual = to_parquet_filter( + _prepare_parquet_predicate( + base_scan.predicate.value, + paths, + base_scan.schema, + base_scan.with_columns, + ), + stream=stream, + ) + if plc_filter is not None and residual is None: + return _read_with_hybrid_scan( base_scan.schema, - base_scan.typ, - base_scan.reader_options, paths, base_scan.with_columns, - skip_rows, - n_rows, - base_scan.row_index, - base_scan.include_file_paths, - base_scan.predicate, - parquet_options, - cached_parquet_info, - context=context, + plc_filter, + bounds.row_groups[0], + stream, + cached_parquet_info[0], + stats_pruning=parquet_options._hybrid_scan_stats_pruning, ) - return _evaluate_scan_task( - base_task, - context=context, - ) + with nvtx_annotate_cudf_polars(message=f"ParquetScan: {', '.join(paths)}"): + return Scan.do_evaluate( + base_scan.schema, + base_scan.typ, + base_scan.reader_options, + paths, + base_scan.with_columns, + bounds.skip_rows, + bounds.n_rows, + base_scan.row_index, + base_scan.include_file_paths, + base_scan.predicate, + parquet_options, + cached_parquet_info, + context=context, + ) StreamingScanTask: TypeAlias = SplitScan | FusedScan | ParquetScanTask diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 33caef738b18..ca5263912cf6 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -527,9 +527,9 @@ def test_attach_cached_parquet_metadata_resolves_row_groups( row_groups = [] for scan in streaming_scan.tasks: assert isinstance(scan, ParquetScanTask) - info = scan.row_groups_and_slice() - assert info is not None - row_groups.append(info[0]) + bounds = scan.task_bounds() + assert bounds is not None + row_groups.append(bounds.row_groups) assert row_groups == [[[0]], [[1]]] @@ -555,7 +555,9 @@ def test_attach_cached_parquet_metadata_leaves_sub_row_group_split_unaligned( for scan in streaming_scan.tasks: assert isinstance(scan, ParquetScanTask) assert isinstance(scan.base_task, SplitScan) - assert scan.row_groups_and_slice() is None + bounds = scan.task_bounds() + assert bounds is not None + assert bounds.row_groups == [] @pytest.mark.parametrize( @@ -589,7 +591,9 @@ def test_attach_cached_parquet_metadata_leaves_sliced_fused_scan_unaligned( for scan in streaming_scan.tasks: assert isinstance(scan, ParquetScanTask) assert isinstance(scan.base_task, FusedScan) - assert scan.row_groups_and_slice() is None + bounds = scan.task_bounds() + assert bounds is not None + assert bounds.row_groups == [] def test_streaming_scan_raises() -> None: From e165d5e908e283b5aedb4ad04360a10b55885592 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 3 Sep 2026 13:32:56 -0700 Subject: [PATCH 20/31] make _split_scan_bounds a method --- .../cudf_polars/cudf_polars/streaming/io.py | 94 +++++++++---------- 1 file changed, 46 insertions(+), 48 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index a1fb8276e8dc..6c50a92c0ac7 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -351,35 +351,6 @@ class ParquetTaskBounds(NamedTuple): n_rows: int -def _split_scan_bounds( - row_group_num_rows: Sequence[int], - split_index: int, - total_splits: int, -) -> ParquetTaskBounds: - """Return row and row-group bounds for a file split.""" - total_row_groups = len(row_group_num_rows) - if total_splits <= total_row_groups: - row_group_stride = total_row_groups // total_splits - row_group_start = row_group_stride * split_index - row_group_stop = ( - total_row_groups - if split_index == total_splits - 1 - else row_group_start + row_group_stride - ) - skip_rows = sum(row_group_num_rows[:row_group_start]) - n_rows = sum(row_group_num_rows[row_group_start:row_group_stop]) - row_groups = [list(range(row_group_start, row_group_stop))] - else: - row_groups = [] - total_rows = sum(row_group_num_rows) - n_rows = total_rows // total_splits - skip_rows = n_rows * split_index - - if split_index == total_splits - 1: - n_rows = -1 - return ParquetTaskBounds(row_groups, skip_rows, n_rows) - - def _cached_parquet_info_for_paths( base_scan: Scan, paths: list[str], @@ -454,6 +425,49 @@ def __init__( f"Unhandled Scan type for file splitting: {base_scan.typ}" ) + def _parquet_task_bounds( + self, + cached_parquet_info: list[CachedParquetInfo] | None, + *, + fetch_missing_metadata: bool, + ) -> ParquetTaskBounds | None: + """Return parquet read bounds for this split task.""" + if len(self.paths) > 1: # pragma: no cover + raise ValueError(f"Expected a single path, got: {self.paths}") + if cached_parquet_info is not None: + row_group_num_rows = cached_parquet_info[0].file_metadata.row_group_num_rows + elif fetch_missing_metadata: + row_group_num_rows = [ + rg["num_rows"] + for rg in plc.io.parquet_metadata.read_parquet_metadata( + plc.io.SourceInfo(self.paths) + ).rowgroup_metadata() + ] + else: + return None + + total_row_groups = len(row_group_num_rows) + if self.total_splits <= total_row_groups: + row_group_stride = total_row_groups // self.total_splits + row_group_start = row_group_stride * self.split_index + row_group_stop = ( + total_row_groups + if self.split_index == self.total_splits - 1 + else row_group_start + row_group_stride + ) + skip_rows = sum(row_group_num_rows[:row_group_start]) + n_rows = sum(row_group_num_rows[row_group_start:row_group_stop]) + row_groups = [list(range(row_group_start, row_group_stop))] + else: + row_groups = [] + total_rows = sum(row_group_num_rows) + n_rows = total_rows // self.total_splits + skip_rows = n_rows * self.split_index + + if self.split_index == self.total_splits - 1: + n_rows = -1 + return ParquetTaskBounds(row_groups, skip_rows, n_rows) + def get_hashable(self) -> Hashable: """Hashable representation of the node.""" return ( @@ -601,25 +615,9 @@ def _get_task_bounds( base_scan = base_task.base_scan if isinstance(base_task, SplitScan): - if len(base_task.paths) > 1: # pragma: no cover - raise ValueError(f"Expected a single path, got: {base_task.paths}") - if cached_parquet_info is not None: - row_group_num_rows = cached_parquet_info[ - 0 - ].file_metadata.row_group_num_rows - elif fetch_missing_metadata: - row_group_num_rows = [ - rg["num_rows"] - for rg in plc.io.parquet_metadata.read_parquet_metadata( - plc.io.SourceInfo(base_task.paths) - ).rowgroup_metadata() - ] - else: - return None - return _split_scan_bounds( - row_group_num_rows, - base_task.split_index, - base_task.total_splits, + return base_task._parquet_task_bounds( + cached_parquet_info, + fetch_missing_metadata=fetch_missing_metadata, ) row_groups: list[list[int]] = [] From 7f6d5e38c9ca9c50f6765469ab96ebbf8c9e0645 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 3 Sep 2026 14:22:05 -0700 Subject: [PATCH 21/31] docstring cleanup --- .../cudf_polars/cudf_polars/streaming/io.py | 99 +++++++++++-------- .../cudf_polars/tests/streaming/test_scan.py | 10 +- 2 files changed, 65 insertions(+), 44 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 6c50a92c0ac7..72ae3d2283b2 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -344,32 +344,24 @@ def _read_with_hybrid_scan( class ParquetTaskBounds(NamedTuple): - """Read bounds for a parquet task.""" + """ + Read bounds for a parquet task. + + ``row_groups=None`` means the task is not row-group aligned. + """ - row_groups: list[list[int]] + row_groups: list[list[int]] | None skip_rows: int n_rows: int -def _cached_parquet_info_for_paths( - base_scan: Scan, - paths: list[str], -) -> list[CachedParquetInfo] | None: - """Return cached parquet metadata matching a scan task's paths.""" - cached_parquet_info = base_scan.cached_parquet_info - if cached_parquet_info is None or cached_parquet_info == []: - return None - if paths == [info.path for info in cached_parquet_info]: - return cached_parquet_info - - cached_by_path = {info.path: info for info in cached_parquet_info} - if not all(path in cached_by_path for path in paths): - return None - return [cached_by_path[path] for path in paths] - - class SplitScan(IR): - """Streaming task describing one split of a parquet file.""" + """ + Streaming task describing one SPLIT_FILES partition of a file. + + This class only describes the split bounds. Parquet execution wraps it in + ``ParquetScanTask``. + """ __slots__ = ( "base_scan", @@ -459,7 +451,7 @@ def _parquet_task_bounds( n_rows = sum(row_group_num_rows[row_group_start:row_group_stop]) row_groups = [list(range(row_group_start, row_group_stop))] else: - row_groups = [] + row_groups = None total_rows = sum(row_group_num_rows) n_rows = total_rows // self.total_splits skip_rows = n_rows * self.split_index @@ -545,7 +537,6 @@ def do_evaluate( context: IRExecutionContext, ) -> DataFrame: """Evaluate and return a dataframe.""" - cached_parquet_info = _cached_parquet_info_for_paths(base_scan, paths) with nvtx_annotate_cudf_polars(message=f"FusedScan: {', '.join(paths)}"): return Scan.do_evaluate( base_scan.schema, @@ -559,7 +550,7 @@ def do_evaluate( base_scan.include_file_paths, base_scan.predicate, parquet_options, - cached_parquet_info, + None, context=context, ) @@ -597,22 +588,36 @@ def __init__( self._non_child_args = (base_task,) self.children = () - def task_bounds(self) -> ParquetTaskBounds | None: + def get_task_bounds( + self, *, fetch_missing_metadata: bool = False + ) -> ParquetTaskBounds | None: """Return parquet read bounds for this task.""" - return self._get_task_bounds( - self.base_task, - _cached_parquet_info_for_paths(self.base_scan, self.paths), - fetch_missing_metadata=False, + return self._task_bounds_from_cached( + self._cached_parquet_info(), + fetch_missing_metadata=fetch_missing_metadata, ) - @staticmethod - def _get_task_bounds( - base_task: SplitScan | FusedScan, + def _cached_parquet_info(self) -> list[CachedParquetInfo] | None: + """Return cached parquet metadata matching this task's paths.""" + cached_parquet_info = self.base_scan.cached_parquet_info + if cached_parquet_info is None or cached_parquet_info == []: + return None + if self.paths == [info.path for info in cached_parquet_info]: + return cached_parquet_info + + cached_by_path = {info.path: info for info in cached_parquet_info} + if not all(path in cached_by_path for path in self.paths): + return None + return [cached_by_path[path] for path in self.paths] + + def _task_bounds_from_cached( + self, cached_parquet_info: list[CachedParquetInfo] | None, *, fetch_missing_metadata: bool, ) -> ParquetTaskBounds | None: - base_scan = base_task.base_scan + base_task = self.base_task + base_scan = self.base_scan if isinstance(base_task, SplitScan): return base_task._parquet_task_bounds( @@ -620,7 +625,7 @@ def _get_task_bounds( fetch_missing_metadata=fetch_missing_metadata, ) - row_groups: list[list[int]] = [] + row_groups: list[list[int]] | None = None if ( cached_parquet_info is not None and base_scan.skip_rows == 0 @@ -648,12 +653,12 @@ def do_evaluate( context: IRExecutionContext, ) -> DataFrame: """Evaluate a parquet scan task.""" - base_scan = base_task.base_scan - paths = base_task.paths + task = cls(base_task) + base_scan = task.base_scan + paths = task.paths parquet_options = base_task.parquet_options - cached_parquet_info = _cached_parquet_info_for_paths(base_scan, paths) - bounds = cls._get_task_bounds( - base_task, + cached_parquet_info = task._cached_parquet_info() + bounds = task._task_bounds_from_cached( cached_parquet_info, fetch_missing_metadata=True, ) @@ -665,6 +670,7 @@ def do_evaluate( # (row_index / include_file_paths). Needs performance investigation. if ( len(paths) == 1 + and bounds.row_groups is not None and len(bounds.row_groups) == 1 and hybrid_scan_eligible( parquet_options, @@ -687,6 +693,12 @@ def do_evaluate( stream=stream, ) if plc_filter is not None and residual is None: + if isinstance(base_task, SplitScan): + split_index = base_task.split_index + total_splits = base_task.total_splits + else: + split_index = 0 + total_splits = 1 return _read_with_hybrid_scan( base_scan.schema, paths, @@ -695,10 +707,19 @@ def do_evaluate( bounds.row_groups[0], stream, cached_parquet_info[0], + split_index=split_index, + total_splits=total_splits, stats_pruning=parquet_options._hybrid_scan_stats_pruning, ) - with nvtx_annotate_cudf_polars(message=f"ParquetScan: {', '.join(paths)}"): + if isinstance(base_task, SplitScan): + nvtx_message = ( + f"SplitScan: {paths[0]} " + f"[{base_task.split_index + 1}/{base_task.total_splits}]" + ) + else: + nvtx_message = f"FusedScan: {', '.join(paths)}" + with nvtx_annotate_cudf_polars(message=nvtx_message): return Scan.do_evaluate( base_scan.schema, base_scan.typ, diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index ca5263912cf6..0b804d2395b2 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -527,7 +527,7 @@ def test_attach_cached_parquet_metadata_resolves_row_groups( row_groups = [] for scan in streaming_scan.tasks: assert isinstance(scan, ParquetScanTask) - bounds = scan.task_bounds() + bounds = scan.get_task_bounds() assert bounds is not None row_groups.append(bounds.row_groups) assert row_groups == [[[0]], [[1]]] @@ -555,9 +555,9 @@ def test_attach_cached_parquet_metadata_leaves_sub_row_group_split_unaligned( for scan in streaming_scan.tasks: assert isinstance(scan, ParquetScanTask) assert isinstance(scan.base_task, SplitScan) - bounds = scan.task_bounds() + bounds = scan.get_task_bounds() assert bounds is not None - assert bounds.row_groups == [] + assert bounds.row_groups is None @pytest.mark.parametrize( @@ -591,9 +591,9 @@ def test_attach_cached_parquet_metadata_leaves_sliced_fused_scan_unaligned( for scan in streaming_scan.tasks: assert isinstance(scan, ParquetScanTask) assert isinstance(scan.base_task, FusedScan) - bounds = scan.task_bounds() + bounds = scan.get_task_bounds() assert bounds is not None - assert bounds.row_groups == [] + assert bounds.row_groups is None def test_streaming_scan_raises() -> None: From 518ed51a8e8954f4156c80f63568901a39ac10f4 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 3 Sep 2026 14:53:33 -0700 Subject: [PATCH 22/31] reuse the metadata we prefatch at do_evaluate time --- .../cudf_polars/cudf_polars/streaming/io.py | 50 +++++++------------ 1 file changed, 19 insertions(+), 31 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 72ae3d2283b2..dbb6ac3d0084 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -420,24 +420,14 @@ def __init__( def _parquet_task_bounds( self, cached_parquet_info: list[CachedParquetInfo] | None, - *, - fetch_missing_metadata: bool, ) -> ParquetTaskBounds | None: """Return parquet read bounds for this split task.""" if len(self.paths) > 1: # pragma: no cover raise ValueError(f"Expected a single path, got: {self.paths}") - if cached_parquet_info is not None: - row_group_num_rows = cached_parquet_info[0].file_metadata.row_group_num_rows - elif fetch_missing_metadata: - row_group_num_rows = [ - rg["num_rows"] - for rg in plc.io.parquet_metadata.read_parquet_metadata( - plc.io.SourceInfo(self.paths) - ).rowgroup_metadata() - ] - else: + if cached_parquet_info is None: return None + row_group_num_rows = cached_parquet_info[0].file_metadata.row_group_num_rows total_row_groups = len(row_group_num_rows) if self.total_splits <= total_row_groups: row_group_stride = total_row_groups // self.total_splits @@ -588,14 +578,9 @@ def __init__( self._non_child_args = (base_task,) self.children = () - def get_task_bounds( - self, *, fetch_missing_metadata: bool = False - ) -> ParquetTaskBounds | None: + def get_task_bounds(self) -> ParquetTaskBounds | None: """Return parquet read bounds for this task.""" - return self._task_bounds_from_cached( - self._cached_parquet_info(), - fetch_missing_metadata=fetch_missing_metadata, - ) + return self._task_bounds_from_cached(self._cached_parquet_info()) def _cached_parquet_info(self) -> list[CachedParquetInfo] | None: """Return cached parquet metadata matching this task's paths.""" @@ -610,20 +595,24 @@ def _cached_parquet_info(self) -> list[CachedParquetInfo] | None: return None return [cached_by_path[path] for path in self.paths] + def _fetch_parquet_info(self) -> list[CachedParquetInfo]: + """Fetch parquet metadata for this task's paths.""" + from cudf_polars.dsl.utils.io import _prefetch_parquet_footers_for_paths + + return _prefetch_parquet_footers_for_paths( + self.paths, + parse_hybrid_metadata=self.base_task.parquet_options.use_hybrid_scan, + ) + def _task_bounds_from_cached( self, cached_parquet_info: list[CachedParquetInfo] | None, - *, - fetch_missing_metadata: bool, ) -> ParquetTaskBounds | None: base_task = self.base_task base_scan = self.base_scan if isinstance(base_task, SplitScan): - return base_task._parquet_task_bounds( - cached_parquet_info, - fetch_missing_metadata=fetch_missing_metadata, - ) + return base_task._parquet_task_bounds(cached_parquet_info) row_groups: list[list[int]] | None = None if ( @@ -658,14 +647,13 @@ def do_evaluate( paths = task.paths parquet_options = base_task.parquet_options cached_parquet_info = task._cached_parquet_info() - bounds = task._task_bounds_from_cached( - cached_parquet_info, - fetch_missing_metadata=True, - ) + if cached_parquet_info is None and isinstance(base_task, SplitScan): + cached_parquet_info = task._fetch_parquet_info() + bounds = task._task_bounds_from_cached(cached_parquet_info) assert bounds is not None - # Hybrid scan reads through the prefetched, shared file metadata, so - # it is only used when footer prefetching is enabled. + # Hybrid scan reads through cached parquet metadata, so it is only used + # when the metadata is available to this task. # TODO: Investigate re-enabling for some of the excluded paths # (row_index / include_file_paths). Needs performance investigation. if ( From 6e1d01c65a86909980a021a3af79084cba5d4c02 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 8 Sep 2026 07:19:25 -0700 Subject: [PATCH 23/31] push staged revisions - dropping nesting --- .../cudf_polars/streaming/actor_graph/io.py | 10 +- .../cudf_polars/cudf_polars/streaming/io.py | 288 ++++++++---------- .../cudf_polars/cudf_polars/utils/config.py | 2 +- .../cudf_polars/tests/streaming/test_scan.py | 27 +- 4 files changed, 147 insertions(+), 180 deletions(-) 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 acd60d0dcd1e..71d6f8ebc819 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -570,15 +570,19 @@ async def read_chunk( br=context.br(), ) stop = time.monotonic_ns() - trace_task = task.base_task if isinstance(task, ParquetScanTask) else task + ir_type = ( + task.trace_ir_type() + if isinstance(task, ParquetScanTask) + else type(task).__name__ + ) log( "IO Task", scope=Scope.IO_TASK.value, start=start, admitted=admitted, stop=stop, - ir_id=trace_task.get_stable_id(), - ir_type=type(trace_task).__name__, + ir_id=task.get_stable_id(), + ir_type=ir_type, sequence_number=seq_num, estimated_output_bytes=estimated_chunk_bytes, reservation_bytes=reservation_bytes, diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index dbb6ac3d0084..b9de833d0178 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -88,8 +88,8 @@ def scan_partition_plan( 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. + # A single file always uses a split parquet task when hybrid scan is enabled, + # so the hybrid reader can be used even when the file 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 @@ -285,10 +285,10 @@ def _read_with_hybrid_scan( 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. + # TODO: Consider implementing page-index stats pruning. For split tasks, 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 fused tasks, 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 = plc.io.parquet_io_utils.fetch_byte_ranges_to_device( @@ -355,114 +355,6 @@ class ParquetTaskBounds(NamedTuple): n_rows: int -class SplitScan(IR): - """ - Streaming task describing one SPLIT_FILES partition of a file. - - This class only describes the split bounds. Parquet execution wraps it in - ``ParquetScanTask``. - """ - - __slots__ = ( - "base_scan", - "parquet_options", - "paths", - "schema", - "split_index", - "total_splits", - ) - _non_child = ( - "base_scan", - "paths", - "split_index", - "total_splits", - "parquet_options", - ) - _n_non_child_args = 5 - base_scan: Scan - """Scan operation this node is based on.""" - paths: list[str] - """File path for this split task.""" - split_index: int - """Index of the current split.""" - total_splits: int - """Total number of splits.""" - parquet_options: ParquetOptions - """Parquet-specific options.""" - - def __init__( - self, - base_scan: Scan, - paths: list[str], - split_index: int, - total_splits: int, - parquet_options: ParquetOptions, - ): - self.schema = base_scan.schema - self.base_scan = base_scan - self.paths = paths - self.split_index = split_index - self.total_splits = total_splits - self._non_child_args = ( - base_scan, - paths, - split_index, - total_splits, - parquet_options, - ) - self.parquet_options = parquet_options - self.children = () - if base_scan.typ not in ("parquet",): # pragma: no cover - raise NotImplementedError( - f"Unhandled Scan type for file splitting: {base_scan.typ}" - ) - - def _parquet_task_bounds( - self, - cached_parquet_info: list[CachedParquetInfo] | None, - ) -> ParquetTaskBounds | None: - """Return parquet read bounds for this split task.""" - if len(self.paths) > 1: # pragma: no cover - raise ValueError(f"Expected a single path, got: {self.paths}") - if cached_parquet_info is None: - return None - - row_group_num_rows = cached_parquet_info[0].file_metadata.row_group_num_rows - total_row_groups = len(row_group_num_rows) - if self.total_splits <= total_row_groups: - row_group_stride = total_row_groups // self.total_splits - row_group_start = row_group_stride * self.split_index - row_group_stop = ( - total_row_groups - if self.split_index == self.total_splits - 1 - else row_group_start + row_group_stride - ) - skip_rows = sum(row_group_num_rows[:row_group_start]) - n_rows = sum(row_group_num_rows[row_group_start:row_group_stop]) - row_groups = [list(range(row_group_start, row_group_stop))] - else: - row_groups = None - total_rows = sum(row_group_num_rows) - n_rows = total_rows // self.total_splits - skip_rows = n_rows * self.split_index - - if self.split_index == self.total_splits - 1: - n_rows = -1 - return ParquetTaskBounds(row_groups, skip_rows, n_rows) - - def get_hashable(self) -> Hashable: - """Hashable representation of the node.""" - return ( - type(self), - tuple(self.schema.items()), - self.base_scan.get_hashable(), - tuple(self.paths), - self.split_index, - self.total_splits, - self.parquet_options, - ) - - class FusedScan(IR): """ Input from one or more complete files read as a single task. @@ -546,38 +438,70 @@ def do_evaluate( class ParquetScanTask(IR): - """Parquet scan task wrapping a generic streaming scan task.""" + """Parquet-specific streaming scan task.""" __slots__ = ( "base_scan", - "base_task", + "parquet_options", "paths", "schema", + "split_index", + "total_splits", ) - _non_child = ("base_task",) - _n_non_child_args = 1 + _non_child = ( + "base_scan", + "paths", + "split_index", + "total_splits", + "parquet_options", + ) + _n_non_child_args = 5 - base_task: SplitScan | FusedScan - """Generic streaming scan task being specialized for parquet.""" base_scan: Scan """Scan operation this task is based on.""" paths: list[str] """File paths assigned to this task.""" + split_index: int | None + """Index of the current split, or None for fused/full-file tasks.""" + total_splits: int | None + """Total number of splits for this file, or None for fused/full-file tasks.""" + parquet_options: ParquetOptions + """Parquet-specific options.""" def __init__( self, - base_task: SplitScan | FusedScan, + base_scan: Scan, + paths: list[str], + split_index: int | None, + total_splits: int | None, + parquet_options: ParquetOptions, ): - base_scan = base_task.base_scan if base_scan.typ != "parquet": # pragma: no cover raise ValueError(f"Expected a parquet scan, got: {base_scan.typ}") - self.base_task = base_task + if (split_index is None) != (total_splits is None): # pragma: no cover + raise ValueError("split_index and total_splits must be set together") + if split_index is not None and len(paths) > 1: # pragma: no cover + raise ValueError(f"Expected a single path for a split task, got: {paths}") self.base_scan = base_scan - self.paths = base_task.paths + self.paths = paths + self.split_index = split_index + self.total_splits = total_splits + self.parquet_options = parquet_options self.schema = base_scan.schema - self._non_child_args = (base_task,) + self._non_child_args = ( + base_scan, + paths, + split_index, + total_splits, + parquet_options, + ) self.children = () + @property + def is_split(self) -> bool: + """Whether this task is one split of a single parquet file.""" + return self.split_index is not None + def get_task_bounds(self) -> ParquetTaskBounds | None: """Return parquet read bounds for this task.""" return self._task_bounds_from_cached(self._cached_parquet_info()) @@ -601,19 +525,50 @@ def _fetch_parquet_info(self) -> list[CachedParquetInfo]: return _prefetch_parquet_footers_for_paths( self.paths, - parse_hybrid_metadata=self.base_task.parquet_options.use_hybrid_scan, + parse_hybrid_metadata=self.parquet_options.use_hybrid_scan, ) - def _task_bounds_from_cached( + def _split_task_bounds( self, cached_parquet_info: list[CachedParquetInfo] | None, ) -> ParquetTaskBounds | None: - base_task = self.base_task - base_scan = self.base_scan + """Return parquet read bounds for a split task.""" + if cached_parquet_info is None: + return None + + assert self.split_index is not None + assert self.total_splits is not None + row_group_num_rows = cached_parquet_info[0].file_metadata.row_group_num_rows + total_row_groups = len(row_group_num_rows) + if self.total_splits <= total_row_groups: + row_group_stride = total_row_groups // self.total_splits + row_group_start = row_group_stride * self.split_index + row_group_stop = ( + total_row_groups + if self.split_index == self.total_splits - 1 + else row_group_start + row_group_stride + ) + skip_rows = sum(row_group_num_rows[:row_group_start]) + n_rows = sum(row_group_num_rows[row_group_start:row_group_stop]) + row_groups = [list(range(row_group_start, row_group_stop))] + else: + row_groups = None + total_rows = sum(row_group_num_rows) + n_rows = total_rows // self.total_splits + skip_rows = n_rows * self.split_index - if isinstance(base_task, SplitScan): - return base_task._parquet_task_bounds(cached_parquet_info) + if self.split_index == self.total_splits - 1: + n_rows = -1 + return ParquetTaskBounds(row_groups, skip_rows, n_rows) + def _task_bounds_from_cached( + self, + cached_parquet_info: list[CachedParquetInfo] | None, + ) -> ParquetTaskBounds | None: + if self.is_split: + return self._split_task_bounds(cached_parquet_info) + + base_scan = self.base_scan row_groups: list[list[int]] | None = None if ( cached_parquet_info is not None @@ -631,23 +586,34 @@ def get_hashable(self) -> Hashable: """Hashable representation of the node.""" return ( type(self), - self.base_task.get_hashable(), + self.base_scan.get_hashable(), + tuple(self.paths), + self.split_index, + self.total_splits, + self.parquet_options, ) + def trace_ir_type(self) -> str: + """Return the task type to use for IO-task tracing.""" + return "SplitScan" if self.is_split else "FusedScan" + @classmethod def do_evaluate( cls, - base_task: SplitScan | FusedScan, + base_scan: Scan, + paths: list[str], + split_index: int | None, + total_splits: int | None, + parquet_options: ParquetOptions, *, context: IRExecutionContext, ) -> DataFrame: """Evaluate a parquet scan task.""" - task = cls(base_task) + task = cls(base_scan, paths, split_index, total_splits, parquet_options) base_scan = task.base_scan paths = task.paths - parquet_options = base_task.parquet_options cached_parquet_info = task._cached_parquet_info() - if cached_parquet_info is None and isinstance(base_task, SplitScan): + if cached_parquet_info is None and task.is_split: cached_parquet_info = task._fetch_parquet_info() bounds = task._task_bounds_from_cached(cached_parquet_info) @@ -681,12 +647,6 @@ def do_evaluate( stream=stream, ) if plc_filter is not None and residual is None: - if isinstance(base_task, SplitScan): - split_index = base_task.split_index - total_splits = base_task.total_splits - else: - split_index = 0 - total_splits = 1 return _read_with_hybrid_scan( base_scan.schema, paths, @@ -695,16 +655,15 @@ def do_evaluate( bounds.row_groups[0], stream, cached_parquet_info[0], - split_index=split_index, - total_splits=total_splits, + split_index=split_index or 0, + total_splits=total_splits or 1, stats_pruning=parquet_options._hybrid_scan_stats_pruning, ) - if isinstance(base_task, SplitScan): - nvtx_message = ( - f"SplitScan: {paths[0]} " - f"[{base_task.split_index + 1}/{base_task.total_splits}]" - ) + if task.is_split: + assert split_index is not None + assert total_splits is not None + nvtx_message = f"SplitScan: {paths[0]} [{split_index + 1}/{total_splits}]" else: nvtx_message = f"FusedScan: {', '.join(paths)}" with nvtx_annotate_cudf_polars(message=nvtx_message): @@ -725,7 +684,7 @@ def do_evaluate( ) -StreamingScanTask: TypeAlias = SplitScan | FusedScan | ParquetScanTask +StreamingScanTask: TypeAlias = FusedScan | ParquetScanTask @lower_ir_node.register(Empty) @@ -813,11 +772,6 @@ def __init__( tasks: Sequence[StreamingScanTask], base_scan: Scan, ): - if base_scan.typ == "parquet": - tasks = [ - task if isinstance(task, ParquetScanTask) else ParquetScanTask(task) - for task in tasks - ] self.base_scan = base_scan self.schema = base_scan.schema self.tasks = tasks @@ -841,12 +795,12 @@ def for_split_files( path_end = math.ceil((local_offset + local_count) / plan.factor) local_paths = base_scan.paths[path_offset:path_end] sindex = local_offset % plan.factor - tasks: list[SplitScan] = [] + tasks: list[StreamingScanTask] = [] splits_created = 0 for path in local_paths: while sindex < plan.factor and splits_created < local_count: tasks.append( - SplitScan( + ParquetScanTask( base_scan, [path], sindex, @@ -874,11 +828,21 @@ def for_fused_files( local_offset, local_count = _rank_slice(partition_count, rank, nranks) paths_start = local_offset * plan.factor paths_end = paths_start + plan.factor * local_count - tasks = [ - FusedScan( - base_scan, - base_scan.paths[offset : offset + plan.factor], - parquet_options, + tasks: list[StreamingScanTask] = [ + ( + ParquetScanTask( + base_scan, + base_scan.paths[offset : offset + plan.factor], + None, + None, + parquet_options, + ) + if base_scan.typ == "parquet" + else FusedScan( + base_scan, + base_scan.paths[offset : offset + plan.factor], + parquet_options, + ) ) for offset in range(paths_start, paths_end, plan.factor) if base_scan.paths[offset : offset + plan.factor] diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 47c8a9f6dfdb..bb506bb2a248 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -354,7 +354,7 @@ class ParquetOptions: improved performance on large datasets with complex filters. Default is False. use_hybrid_scan - Whether to use the two-pass ``HybridScanReader`` for ``SplitScan`` + Whether to use the two-pass ``HybridScanReader`` for split parquet tasks when a predicate can be pushed down to a parquet filter. Default is False. """ diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 0b804d2395b2..dcfbdbce4cbb 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -34,7 +34,6 @@ from cudf_polars.streaming.io import ( FusedScan, ParquetScanTask, - SplitScan, StreamingScan, expand_scan_for_rank, scan_partition_plan, @@ -468,7 +467,8 @@ def test_expand_scan_for_rank_fused_and_single_read( streaming_scan.tasks, expected_path_groups, strict=True ): assert isinstance(scan, ParquetScanTask) - assert isinstance(scan.base_task, FusedScan) + assert scan.split_index is None + assert scan.total_splits is None assert scan.paths == expected_paths @@ -499,9 +499,8 @@ def test_expand_scan_for_rank_split_files( streaming_scan.tasks, expected_splits, strict=True ): assert isinstance(scan, ParquetScanTask) - assert isinstance(scan.base_task, SplitScan) - assert scan.base_task.split_index == split_index - assert scan.base_task.total_splits == total_splits + assert scan.split_index == split_index + assert scan.total_splits == total_splits assert scan.paths == ["file.parquet"] @@ -554,7 +553,7 @@ def test_attach_cached_parquet_metadata_leaves_sub_row_group_split_unaligned( for scan in streaming_scan.tasks: assert isinstance(scan, ParquetScanTask) - assert isinstance(scan.base_task, SplitScan) + assert scan.is_split bounds = scan.get_task_bounds() assert bounds is not None assert bounds.row_groups is None @@ -590,7 +589,7 @@ def test_attach_cached_parquet_metadata_leaves_sliced_fused_scan_unaligned( for scan in streaming_scan.tasks: assert isinstance(scan, ParquetScanTask) - assert isinstance(scan.base_task, FusedScan) + assert not scan.is_split bounds = scan.get_task_bounds() assert bounds is not None assert bounds.row_groups is None @@ -751,12 +750,12 @@ def test_fused_scan_identity_equality() -> None: assert a != c -def test_split_scan_identity_equality() -> None: +def test_parquet_split_task_identity_equality() -> None: base = _make_parquet_scan(["a.parquet"]) - a = SplitScan(base, base.paths, 0, 4, base.parquet_options) - b = SplitScan(base, base.paths, 0, 4, base.parquet_options) - c = SplitScan(base, base.paths, 1, 4, base.parquet_options) + a = ParquetScanTask(base, base.paths, 0, 4, base.parquet_options) + b = ParquetScanTask(base, base.paths, 0, 4, base.parquet_options) + c = ParquetScanTask(base, base.paths, 1, 4, base.parquet_options) assert a == b assert hash(a) == hash(b) @@ -765,21 +764,21 @@ def test_split_scan_identity_equality() -> None: def test_streaming_scan_identity_equality() -> None: base = _make_parquet_scan(["a.parquet"]) - split = SplitScan( + split = ParquetScanTask( base, base.paths, 0, 2, base.parquet_options, ) - split_same = SplitScan( + split_same = ParquetScanTask( base, base.paths, 0, 2, base.parquet_options, ) - split_diff = SplitScan( + split_diff = ParquetScanTask( base, base.paths, 1, From 95d27b33e4ed070f57575eb661743133e89ed53c Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 8 Sep 2026 07:34:35 -0700 Subject: [PATCH 24/31] require split_index/num_splits in all cases to be consistent --- .../cudf_polars/cudf_polars/streaming/io.py | 68 +++++++++++-------- .../cudf_polars/tests/streaming/test_scan.py | 4 +- 2 files changed, 42 insertions(+), 30 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index b9de833d0178..8309bba7306b 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -461,10 +461,10 @@ class ParquetScanTask(IR): """Scan operation this task is based on.""" paths: list[str] """File paths assigned to this task.""" - split_index: int | None - """Index of the current split, or None for fused/full-file tasks.""" - total_splits: int | None - """Total number of splits for this file, or None for fused/full-file tasks.""" + split_index: int + """Index of the current split, or 0 for non-split tasks.""" + total_splits: int + """Total number of splits for a split file, or 1 for non-split tasks.""" parquet_options: ParquetOptions """Parquet-specific options.""" @@ -472,15 +472,19 @@ def __init__( self, base_scan: Scan, paths: list[str], - split_index: int | None, - total_splits: int | None, + split_index: int, + total_splits: int, parquet_options: ParquetOptions, ): if base_scan.typ != "parquet": # pragma: no cover raise ValueError(f"Expected a parquet scan, got: {base_scan.typ}") - if (split_index is None) != (total_splits is None): # pragma: no cover - raise ValueError("split_index and total_splits must be set together") - if split_index is not None and len(paths) > 1: # pragma: no cover + if total_splits < 1: # pragma: no cover + raise ValueError(f"Expected at least one split, got: {total_splits}") + if not 0 <= split_index < total_splits: # pragma: no cover + raise ValueError( + f"Expected split_index in [0, {total_splits}), got: {split_index}" + ) + if total_splits > 1 and len(paths) > 1: # pragma: no cover raise ValueError(f"Expected a single path for a split task, got: {paths}") self.base_scan = base_scan self.paths = paths @@ -499,8 +503,8 @@ def __init__( @property def is_split(self) -> bool: - """Whether this task is one split of a single parquet file.""" - return self.split_index is not None + """Whether this task is one of multiple splits of a single parquet file.""" + return self.total_splits > 1 def get_task_bounds(self) -> ParquetTaskBounds | None: """Return parquet read bounds for this task.""" @@ -530,14 +534,9 @@ def _fetch_parquet_info(self) -> list[CachedParquetInfo]: def _split_task_bounds( self, - cached_parquet_info: list[CachedParquetInfo] | None, - ) -> ParquetTaskBounds | None: + cached_parquet_info: list[CachedParquetInfo], + ) -> ParquetTaskBounds: """Return parquet read bounds for a split task.""" - if cached_parquet_info is None: - return None - - assert self.split_index is not None - assert self.total_splits is not None row_group_num_rows = cached_parquet_info[0].file_metadata.row_group_num_rows total_row_groups = len(row_group_num_rows) if self.total_splits <= total_row_groups: @@ -566,7 +565,11 @@ def _task_bounds_from_cached( cached_parquet_info: list[CachedParquetInfo] | None, ) -> ParquetTaskBounds | None: if self.is_split: - return self._split_task_bounds(cached_parquet_info) + return ( + None + if cached_parquet_info is None + else self._split_task_bounds(cached_parquet_info) + ) base_scan = self.base_scan row_groups: list[list[int]] | None = None @@ -602,8 +605,8 @@ def do_evaluate( cls, base_scan: Scan, paths: list[str], - split_index: int | None, - total_splits: int | None, + split_index: int, + total_splits: int, parquet_options: ParquetOptions, *, context: IRExecutionContext, @@ -613,7 +616,18 @@ def do_evaluate( base_scan = task.base_scan paths = task.paths cached_parquet_info = task._cached_parquet_info() - if cached_parquet_info is None and task.is_split: + if cached_parquet_info is None and ( + task.is_split + or ( + parquet_options.use_hybrid_scan + and len(paths) == 1 + and base_scan.skip_rows == 0 + and base_scan.n_rows == -1 + and base_scan.row_index is None + and base_scan.include_file_paths is None + and base_scan.predicate is not None + ) + ): cached_parquet_info = task._fetch_parquet_info() bounds = task._task_bounds_from_cached(cached_parquet_info) @@ -655,14 +669,12 @@ def do_evaluate( bounds.row_groups[0], stream, cached_parquet_info[0], - split_index=split_index or 0, - total_splits=total_splits or 1, + split_index=split_index, + total_splits=total_splits, stats_pruning=parquet_options._hybrid_scan_stats_pruning, ) if task.is_split: - assert split_index is not None - assert total_splits is not None nvtx_message = f"SplitScan: {paths[0]} [{split_index + 1}/{total_splits}]" else: nvtx_message = f"FusedScan: {', '.join(paths)}" @@ -833,8 +845,8 @@ def for_fused_files( ParquetScanTask( base_scan, base_scan.paths[offset : offset + plan.factor], - None, - None, + 0, + 1, parquet_options, ) if base_scan.typ == "parquet" diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index dcfbdbce4cbb..e98384094fee 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -467,8 +467,8 @@ def test_expand_scan_for_rank_fused_and_single_read( streaming_scan.tasks, expected_path_groups, strict=True ): assert isinstance(scan, ParquetScanTask) - assert scan.split_index is None - assert scan.total_splits is None + assert scan.split_index == 0 + assert scan.total_splits == 1 assert scan.paths == expected_paths From 9cf1ea1d7c815d9ff46ca051ba57390f6dab9568 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 8 Sep 2026 07:53:37 -0700 Subject: [PATCH 25/31] move to simple ParquetScanTask(ScanTask) --- .../cudf_polars/streaming/actor_graph/io.py | 11 +- .../cudf_polars/cudf_polars/streaming/io.py | 146 +++++++----------- .../cudf_polars/tests/streaming/test_scan.py | 50 ++++-- 3 files changed, 95 insertions(+), 112 deletions(-) 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 71d6f8ebc819..2c3f8c7bf577 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -39,7 +39,7 @@ send_metadata, ) from cudf_polars.streaming.io import ( - ParquetScanTask, + ScanTask, StreamingScan, StreamingSink, _prepare_sink_directory, @@ -61,7 +61,6 @@ IOPartitionPlan, PartitionInfo, ) - from cudf_polars.streaming.io import StreamingScanTask from cudf_polars.utils.config import MaxConcurrentIOTasks @@ -571,9 +570,7 @@ async def read_chunk( ) stop = time.monotonic_ns() ir_type = ( - task.trace_ir_type() - if isinstance(task, ParquetScanTask) - else type(task).__name__ + task.trace_ir_type() if isinstance(task, ScanTask) else type(task).__name__ ) log( "IO Task", @@ -619,7 +616,7 @@ async def scan_node( Estimated retained output size of each chunk in bytes. Used to estimate peak memory for admission before launching each read. """ - tasks: Sequence[StreamingScanTask] = ir.tasks + tasks: Sequence[ScanTask] = ir.tasks async with shutdown_on_error( context, ch_out, trace_ir=ir, ir_context=ir_context @@ -657,7 +654,7 @@ async def scan_node( lineariser = Lineariser(context, ch_out, num_producers) # Assign tasks to producers using round-robin - producer_tasks: list[list[tuple[int, StreamingScanTask]]] = [ + producer_tasks: list[list[tuple[int, ScanTask]]] = [ [] for _ in range(num_producers) ] for task_idx, task in enumerate(tasks): diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 8309bba7306b..78e1edaa72a0 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -11,7 +11,7 @@ import statistics from collections import defaultdict from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal, NamedTuple, Self, TypeAlias, overload +from typing import TYPE_CHECKING, Any, Literal, NamedTuple, Self, overload import polars as pl @@ -355,30 +355,33 @@ class ParquetTaskBounds(NamedTuple): n_rows: int -class FusedScan(IR): - """ - Input from one or more complete files read as a single task. - - Covers both FUSED_FILES (N > 1 small files grouped together) and - SINGLE_FILE (N = 1). - """ +class ScanTask(IR): + """Generic streaming scan task.""" __slots__ = ( "base_scan", "parquet_options", "paths", "schema", + "split_index", + "total_splits", ) _non_child = ( "base_scan", "paths", + "split_index", + "total_splits", "parquet_options", ) - _n_non_child_args = 3 + _n_non_child_args = 5 base_scan: Scan - """Scan operation this node is based on.""" + """Scan operation this task is based on.""" paths: list[str] """File paths assigned to this task.""" + split_index: int + """Index of the current split, or 0 for non-split tasks.""" + total_splits: int + """Total number of splits for a split file, or 1 for non-split tasks.""" parquet_options: ParquetOptions """Parquet-specific options.""" @@ -386,19 +389,36 @@ def __init__( self, base_scan: Scan, paths: list[str], + split_index: int, + total_splits: int, parquet_options: ParquetOptions, ): + if total_splits < 1: # pragma: no cover + raise ValueError(f"Expected at least one split, got: {total_splits}") + if not 0 <= split_index < total_splits: # pragma: no cover + raise ValueError( + f"Expected split_index in [0, {total_splits}), got: {split_index}" + ) self.schema = base_scan.schema self.base_scan = base_scan self.paths = paths + self.split_index = split_index + self.total_splits = total_splits self.parquet_options = parquet_options self._non_child_args = ( base_scan, paths, + split_index, + total_splits, parquet_options, ) self.children = () + @property + def is_split(self) -> bool: + """Whether this task is one of multiple splits of a single file.""" + return self.total_splits > 1 + def get_hashable(self) -> Hashable: """Hashable representation of the node.""" return ( @@ -406,20 +426,32 @@ def get_hashable(self) -> Hashable: tuple(self.schema.items()), self.base_scan.get_hashable(), tuple(self.paths), + self.split_index, + self.total_splits, self.parquet_options, ) + def trace_ir_type(self) -> str: + """Return the task type to use for IO-task tracing.""" + return "SplitScan" if self.is_split else type(self).__name__ + @classmethod def do_evaluate( cls, base_scan: Scan, paths: list[str], + split_index: int, + total_splits: int, parquet_options: ParquetOptions, *, context: IRExecutionContext, ) -> DataFrame: """Evaluate and return a dataframe.""" - with nvtx_annotate_cudf_polars(message=f"FusedScan: {', '.join(paths)}"): + if total_splits > 1: # pragma: no cover + raise NotImplementedError( + f"File splitting is not implemented for {base_scan.typ} scans." + ) + with nvtx_annotate_cudf_polars(message=f"ScanTask: {', '.join(paths)}"): return Scan.do_evaluate( base_scan.schema, base_scan.typ, @@ -437,36 +469,12 @@ def do_evaluate( ) -class ParquetScanTask(IR): +class ParquetScanTask(ScanTask): """Parquet-specific streaming scan task.""" - __slots__ = ( - "base_scan", - "parquet_options", - "paths", - "schema", - "split_index", - "total_splits", - ) - _non_child = ( - "base_scan", - "paths", - "split_index", - "total_splits", - "parquet_options", - ) - _n_non_child_args = 5 - - base_scan: Scan - """Scan operation this task is based on.""" - paths: list[str] - """File paths assigned to this task.""" - split_index: int - """Index of the current split, or 0 for non-split tasks.""" - total_splits: int - """Total number of splits for a split file, or 1 for non-split tasks.""" - parquet_options: ParquetOptions - """Parquet-specific options.""" + __slots__ = () + _non_child = ScanTask._non_child + _n_non_child_args = ScanTask._n_non_child_args def __init__( self, @@ -478,33 +486,9 @@ def __init__( ): if base_scan.typ != "parquet": # pragma: no cover raise ValueError(f"Expected a parquet scan, got: {base_scan.typ}") - if total_splits < 1: # pragma: no cover - raise ValueError(f"Expected at least one split, got: {total_splits}") - if not 0 <= split_index < total_splits: # pragma: no cover - raise ValueError( - f"Expected split_index in [0, {total_splits}), got: {split_index}" - ) if total_splits > 1 and len(paths) > 1: # pragma: no cover raise ValueError(f"Expected a single path for a split task, got: {paths}") - self.base_scan = base_scan - self.paths = paths - self.split_index = split_index - self.total_splits = total_splits - self.parquet_options = parquet_options - self.schema = base_scan.schema - self._non_child_args = ( - base_scan, - paths, - split_index, - total_splits, - parquet_options, - ) - self.children = () - - @property - def is_split(self) -> bool: - """Whether this task is one of multiple splits of a single parquet file.""" - return self.total_splits > 1 + super().__init__(base_scan, paths, split_index, total_splits, parquet_options) def get_task_bounds(self) -> ParquetTaskBounds | None: """Return parquet read bounds for this task.""" @@ -585,21 +569,6 @@ def _task_bounds_from_cached( ] return ParquetTaskBounds(row_groups, base_scan.skip_rows, base_scan.n_rows) - def get_hashable(self) -> Hashable: - """Hashable representation of the node.""" - return ( - type(self), - self.base_scan.get_hashable(), - tuple(self.paths), - self.split_index, - self.total_splits, - self.parquet_options, - ) - - def trace_ir_type(self) -> str: - """Return the task type to use for IO-task tracing.""" - return "SplitScan" if self.is_split else "FusedScan" - @classmethod def do_evaluate( cls, @@ -677,7 +646,7 @@ def do_evaluate( if task.is_split: nvtx_message = f"SplitScan: {paths[0]} [{split_index + 1}/{total_splits}]" else: - nvtx_message = f"FusedScan: {', '.join(paths)}" + nvtx_message = f"ParquetScanTask: {', '.join(paths)}" with nvtx_annotate_cudf_polars(message=nvtx_message): return Scan.do_evaluate( base_scan.schema, @@ -696,9 +665,6 @@ def do_evaluate( ) -StreamingScanTask: TypeAlias = FusedScan | ParquetScanTask - - @lower_ir_node.register(Empty) def _( ir: Empty, rec: LowerIRTransformer @@ -777,11 +743,11 @@ class StreamingScan(IR): ) _n_non_child_args = 2 base_scan: Scan - tasks: Sequence[StreamingScanTask] + tasks: Sequence[ScanTask] def __init__( self, - tasks: Sequence[StreamingScanTask], + tasks: Sequence[ScanTask], base_scan: Scan, ): self.base_scan = base_scan @@ -807,7 +773,7 @@ def for_split_files( path_end = math.ceil((local_offset + local_count) / plan.factor) local_paths = base_scan.paths[path_offset:path_end] sindex = local_offset % plan.factor - tasks: list[StreamingScanTask] = [] + tasks: list[ScanTask] = [] splits_created = 0 for path in local_paths: while sindex < plan.factor and splits_created < local_count: @@ -840,7 +806,7 @@ def for_fused_files( local_offset, local_count = _rank_slice(partition_count, rank, nranks) paths_start = local_offset * plan.factor paths_end = paths_start + plan.factor * local_count - tasks: list[StreamingScanTask] = [ + tasks: list[ScanTask] = [ ( ParquetScanTask( base_scan, @@ -850,9 +816,11 @@ def for_fused_files( parquet_options, ) if base_scan.typ == "parquet" - else FusedScan( + else ScanTask( base_scan, base_scan.paths[offset : offset + plan.factor], + 0, + 1, parquet_options, ) ) @@ -869,7 +837,7 @@ def get_hashable(self) -> Hashable: @classmethod def do_evaluate( cls, - tasks: Sequence[StreamingScanTask], + tasks: Sequence[ScanTask], base_scan: Scan, *, context: IRExecutionContext, diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index e98384094fee..35588d9fd978 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -32,8 +32,8 @@ StatsCollector, ) from cudf_polars.streaming.io import ( - FusedScan, ParquetScanTask, + ScanTask, StreamingScan, expand_scan_for_rank, scan_partition_plan, @@ -179,8 +179,8 @@ def recording_prefetch( ) scan = _make_parquet_scan(paths) - fused = FusedScan(scan, paths, scan.parquet_options) - streaming_scan = StreamingScan([fused], scan) + task = ParquetScanTask(scan, paths, 0, 1, scan.parquet_options) + streaming_scan = StreamingScan([task], scan) result = prefetch_parquet_file_metadata_for_ir( streaming_scan, py_executor=None, stats=stats @@ -195,8 +195,8 @@ def test_prefetch_parquet_file_metadata_remote_only(tmp_path, df) -> None: local_path = str(next(tmp_path.glob("*.parquet"))) scan = _make_parquet_scan([local_path]) - fused = FusedScan(scan, scan.paths, scan.parquet_options) - streaming_scan = StreamingScan([fused], scan) + task = ParquetScanTask(scan, scan.paths, 0, 1, scan.parquet_options) + streaming_scan = StreamingScan([task], scan) # Local paths are skipped entirely when remote_only=True. result = prefetch_parquet_file_metadata_for_ir( @@ -420,6 +420,24 @@ def _make_parquet_scan( ) +def _make_csv_scan(paths: list[str]) -> Scan: + return Scan( + {"x": DataType(pl.Int64())}, + "csv", + {}, + None, + paths, + None, + 0, + -1, + None, + None, + None, + ParquetOptions(), + None, + ) + + @pytest.mark.parametrize( "plan,paths,rank,nranks,expected_path_groups", [ @@ -597,11 +615,11 @@ def test_attach_cached_parquet_metadata_leaves_sliced_fused_scan_unaligned( def test_streaming_scan_raises() -> None: # This isn't reachable by normal cudf-polars usage. - scan = _make_parquet_scan(["file.parquet"]) - fused = FusedScan(scan, scan.paths, scan.parquet_options) + scan = _make_csv_scan(["file.csv"]) + task = ScanTask(scan, scan.paths, 0, 1, scan.parquet_options) ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): - StreamingScan.do_evaluate([fused], scan, context=ctx) + StreamingScan.do_evaluate([task], scan, context=ctx) @pytest.mark.parametrize( @@ -675,11 +693,11 @@ def test_streaming_scan_missing_prefetch_metadata_raises() -> None: scan = _make_parquet_scan( ["file.parquet"], parquet_options=ParquetOptions(prefetch_file_metadata=True) ) - fused = FusedScan(scan, scan.paths, scan.parquet_options) + task = ParquetScanTask(scan, scan.paths, 0, 1, scan.parquet_options) ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): - StreamingScan.do_evaluate([fused], scan, context=ctx) + StreamingScan.do_evaluate([task], scan, context=ctx) def test_prefetch_file_metadata_join( @@ -737,13 +755,13 @@ def test_prefetch_file_metadata_with_cached_scan_parent_nodes( assert_gpu_result_equal(q, engine=engine) -def test_fused_scan_identity_equality() -> None: - base = _make_parquet_scan(["a.parquet", "b.parquet"]) - paths = ["a.parquet"] +def test_scan_task_identity_equality() -> None: + base = _make_csv_scan(["a.csv", "b.csv"]) + paths = ["a.csv"] - a = FusedScan(base, paths, base.parquet_options) - b = FusedScan(base, paths, base.parquet_options) - c = FusedScan(base, ["b.parquet"], base.parquet_options) + a = ScanTask(base, paths, 0, 1, base.parquet_options) + b = ScanTask(base, paths, 0, 1, base.parquet_options) + c = ScanTask(base, ["b.csv"], 0, 1, base.parquet_options) assert a == b assert hash(a) == hash(b) From c7161b8610864daeb58df7d282d8796067a9f83f Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 8 Sep 2026 08:51:46 -0700 Subject: [PATCH 26/31] fix CI --- .../cudf_polars/cudf_polars/dsl/utils/io.py | 15 ++++--- .../cudf_polars/tests/streaming/test_scan.py | 44 ++++++++++++++++++- 2 files changed, 53 insertions(+), 6 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index dd3f7fe480a8..010e65a3c463 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -243,7 +243,7 @@ def attach_cached_parquet_metadata( cached_parquet_info_map: dict[str, CachedParquetInfo], ) -> None: """ - Attach prefetched metadata to parquet scan nodes. + Attach prefetched metadata to parquet scan tasks. This is an optimization only and does not affect IR identity. @@ -257,9 +257,14 @@ def attach_cached_parquet_metadata( for node in traversal([root]): if isinstance(node, StreamingScan) and node.base_scan.typ == "parquet": base_scan = node.base_scan - if not all(path in cached_parquet_info_map for path in base_scan.paths): + task_paths = {path for task in node.tasks for path in task.paths} + cached_paths = [ + path + for path in base_scan.paths + if path in task_paths and path in cached_parquet_info_map + ] + cached = [cached_parquet_info_map[path] for path in cached_paths] + if not cached: continue - cached = [cached_parquet_info_map[path] for path in base_scan.paths] - Scan._validate_cached_parquet_info(base_scan.paths, cached) + Scan._validate_cached_parquet_info(cached_paths, cached) base_scan.cached_parquet_info = cached - base_scan._non_child_args = (*base_scan._non_child_args[:-1], cached) diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 35588d9fd978..1867e6cfd766 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -424,7 +424,20 @@ def _make_csv_scan(paths: list[str]) -> Scan: return Scan( {"x": DataType(pl.Int64())}, "csv", - {}, + { + "has_header": True, + "schema": None, + "skip_rows": 0, + "skip_rows_after_header": 0, + "parse_options": { + "separator": ord(","), + "quote_char": ord('"'), + "eol_char": ord("\n"), + "null_values": None, + "comment_prefix": None, + "decimal_comma": False, + }, + }, None, paths, None, @@ -550,6 +563,35 @@ def test_attach_cached_parquet_metadata_resolves_row_groups( assert row_groups == [[[0]], [[1]]] +def test_attach_cached_parquet_metadata_uses_rank_local_tasks( + tmp_path: Path, +) -> None: + paths = [str(tmp_path / f"part-{i}.parquet") for i in range(4)] + for path in paths: + pl.DataFrame({"x": range(4)}).write_parquet(path, row_group_size=2) + + base = _make_parquet_scan(paths) + streaming_scan = StreamingScan.for_fused_files( + base, + IOPartitionPlan(1, IOPartitionFlavor.SINGLE_FILE), + partition_count=4, + rank=0, + nranks=2, + parquet_options=base.parquet_options, + ) + + cached = prefetch_parquet_file_metadata_for_ir(streaming_scan, None) + attach_cached_parquet_metadata(streaming_scan, cached) + + assert base.cached_parquet_info is not None + assert [info.path for info in base.cached_parquet_info] == paths[:2] + for scan in streaming_scan.tasks: + assert isinstance(scan, ParquetScanTask) + bounds = scan.get_task_bounds() + assert bounds is not None + assert bounds.row_groups == [[0, 1]] + + def test_attach_cached_parquet_metadata_leaves_sub_row_group_split_unaligned( tmp_path: Path, ) -> None: From 20f28a16242b0df9268ba13a185ce576d46f4858 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 8 Sep 2026 09:57:27 -0700 Subject: [PATCH 27/31] fix code-cov --- python/cudf_polars/tests/streaming/test_scan.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 1867e6cfd766..4fc87b266e08 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -203,6 +203,8 @@ def test_prefetch_parquet_file_metadata_remote_only(tmp_path, df) -> None: streaming_scan, py_executor=None, stats=None, remote_only=True ) assert result == {} + attach_cached_parquet_metadata(streaming_scan, result) + assert scan.cached_parquet_info is None # The same local path is prefetched when remote_only=False (the default). result = prefetch_parquet_file_metadata_for_ir( From 848337f4854123be6fba4694afd496ae1b1fb3d2 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 8 Sep 2026 11:30:55 -0700 Subject: [PATCH 28/31] address comments --- .../cudf_polars/cudf_polars/streaming/io.py | 45 +++++++++++++------ .../cudf_polars/tests/streaming/test_scan.py | 22 ++++++--- 2 files changed, 46 insertions(+), 21 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 78e1edaa72a0..ece8e6b755f0 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -211,10 +211,25 @@ def hybrid_scan_eligible( include_file_paths: str | None, predicate: NamedExpr | None, ) -> bool: - """Whether a parquet split is eligible for the HybridScanReader path.""" + """Whether cached parquet metadata can use the HybridScanReader path.""" + return cached_parquet_info is not None and _hybrid_scan_preconditions( + parquet_options, + row_index=row_index, + include_file_paths=include_file_paths, + predicate=predicate, + ) + + +def _hybrid_scan_preconditions( + parquet_options: ParquetOptions, + *, + row_index: tuple[str, int] | None, + include_file_paths: str | None, + predicate: NamedExpr | None, +) -> bool: + """Whether scan options allow hybrid scan if metadata is available.""" return ( 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 @@ -343,9 +358,9 @@ def _read_with_hybrid_scan( return DataFrame(columns, stream=stream).select(list(schema.keys())) -class ParquetTaskBounds(NamedTuple): +class ParquetScanTaskBounds(NamedTuple): """ - Read bounds for a parquet task. + Read bounds for a parquet scan task. ``row_groups=None`` means the task is not row-group aligned. """ @@ -490,7 +505,7 @@ def __init__( raise ValueError(f"Expected a single path for a split task, got: {paths}") super().__init__(base_scan, paths, split_index, total_splits, parquet_options) - def get_task_bounds(self) -> ParquetTaskBounds | None: + def get_task_bounds(self) -> ParquetScanTaskBounds | None: """Return parquet read bounds for this task.""" return self._task_bounds_from_cached(self._cached_parquet_info()) @@ -519,7 +534,7 @@ def _fetch_parquet_info(self) -> list[CachedParquetInfo]: def _split_task_bounds( self, cached_parquet_info: list[CachedParquetInfo], - ) -> ParquetTaskBounds: + ) -> ParquetScanTaskBounds: """Return parquet read bounds for a split task.""" row_group_num_rows = cached_parquet_info[0].file_metadata.row_group_num_rows total_row_groups = len(row_group_num_rows) @@ -542,12 +557,12 @@ def _split_task_bounds( if self.split_index == self.total_splits - 1: n_rows = -1 - return ParquetTaskBounds(row_groups, skip_rows, n_rows) + return ParquetScanTaskBounds(row_groups, skip_rows, n_rows) def _task_bounds_from_cached( self, cached_parquet_info: list[CachedParquetInfo] | None, - ) -> ParquetTaskBounds | None: + ) -> ParquetScanTaskBounds | None: if self.is_split: return ( None @@ -567,7 +582,7 @@ def _task_bounds_from_cached( list(range(len(info.file_metadata.row_group_num_rows))) for info in cached_parquet_info ] - return ParquetTaskBounds(row_groups, base_scan.skip_rows, base_scan.n_rows) + return ParquetScanTaskBounds(row_groups, base_scan.skip_rows, base_scan.n_rows) @classmethod def do_evaluate( @@ -588,13 +603,15 @@ def do_evaluate( if cached_parquet_info is None and ( task.is_split or ( - parquet_options.use_hybrid_scan - and len(paths) == 1 + len(paths) == 1 and base_scan.skip_rows == 0 and base_scan.n_rows == -1 - and base_scan.row_index is None - and base_scan.include_file_paths is None - and base_scan.predicate is not None + and _hybrid_scan_preconditions( + parquet_options, + row_index=base_scan.row_index, + include_file_paths=base_scan.include_file_paths, + predicate=base_scan.predicate, + ) ) ): cached_parquet_info = task._fetch_parquet_info() diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 4fc87b266e08..1d3148edb96c 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -732,16 +732,24 @@ def test_scan_path_mismatch_raises() -> None: ) -def test_streaming_scan_missing_prefetch_metadata_raises() -> None: - # This isn't reachable by polars' public API, so we test it directly. +def test_parquet_split_task_fetches_missing_metadata(tmp_path: Path) -> None: + source = tmp_path / "data.parquet" + pl.DataFrame({"x": range(4)}).write_parquet(source, row_group_size=2) + scan = _make_parquet_scan( - ["file.parquet"], parquet_options=ParquetOptions(prefetch_file_metadata=True) + [str(source)], parquet_options=ParquetOptions(prefetch_file_metadata=True) ) - task = ParquetScanTask(scan, scan.paths, 0, 1, scan.parquet_options) - ctx = IRExecutionContext() - with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): - StreamingScan.do_evaluate([task], scan, context=ctx) + result = ParquetScanTask.do_evaluate( + scan, + scan.paths, + 0, + 2, + scan.parquet_options, + context=IRExecutionContext(), + ) + + assert result.to_polars().to_dict(as_series=False) == {"x": [0, 1]} def test_prefetch_file_metadata_join( From aa6fdbe958232d506548bf82c812744aa3051b00 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 8 Sep 2026 11:42:16 -0700 Subject: [PATCH 29/31] small revision --- .../cudf_polars/cudf_polars/streaming/io.py | 34 ++++++++++--------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index ece8e6b755f0..c34ef5745e18 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -522,13 +522,15 @@ def _cached_parquet_info(self) -> list[CachedParquetInfo] | None: return None return [cached_by_path[path] for path in self.paths] - def _fetch_parquet_info(self) -> list[CachedParquetInfo]: + def _fetch_parquet_info( + self, *, parse_hybrid_metadata: bool = False + ) -> list[CachedParquetInfo]: """Fetch parquet metadata for this task's paths.""" from cudf_polars.dsl.utils.io import _prefetch_parquet_footers_for_paths return _prefetch_parquet_footers_for_paths( self.paths, - parse_hybrid_metadata=self.parquet_options.use_hybrid_scan, + parse_hybrid_metadata=parse_hybrid_metadata, ) def _split_task_bounds( @@ -600,21 +602,21 @@ def do_evaluate( base_scan = task.base_scan paths = task.paths cached_parquet_info = task._cached_parquet_info() - if cached_parquet_info is None and ( - task.is_split - or ( - len(paths) == 1 - and base_scan.skip_rows == 0 - and base_scan.n_rows == -1 - and _hybrid_scan_preconditions( - parquet_options, - row_index=base_scan.row_index, - include_file_paths=base_scan.include_file_paths, - predicate=base_scan.predicate, - ) + should_try_hybrid_scan = ( + len(paths) == 1 + and base_scan.skip_rows == 0 + and base_scan.n_rows == -1 + and _hybrid_scan_preconditions( + parquet_options, + row_index=base_scan.row_index, + include_file_paths=base_scan.include_file_paths, + predicate=base_scan.predicate, + ) + ) + if cached_parquet_info is None and (task.is_split or should_try_hybrid_scan): + cached_parquet_info = task._fetch_parquet_info( + parse_hybrid_metadata=should_try_hybrid_scan ) - ): - cached_parquet_info = task._fetch_parquet_info() bounds = task._task_bounds_from_cached(cached_parquet_info) assert bounds is not None From d6d264aebe8c38edf46b469e102d9950b83ee251 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 8 Sep 2026 12:05:46 -0700 Subject: [PATCH 30/31] address more comments --- .../cudf_polars/cudf_polars/streaming/io.py | 156 +++++------------- .../cudf_polars/tests/streaming/test_scan.py | 8 +- 2 files changed, 44 insertions(+), 120 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index c34ef5745e18..39ea25abc7d4 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -11,7 +11,7 @@ import statistics from collections import defaultdict from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal, NamedTuple, Self, overload +from typing import TYPE_CHECKING, Any, Literal, NamedTuple, overload import polars as pl @@ -183,44 +183,47 @@ def expand_scan_for_rank( StreamingScan Rank-local streaming scan. """ + local_offset, local_count = _rank_slice(partition_count, rank, nranks) if plan.flavor == IOPartitionFlavor.SPLIT_FILES: - return StreamingScan.for_split_files( - ir, - plan, - partition_count, - rank=rank, - nranks=nranks, - parquet_options=parquet_options, - ) + path_offset = local_offset // plan.factor + path_end = math.ceil((local_offset + local_count) / plan.factor) + local_paths = ir.paths[path_offset:path_end] + sindex = local_offset % plan.factor + tasks: list[ScanTask] = [] + splits_created = 0 + for path in local_paths: + while sindex < plan.factor and splits_created < local_count: + tasks.append( + ParquetScanTask( + ir, + [path], + sindex, + plan.factor, + parquet_options, + ) + ) + sindex += 1 + splits_created += 1 + sindex = 0 else: - return StreamingScan.for_fused_files( - ir, - plan, - partition_count, - rank=rank, - nranks=nranks, - parquet_options=parquet_options, - ) + paths_start = local_offset * plan.factor + paths_end = paths_start + plan.factor * local_count + task_type = ParquetScanTask if ir.typ == "parquet" else ScanTask + tasks = [ + task_type( + ir, + ir.paths[offset : offset + plan.factor], + 0, + 1, + parquet_options, + ) + for offset in range(paths_start, paths_end, plan.factor) + if ir.paths[offset : offset + plan.factor] + ] + return StreamingScan(tasks, ir) def hybrid_scan_eligible( - parquet_options: ParquetOptions, - *, - cached_parquet_info: list[CachedParquetInfo] | None, - row_index: tuple[str, int] | None, - include_file_paths: str | None, - predicate: NamedExpr | None, -) -> bool: - """Whether cached parquet metadata can use the HybridScanReader path.""" - return cached_parquet_info is not None and _hybrid_scan_preconditions( - parquet_options, - row_index=row_index, - include_file_paths=include_file_paths, - predicate=predicate, - ) - - -def _hybrid_scan_preconditions( parquet_options: ParquetOptions, *, row_index: tuple[str, int] | None, @@ -606,7 +609,7 @@ def do_evaluate( len(paths) == 1 and base_scan.skip_rows == 0 and base_scan.n_rows == -1 - and _hybrid_scan_preconditions( + and hybrid_scan_eligible( parquet_options, row_index=base_scan.row_index, include_file_paths=base_scan.include_file_paths, @@ -625,16 +628,10 @@ def do_evaluate( # TODO: Investigate re-enabling for some of the excluded paths # (row_index / include_file_paths). Needs performance investigation. if ( - len(paths) == 1 + should_try_hybrid_scan and bounds.row_groups is not None and len(bounds.row_groups) == 1 - and hybrid_scan_eligible( - parquet_options, - cached_parquet_info=cached_parquet_info, - row_index=base_scan.row_index, - include_file_paths=base_scan.include_file_paths, - predicate=base_scan.predicate, - ) + and cached_parquet_info is not None ): assert base_scan.predicate is not None assert cached_parquet_info is not None @@ -775,79 +772,6 @@ def __init__( self._non_child_args = (tasks, base_scan) self.children = () - @classmethod - def for_split_files( - cls, - base_scan: Scan, - plan: IOPartitionPlan, - partition_count: int, - *, - rank: int, - nranks: int, - parquet_options: ParquetOptions, - ) -> Self: - """Construct a StreamingScan where each file is split into factor partitions.""" - local_offset, local_count = _rank_slice(partition_count, rank, nranks) - path_offset = local_offset // plan.factor - path_end = math.ceil((local_offset + local_count) / plan.factor) - local_paths = base_scan.paths[path_offset:path_end] - sindex = local_offset % plan.factor - tasks: list[ScanTask] = [] - splits_created = 0 - for path in local_paths: - while sindex < plan.factor and splits_created < local_count: - tasks.append( - ParquetScanTask( - base_scan, - [path], - sindex, - plan.factor, - parquet_options, - ) - ) - sindex += 1 - splits_created += 1 - sindex = 0 - return cls(tasks, base_scan) - - @classmethod - def for_fused_files( - cls, - base_scan: Scan, - plan: IOPartitionPlan, - partition_count: int, - *, - rank: int, - nranks: int, - parquet_options: ParquetOptions, - ) -> Self: - """Construct a StreamingScan where factor files are grouped into one partition.""" - local_offset, local_count = _rank_slice(partition_count, rank, nranks) - paths_start = local_offset * plan.factor - paths_end = paths_start + plan.factor * local_count - tasks: list[ScanTask] = [ - ( - ParquetScanTask( - base_scan, - base_scan.paths[offset : offset + plan.factor], - 0, - 1, - parquet_options, - ) - if base_scan.typ == "parquet" - else ScanTask( - base_scan, - base_scan.paths[offset : offset + plan.factor], - 0, - 1, - parquet_options, - ) - ) - for offset in range(paths_start, paths_end, plan.factor) - if base_scan.paths[offset : offset + plan.factor] - ] - return cls(tasks, base_scan) - def get_hashable(self) -> Hashable: """Hashable representation of the node.""" # We don't need to include base_scan / schema, since it's in all the scan tasks. diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 1d3148edb96c..6b0818356da6 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -544,7 +544,7 @@ def test_attach_cached_parquet_metadata_resolves_row_groups( pl.DataFrame({"x": range(4)}).write_parquet(source, row_group_size=2) base = _make_parquet_scan([str(source)]) - streaming_scan = StreamingScan.for_split_files( + streaming_scan = expand_scan_for_rank( base, IOPartitionPlan(2, IOPartitionFlavor.SPLIT_FILES), partition_count=2, @@ -573,7 +573,7 @@ def test_attach_cached_parquet_metadata_uses_rank_local_tasks( pl.DataFrame({"x": range(4)}).write_parquet(path, row_group_size=2) base = _make_parquet_scan(paths) - streaming_scan = StreamingScan.for_fused_files( + streaming_scan = expand_scan_for_rank( base, IOPartitionPlan(1, IOPartitionFlavor.SINGLE_FILE), partition_count=4, @@ -601,7 +601,7 @@ def test_attach_cached_parquet_metadata_leaves_sub_row_group_split_unaligned( pl.DataFrame({"x": range(4)}).write_parquet(source, row_group_size=2) base = _make_parquet_scan([str(source)]) - streaming_scan = StreamingScan.for_split_files( + streaming_scan = expand_scan_for_rank( base, IOPartitionPlan(4, IOPartitionFlavor.SPLIT_FILES), partition_count=4, @@ -637,7 +637,7 @@ def test_attach_cached_parquet_metadata_leaves_sliced_fused_scan_unaligned( base = _make_parquet_scan( [str(source)], skip_rows=skip_rows, n_rows=n_rows, row_index=row_index ) - streaming_scan = StreamingScan.for_fused_files( + streaming_scan = expand_scan_for_rank( base, IOPartitionPlan(1, IOPartitionFlavor.SINGLE_READ), partition_count=1, From f13969469ee84e16ed4efe1e3cfc2bf10ef4468f Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 8 Sep 2026 12:44:34 -0700 Subject: [PATCH 31/31] remove parquet_options from ScanTask --- .../cudf_polars/cudf_polars/streaming/io.py | 76 ++++++++++--------- .../cudf_polars/tests/streaming/test_scan.py | 34 ++++++++- 2 files changed, 72 insertions(+), 38 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 39ea25abc7d4..93fb69802ce2 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -11,7 +11,7 @@ import statistics from collections import defaultdict from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal, NamedTuple, overload +from typing import TYPE_CHECKING, Any, ClassVar, Literal, NamedTuple, overload import polars as pl @@ -208,18 +208,15 @@ def expand_scan_for_rank( else: paths_start = local_offset * plan.factor paths_end = paths_start + plan.factor * local_count - task_type = ParquetScanTask if ir.typ == "parquet" else ScanTask - tasks = [ - task_type( - ir, - ir.paths[offset : offset + plan.factor], - 0, - 1, - parquet_options, - ) - for offset in range(paths_start, paths_end, plan.factor) - if ir.paths[offset : offset + plan.factor] - ] + tasks = [] + for offset in range(paths_start, paths_end, plan.factor): + paths = ir.paths[offset : offset + plan.factor] + if not paths: + continue + if ir.typ == "parquet": + tasks.append(ParquetScanTask(ir, paths, 0, 1, parquet_options)) + else: + tasks.append(ScanTask(ir, paths, 0, 1)) return StreamingScan(tasks, ir) @@ -378,20 +375,18 @@ class ScanTask(IR): __slots__ = ( "base_scan", - "parquet_options", "paths", "schema", "split_index", "total_splits", ) - _non_child = ( + _non_child: ClassVar[tuple[str, ...]] = ( "base_scan", "paths", "split_index", "total_splits", - "parquet_options", ) - _n_non_child_args = 5 + _n_non_child_args = 4 base_scan: Scan """Scan operation this task is based on.""" paths: list[str] @@ -400,8 +395,6 @@ class ScanTask(IR): """Index of the current split, or 0 for non-split tasks.""" total_splits: int """Total number of splits for a split file, or 1 for non-split tasks.""" - parquet_options: ParquetOptions - """Parquet-specific options.""" def __init__( self, @@ -409,11 +402,10 @@ def __init__( paths: list[str], split_index: int, total_splits: int, - parquet_options: ParquetOptions, ): - if total_splits < 1: # pragma: no cover + if total_splits < 1: raise ValueError(f"Expected at least one split, got: {total_splits}") - if not 0 <= split_index < total_splits: # pragma: no cover + if not 0 <= split_index < total_splits: raise ValueError( f"Expected split_index in [0, {total_splits}), got: {split_index}" ) @@ -422,13 +414,11 @@ def __init__( self.paths = paths self.split_index = split_index self.total_splits = total_splits - self.parquet_options = parquet_options self._non_child_args = ( base_scan, paths, split_index, total_splits, - parquet_options, ) self.children = () @@ -446,7 +436,6 @@ def get_hashable(self) -> Hashable: tuple(self.paths), self.split_index, self.total_splits, - self.parquet_options, ) def trace_ir_type(self) -> str: @@ -460,7 +449,6 @@ def do_evaluate( paths: list[str], split_index: int, total_splits: int, - parquet_options: ParquetOptions, *, context: IRExecutionContext, ) -> DataFrame: @@ -481,7 +469,7 @@ def do_evaluate( base_scan.row_index, base_scan.include_file_paths, base_scan.predicate, - parquet_options, + base_scan.parquet_options, None, context=context, ) @@ -490,9 +478,12 @@ def do_evaluate( class ParquetScanTask(ScanTask): """Parquet-specific streaming scan task.""" - __slots__ = () - _non_child = ScanTask._non_child - _n_non_child_args = ScanTask._n_non_child_args + __slots__ = ("parquet_options",) + _non_child: ClassVar[tuple[str, ...]] = ( + *ScanTask._non_child, + "parquet_options", + ) + _n_non_child_args = 5 def __init__( self, @@ -502,11 +493,28 @@ def __init__( total_splits: int, parquet_options: ParquetOptions, ): - if base_scan.typ != "parquet": # pragma: no cover + if base_scan.typ != "parquet": raise ValueError(f"Expected a parquet scan, got: {base_scan.typ}") - if total_splits > 1 and len(paths) > 1: # pragma: no cover + if total_splits > 1 and len(paths) > 1: raise ValueError(f"Expected a single path for a split task, got: {paths}") - super().__init__(base_scan, paths, split_index, total_splits, parquet_options) + super().__init__(base_scan, paths, split_index, total_splits) + self.parquet_options = parquet_options + self._non_child_args = ( + *self._non_child_args, + parquet_options, + ) + + def get_hashable(self) -> Hashable: + """Hashable representation of the node.""" + return ( + type(self), + tuple(self.schema.items()), + self.base_scan.get_hashable(), + tuple(self.paths), + self.split_index, + self.total_splits, + self.parquet_options, + ) def get_task_bounds(self) -> ParquetScanTaskBounds | None: """Return parquet read bounds for this task.""" @@ -590,7 +598,7 @@ def _task_bounds_from_cached( return ParquetScanTaskBounds(row_groups, base_scan.skip_rows, base_scan.n_rows) @classmethod - def do_evaluate( + def do_evaluate( # type: ignore[override] cls, base_scan: Scan, paths: list[str], diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 6b0818356da6..8510d8a3bba9 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -660,7 +660,7 @@ def test_attach_cached_parquet_metadata_leaves_sliced_fused_scan_unaligned( def test_streaming_scan_raises() -> None: # This isn't reachable by normal cudf-polars usage. scan = _make_csv_scan(["file.csv"]) - task = ScanTask(scan, scan.paths, 0, 1, scan.parquet_options) + task = ScanTask(scan, scan.paths, 0, 1) ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): StreamingScan.do_evaluate([task], scan, context=ctx) @@ -811,15 +811,41 @@ def test_scan_task_identity_equality() -> None: base = _make_csv_scan(["a.csv", "b.csv"]) paths = ["a.csv"] - a = ScanTask(base, paths, 0, 1, base.parquet_options) - b = ScanTask(base, paths, 0, 1, base.parquet_options) - c = ScanTask(base, ["b.csv"], 0, 1, base.parquet_options) + a = ScanTask(base, paths, 0, 1) + b = ScanTask(base, paths, 0, 1) + c = ScanTask(base, ["b.csv"], 0, 1) assert a == b assert hash(a) == hash(b) assert a != c +def test_scan_task_validates_split_bounds() -> None: + base = _make_csv_scan(["a.csv"]) + + with pytest.raises(ValueError, match=r"Expected at least one split"): + ScanTask(base, ["a.csv"], 0, 0) + + with pytest.raises(ValueError, match=r"Expected split_index in"): + ScanTask(base, ["a.csv"], 1, 1) + + +def test_parquet_scan_task_validates_inputs() -> None: + csv_scan = _make_csv_scan(["a.csv"]) + with pytest.raises(ValueError, match=r"Expected a parquet scan"): + ParquetScanTask(csv_scan, csv_scan.paths, 0, 1, csv_scan.parquet_options) + + parquet_scan = _make_parquet_scan(["a.parquet", "b.parquet"]) + with pytest.raises(ValueError, match=r"Expected a single path for a split task"): + ParquetScanTask( + parquet_scan, + parquet_scan.paths, + 0, + 2, + parquet_scan.parquet_options, + ) + + def test_parquet_split_task_identity_equality() -> None: base = _make_parquet_scan(["a.parquet"])