From 9db0749d0e86a1145f21cc965e84fd5bcf96ac03 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Wed, 10 Jun 2026 16:12:05 +0000 Subject: [PATCH 1/4] Add FusedScan and nxtx annotations for FusedScan and SplitScan --- .../cudf_polars/streaming/actor_graph/io.py | 4 +- .../cudf_polars/cudf_polars/streaming/io.py | 161 ++++++++++++++---- .../cudf_polars/tests/streaming/test_scan.py | 12 +- 3 files changed, 136 insertions(+), 41 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 822a84d38c75..547f1e6be53c 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -63,7 +63,7 @@ PartitionInfo, StatsCollector, ) - from cudf_polars.streaming.io import SplitScan + from cudf_polars.streaming.io import FusedScan, SplitScan class Lineariser: @@ -408,7 +408,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, Scan | SplitScan]]] = [ + producer_tasks: list[list[tuple[int, SplitScan | FusedScan]]] = [ [] for _ in range(num_producers) ] for task_idx, scan in enumerate(scans): diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index d56c31ce379b..74006e601620 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -112,9 +112,9 @@ def expand_scan_for_rank( rank: int, nranks: int, parquet_options: ParquetOptions, -) -> list[Scan | SplitScan]: +) -> list[SplitScan | FusedScan]: """ - Expand a Scan node into rank-local Scan and SplitScan operations. + Expand a Scan node into rank-local SplitScan and FusedScan operations. Parameters ---------- @@ -131,10 +131,10 @@ def expand_scan_for_rank( Returns ------- - list[Scan | SplitScan] + list[SplitScan | FusedScan] Rank-local scan operations. """ - scans: list[Scan | SplitScan] = [] + scans: list[SplitScan | FusedScan] = [] if plan.flavor == IOPartitionFlavor.SPLIT_FILES: count = plan.factor * len(ir.paths) local_count = math.ceil(count / nranks) @@ -182,23 +182,8 @@ def expand_scan_for_rank( paths_offset_end = paths_offset_start + plan.factor * local_count for offset in range(paths_offset_start, paths_offset_end, plan.factor): local_paths = ir.paths[offset : offset + plan.factor] - if len(local_paths) > 0: # Only add scan if there are paths - scans.append( - Scan( - ir.schema, - ir.typ, - ir.reader_options, - ir.cloud_options, - local_paths, - ir.with_columns, - ir.skip_rows, - ir.n_rows, - ir.row_index, - ir.include_file_paths, - ir.predicate, - parquet_options, - ) - ) + if len(local_paths) > 0: + scans.append(FusedScan(ir.schema, ir, local_paths, parquet_options)) return scans @@ -207,7 +192,7 @@ class SplitScan(IR): """ Input from a split file. - This class wraps a single-file `Scan` object. At + 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. @@ -337,20 +322,124 @@ def do_evaluate( n_rows = -1 # Perform the partial read - return Scan.do_evaluate( - schema, - typ, - reader_options, + with nvtx_annotate_cudf_polars( + message=f"SplitScan: {Path(paths[0]).name} [{split_index + 1}/{total_splits}]" + ): + return Scan.do_evaluate( + schema, + typ, + reader_options, + paths, + with_columns, + skip_rows, + n_rows, + row_index, + include_file_paths, + predicate, + parquet_options, + context=context, + ) + + +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). The ``paths`` attribute holds the file group + assigned to this task; the remaining read options come from + ``base_scan``. + """ + + __slots__ = ( + "base_scan", + "parquet_options", + "paths", + "schema", + ) + _non_child = ( + "schema", + "base_scan", + "paths", + "parquet_options", + ) + _n_non_child_args = 11 + 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.""" + + def __init__( + self, + schema: Schema, + base_scan: Scan, + paths: list[str], + parquet_options: ParquetOptions, + ): + self.schema = schema + self.base_scan = base_scan + self.paths = paths + self.parquet_options = parquet_options + self._non_child_args = ( + 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, - context=context, ) + self.children = () + + 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.parquet_options, + ) + + @classmethod + 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: NamedExpr | None, + parquet_options: ParquetOptions, + *, + context: IRExecutionContext, + ) -> DataFrame: + """Evaluate and return a dataframe.""" + with nvtx_annotate_cudf_polars(message=f"FusedScan: {len(paths)} files"): + return Scan.do_evaluate( + schema, + typ, + reader_options, + paths, + with_columns, + skip_rows, + n_rows, + row_index, + include_file_paths, + predicate, + parquet_options, + context=context, + ) @lower_ir_node.register(Empty) @@ -484,10 +573,10 @@ class StreamingScan(IR): "base_scan", ) _n_non_child_args = 2 - scans: list[Scan | SplitScan] + scans: list[SplitScan | FusedScan] base_scan: Scan - def __init__(self, scans: list[Scan | SplitScan], base_scan: Scan): + def __init__(self, scans: list[SplitScan | FusedScan], base_scan: Scan): self.scans = scans self.base_scan = base_scan self.schema = base_scan.schema @@ -502,7 +591,7 @@ def get_hashable(self) -> Hashable: @classmethod def do_evaluate( cls, - scans: list[Scan | SplitScan], + scans: list[SplitScan | FusedScan], 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 c9ccb13202dc..7ea69d4296d9 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -14,7 +14,12 @@ from cudf_polars.dsl.ir import IRExecutionContext, Scan from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.base import IOPartitionFlavor, IOPartitionPlan -from cudf_polars.streaming.io import SplitScan, StreamingScan, expand_scan_for_rank +from cudf_polars.streaming.io import ( + FusedScan, + SplitScan, + StreamingScan, + expand_scan_for_rank, +) from cudf_polars.streaming.parallel import lower_ir_graph from cudf_polars.streaming.statistics import collect_statistics from cudf_polars.testing.asserts import assert_gpu_result_equal @@ -217,7 +222,7 @@ def test_expand_scan_for_rank_fused_and_single_read( parquet_options=ParquetOptions(), ) for scan, expected_paths in zip(scans, expected_path_groups, strict=True): - assert isinstance(scan, Scan) + assert isinstance(scan, FusedScan) assert scan.paths == expected_paths @@ -251,6 +256,7 @@ def test_expand_scan_for_rank_split_files( 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) ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): - StreamingScan.do_evaluate([scan], scan, context=ctx) + StreamingScan.do_evaluate([fused], scan, context=ctx) From b6bc1c13e0456af617fcb18d821ff1b3f6b33851 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Wed, 10 Jun 2026 17:32:10 +0000 Subject: [PATCH 2/4] Have both streaming scan class take paths --- .../cudf_polars/cudf_polars/streaming/io.py | 34 +++++++------------ .../cudf_polars/tests/streaming/test_scan.py | 6 ++-- 2 files changed, 16 insertions(+), 24 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 74006e601620..dccb77d887d0 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -146,25 +146,12 @@ def expand_scan_for_rank( sindex = local_offset % plan.factor splits_created = 0 for path in local_paths: - base_scan = Scan( - ir.schema, - ir.typ, - ir.reader_options, - ir.cloud_options, - [path], - ir.with_columns, - ir.skip_rows, - ir.n_rows, - ir.row_index, - ir.include_file_paths, - ir.predicate, - parquet_options, - ) while sindex < plan.factor and splits_created < local_count: scans.append( SplitScan( ir.schema, - base_scan, + ir, + [path], sindex, plan.factor, parquet_options, @@ -201,6 +188,7 @@ class SplitScan(IR): __slots__ = ( "base_scan", "parquet_options", + "paths", "schema", "split_index", "total_splits", @@ -208,6 +196,7 @@ class SplitScan(IR): _non_child = ( "schema", "base_scan", + "paths", "split_index", "total_splits", "parquet_options", @@ -215,6 +204,8 @@ class SplitScan(IR): _n_non_child_args = 13 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 @@ -226,12 +217,14 @@ def __init__( self, schema: Schema, base_scan: Scan, + paths: list[str], split_index: int, total_splits: int, parquet_options: ParquetOptions, ): self.schema = schema self.base_scan = base_scan + self.paths = paths self.split_index = split_index self.total_splits = total_splits self._non_child_args = ( @@ -240,14 +233,14 @@ def __init__( base_scan.schema, base_scan.typ, base_scan.reader_options, - base_scan.paths, + 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, - base_scan.parquet_options, + parquet_options, ) self.parquet_options = parquet_options self.children = () @@ -346,9 +339,7 @@ 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). The ``paths`` attribute holds the file group - assigned to this task; the remaining read options come from - ``base_scan``. + SINGLE_FILE (N = 1). """ __slots__ = ( @@ -425,7 +416,8 @@ def do_evaluate( context: IRExecutionContext, ) -> DataFrame: """Evaluate and return a dataframe.""" - with nvtx_annotate_cudf_polars(message=f"FusedScan: {len(paths)} files"): + names = ", ".join(Path(p).name for p in paths) + with nvtx_annotate_cudf_polars(message=f"FusedScan: {names}"): return Scan.do_evaluate( schema, typ, diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 7ea69d4296d9..5ddf0868513c 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -223,7 +223,7 @@ def test_expand_scan_for_rank_fused_and_single_read( ) for scan, expected_paths in zip(scans, expected_path_groups, strict=True): assert isinstance(scan, FusedScan) - assert scan.paths == expected_paths + assert scan.base_scan.paths == expected_paths @pytest.mark.parametrize( @@ -250,13 +250,13 @@ def test_expand_scan_for_rank_split_files( assert isinstance(scan, SplitScan) assert scan.split_index == split_index assert scan.total_splits == total_splits - assert scan.base_scan.paths == ["file.parquet"] + assert scan.paths == ["file.parquet"] 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.schema, scan, scan.parquet_options) ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): StreamingScan.do_evaluate([fused], scan, context=ctx) From c30e48829a029d69ca7488eb760bd14d1066b4f3 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Wed, 10 Jun 2026 17:34:25 +0000 Subject: [PATCH 3/4] missed --- python/cudf_polars/tests/streaming/test_scan.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 5ddf0868513c..db8492b2c6ee 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -256,7 +256,7 @@ def test_expand_scan_for_rank_split_files( 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.parquet_options) + fused = FusedScan(scan.schema, scan, scan.paths, scan.parquet_options) ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): StreamingScan.do_evaluate([fused], scan, context=ctx) From d3f4086a32eccb0414b4ed95d51d28a912c92da6 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Wed, 10 Jun 2026 22:58:26 +0000 Subject: [PATCH 4/4] address review, add absolute path prefix to nvtx annotation --- python/cudf_polars/cudf_polars/streaming/io.py | 17 ++++++++++++++--- python/cudf_polars/tests/streaming/test_scan.py | 2 +- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index dccb77d887d0..4162a08791d1 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -249,6 +249,18 @@ def __init__( f"Unhandled Scan type for file splitting: {base_scan.typ}" ) + 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, + ) + @classmethod def do_evaluate( cls, @@ -316,7 +328,7 @@ def do_evaluate( # Perform the partial read with nvtx_annotate_cudf_polars( - message=f"SplitScan: {Path(paths[0]).name} [{split_index + 1}/{total_splits}]" + message=f"SplitScan: {paths[0]} [{split_index + 1}/{total_splits}]" ): return Scan.do_evaluate( schema, @@ -416,8 +428,7 @@ def do_evaluate( context: IRExecutionContext, ) -> DataFrame: """Evaluate and return a dataframe.""" - names = ", ".join(Path(p).name for p in paths) - with nvtx_annotate_cudf_polars(message=f"FusedScan: {names}"): + with nvtx_annotate_cudf_polars(message=f"FusedScan: {', '.join(paths)}"): return Scan.do_evaluate( schema, typ, diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index db8492b2c6ee..cf6e14d31bca 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -223,7 +223,7 @@ def test_expand_scan_for_rank_fused_and_single_read( ) for scan, expected_paths in zip(scans, expected_path_groups, strict=True): assert isinstance(scan, FusedScan) - assert scan.base_scan.paths == expected_paths + assert scan.paths == expected_paths @pytest.mark.parametrize(