Refactor streaming scan construction into StreamingScan static methods

python/cudf_polars/cudf_polars/streaming/actor_graph/io.py

@@ -51,6 +51,8 @@
 )
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     from rapidsmpf.communicator.communicator import Communicator
     from rapidsmpf.streaming.core.channel import Channel
     from rapidsmpf.streaming.core.context import Context
@@ -370,7 +372,7 @@ async def scan_node(
         Estimated size of each chunk in bytes. Used for memory reservation
         with block spilling to avoid thrashing.
     """
-    scans = ir.scans
+    scans: Sequence[SplitScan] | Sequence[FusedScan] = ir.scans
 
     async with shutdown_on_error(
         context, ch_out, trace_ir=ir, ir_context=ir_context
@@ -413,7 +415,8 @@ async def scan_node(
         ]
         for task_idx, scan in enumerate(scans):
             producer_id = task_idx % num_producers
-            producer_tasks[producer_id].append((task_idx, scan))
+            # mypy resolves __iter__ on union-of-sequences to the common base (IR)
+            producer_tasks[producer_id].append((task_idx, scan))  # type: ignore[arg-type]
 
         async def _producer(producer_id: int, ch_out: Channel) -> None:
             for task_idx, scan in producer_tasks[producer_id]:

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, overload
+from typing import TYPE_CHECKING, Any, Literal, Self, overload
 
 import polars as pl
 
@@ -38,7 +38,7 @@
 from cudf_polars.utils.versions import POLARS_VERSION_LT_137
 
 if TYPE_CHECKING:
-    from collections.abc import Hashable, MutableMapping
+    from collections.abc import Hashable, MutableMapping, Sequence
 
     from cudf_polars.containers import DataFrame
     from cudf_polars.dsl.expr import NamedExpr
@@ -109,23 +109,32 @@ def scan_partition_plan(
     return IOPartitionPlan(1, IOPartitionFlavor.SINGLE_FILE)
 
 
+def _rank_slice(total: int, rank: int, nranks: int) -> tuple[int, int]:
+    """Return the partition range owned by this rank."""
+    count = math.ceil(total / nranks)
+    return count * rank, count
+
+
 def expand_scan_for_rank(
     ir: Scan,
     plan: IOPartitionPlan,
+    partition_count: int,
     *,
     rank: int,
     nranks: int,
     parquet_options: ParquetOptions,
-) -> list[SplitScan | FusedScan]:
+) -> StreamingScan:
     """
-    Expand a Scan node into rank-local SplitScan and FusedScan operations.
+    Expand a Scan node into a rank-local StreamingScan.
 
     Parameters
     ----------
     ir
         The Scan node to expand.
     plan
         The IO partitioning plan for the scan.
+    partition_count
+        Total number of partitions across all ranks.
     rank
         Rank of the current worker.
     nranks
@@ -135,48 +144,27 @@ def expand_scan_for_rank(
 
     Returns
     -------
-    list[SplitScan | FusedScan]
-        Rank-local scan operations.
+    StreamingScan
+        Rank-local streaming scan.
     """
-    scans: list[SplitScan | FusedScan] = []
     if plan.flavor == IOPartitionFlavor.SPLIT_FILES:
-        count = plan.factor * len(ir.paths)
-        local_count = math.ceil(count / nranks)
-        local_offset = local_count * rank
-        path_offset = local_offset // plan.factor
-        path_end = math.ceil((local_offset + local_count) / plan.factor)
-        path_count = path_end - path_offset
-        local_paths = ir.paths[path_offset : path_offset + path_count]
-        sindex = local_offset % plan.factor
-        splits_created = 0
-        for path in local_paths:
-            while sindex < plan.factor and splits_created < local_count:
-                scans.append(
-                    SplitScan(
-                        ir.schema,
-                        ir,
-                        [path],
-                        sindex,
-                        plan.factor,
-                        parquet_options,
-                    )
-                )
-                sindex += 1
-                splits_created += 1
-            sindex = 0
-
+        return StreamingScan.for_split_files(
+            ir,
+            plan,
+            partition_count,
+            rank=rank,
+            nranks=nranks,
+            parquet_options=parquet_options,
+        )
     else:
-        count = math.ceil(len(ir.paths) / plan.factor)
-        local_count = math.ceil(count / nranks)
-        local_offset = local_count * rank
-        paths_offset_start = local_offset * plan.factor
-        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:
-                scans.append(FusedScan(ir.schema, ir, local_paths, parquet_options))
-
-    return scans
+        return StreamingScan.for_fused_files(
+            ir,
+            plan,
+            partition_count,
+            rank=rank,
+            nranks=nranks,
+            parquet_options=parquet_options,
+        )
 
 
 class SplitScan(IR):
@@ -556,14 +544,14 @@ def _(
     ):
         parquet_options = dataclasses.replace(parquet_options, chunked=False)
 
-    scans = expand_scan_for_rank(
+    new_ir = expand_scan_for_rank(
         ir,
         plan,
+        count,
         rank=rec.state["rank"],
         nranks=rec.state["nranks"],
         parquet_options=parquet_options,
     )
-    new_ir = StreamingScan(scans, ir)
     return new_ir, {new_ir: PartitionInfo(count=count, io_plan=plan)}
 
 
@@ -580,16 +568,81 @@ class StreamingScan(IR):
         "base_scan",
     )
     _n_non_child_args = 2
-    scans: list[SplitScan | FusedScan]
+    scans: Sequence[SplitScan] | Sequence[FusedScan]
     base_scan: Scan
 
-    def __init__(self, scans: list[SplitScan | FusedScan], base_scan: Scan):
+    def __init__(
+        self, scans: Sequence[SplitScan] | Sequence[FusedScan], base_scan: Scan
+    ):
         self.scans = scans
         self.base_scan = base_scan
         self.schema = base_scan.schema
         self._non_child_args = (scans, 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
+        scans: list[SplitScan] = []
+        splits_created = 0
+        for path in local_paths:
+            while sindex < plan.factor and splits_created < local_count:
+                scans.append(
+                    SplitScan(
+                        base_scan.schema,
+                        base_scan,
+                        [path],
+                        sindex,
+                        plan.factor,
+                        parquet_options,
+                    )
+                )
+                sindex += 1
+                splits_created += 1
+            sindex = 0
+        return cls(scans, 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
+        scans = [
+            FusedScan(
+                base_scan.schema,
+                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]
+        ]
+        return cls(scans, 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.
@@ -598,7 +651,7 @@ def get_hashable(self) -> Hashable:
     @classmethod
     def do_evaluate(
         cls,
-        scans: list[SplitScan | FusedScan],
+        scans: Sequence[SplitScan] | Sequence[FusedScan],
         base_scan: Scan,
         *,
         context: IRExecutionContext,

python/cudf_polars/tests/streaming/test_scan.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import math
 from typing import TYPE_CHECKING
 
 import pytest
@@ -231,14 +232,18 @@ def test_expand_scan_for_rank_fused_and_single_read(
     nranks: int,
     expected_path_groups: list[list[str]],
 ) -> None:
-    scans = expand_scan_for_rank(
+    partition_count = math.ceil(len(paths) / plan.factor)
+    streaming_scan = expand_scan_for_rank(
         _make_parquet_scan(paths),
         plan,
+        partition_count,
         rank=rank,
         nranks=nranks,
         parquet_options=ParquetOptions(),
     )
-    for scan, expected_paths in zip(scans, expected_path_groups, strict=True):
+    for scan, expected_paths in zip(
+        streaming_scan.scans, expected_path_groups, strict=True
+    ):
         assert isinstance(scan, FusedScan)
         assert scan.paths == expected_paths
 
@@ -255,15 +260,20 @@ def test_expand_scan_for_rank_split_files(
     expected_splits: list[tuple[int, int]],
 ) -> None:
     plan = IOPartitionPlan(4, IOPartitionFlavor.SPLIT_FILES)
-    scans = expand_scan_for_rank(
-        _make_parquet_scan(["file.parquet"]),
+    paths = ["file.parquet"]
+    partition_count = plan.factor * len(paths)
+    streaming_scan = expand_scan_for_rank(
+        _make_parquet_scan(paths),
         plan,
+        partition_count,
         rank=rank,
         nranks=2,
         parquet_options=ParquetOptions(),
     )
-    assert len(scans) == len(expected_splits)
-    for scan, (split_index, total_splits) in zip(scans, expected_splits, strict=True):
+    assert len(streaming_scan.scans) == len(expected_splits)
+    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