Add FusedScan and nxtx annotations for FusedScan and SplitScan

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):

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)
@@ -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,
@@ -182,23 +169,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 +179,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.
@@ -216,20 +188,24 @@ class SplitScan(IR):
     __slots__ = (
         "base_scan",
         "parquet_options",
+        "paths",
         "schema",
         "split_index",
         "total_splits",
     )
     _non_child = (
         "schema",
         "base_scan",
+        "paths",
         "split_index",
         "total_splits",
         "parquet_options",
     )
     _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
@@ -241,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 = (
@@ -255,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 = ()
@@ -271,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,
@@ -337,20 +327,122 @@ 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: {paths[0]} [{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).
+    """
+
+    __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: {', '.join(paths)}"):
+            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 +576,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 +594,7 @@ def get_hashable(self) -> Hashable:
     @classmethod
     def do_evaluate(
         cls,
-        scans: list[Scan | SplitScan],
+        scans: list[SplitScan | FusedScan],
         base_scan: Scan,
         *,
         context: IRExecutionContext,

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
 
 
@@ -245,12 +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)
     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)