From a4477755fbb18d86b84b5c83fd9913b498e81618 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Tue, 2 Jun 2026 14:47:55 -0700
Subject: [PATCH 01/20] Refactor dynamic Scan node lowering

This refactors how we dynamically generate `Scan` nodes with the
rapidsmpf streaming runtime.

Previously, these nodes were generated inside
`cudf_polars.streaming.actor_grpah.io.scan_node`. Now, we treat it more
like a (rank-specific) lowering stage.

This is primarily motivated by awkwardness in using the dynamically
generated Scan nodes in a couple spots:

1. prefetching parquet metadata(https://github.com/rapidsai/cudf/pull/22700)
2. Quent tracing

I've added a new `StreamingScan` IR node that just holds references to
the `list[Scan]` nodes. This might be overkill for what we need, but
it's at least consistent with how we do the main lowering, which I like.
---
 python/cudf_polars/cudf_polars/engine/core.py |   7 +
 .../cudf_polars/streaming/actor_graph/io.py   | 258 ++++++++++++------
 .../cudf_polars/cudf_polars/streaming/io.py   |  28 ++
 3 files changed, 211 insertions(+), 82 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py
index 5fddabcf78df..7d703b642339 100644
--- a/python/cudf_polars/cudf_polars/engine/core.py
+++ b/python/cudf_polars/cudf_polars/engine/core.py
@@ -692,6 +692,13 @@ def evaluate_on_rank(
         # so we only log it once.
         log_query_plan(ir, config_options)
 
+    from cudf_polars.streaming.actor_graph.io import io_lower_ir_graph
+
+    ir, partition_info = io_lower_ir_graph(
+        ir, partition_info, comm, config_options.parquet_options
+    )
+    # TODO: log this query plan too?
+
     with ReserveOpIDs(ir, config_options) as collective_id_map:
         return execute_ir_on_rank(
             ctx,
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 78672bb0cab0..c4a7ee0cbd90 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
@@ -7,6 +7,7 @@
 import asyncio
 import dataclasses
 import math
+from functools import singledispatch
 from typing import TYPE_CHECKING, Any
 
 from rapidsmpf.memory.memory_reservation import opaque_memory_usage
@@ -49,6 +50,7 @@
 )
 from cudf_polars.streaming.io import (
     SplitScan,
+    StreamingScan,
     StreamingSink,
     _prepare_sink_directory,
     _sink_to_file,
@@ -56,6 +58,8 @@
 from cudf_polars.streaming.utils import _dynamic_planning_on
 
 if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
     from rapidsmpf.communicator.communicator import Communicator
     from rapidsmpf.streaming.core.channel import Channel
     from rapidsmpf.streaming.core.context import Context
@@ -68,7 +72,173 @@
         PartitionInfo,
         StatsCollector,
     )
-    from cudf_polars.utils.config import ParquetOptions
+    from cudf_polars.utils.config import (
+        ParquetOptions,
+    )
+
+from typing import TYPE_CHECKING, TypeAlias, TypedDict
+
+from cudf_polars.dsl.traversal import CachingVisitor
+from cudf_polars.typing import GenericTransformer
+
+
+class IOLowerIRState(TypedDict):
+    """State used for lowering IR nodes."""
+
+    partition_info: MutableMapping[IR, PartitionInfo]
+    rank: int
+    nranks: int
+    parquet_options: ParquetOptions
+
+    # config_options: ConfigOptions[StreamingExecutor]
+    # stats: StatsCollector
+
+
+IOLowerIRTransformer: TypeAlias = GenericTransformer[
+    "IR", "tuple[IR, MutableMapping[IR, PartitionInfo]]", IOLowerIRState
+]
+
+
+def io_lower_ir_graph(
+    ir: IR,
+    partition_info: MutableMapping[IR, PartitionInfo],
+    comm: Communicator,
+    parquet_options: ParquetOptions,
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    """
+    Rewrite an IR graph and extract partitioning information.
+
+    Parameters
+    ----------
+    ir
+        Root of the graph to rewrite.
+    partition_info
+        Partition information for the graph. Mutate this in place for new IR nodes.
+    comm
+        Communicator for the graph. The output IR graph will contain nodes unique
+        to this rank.
+    parquet_options
+        Parquet reader options.
+
+    Returns
+    -------
+    new_ir, partition_info
+        The rewritten graph and a mapping from unique nodes
+        in the new graph to associated partitioning information.
+
+    Notes
+    -----
+    This function traverses the unique nodes of the graph with
+    root `ir`, and applies :func:`lower_ir_node` to each node.
+
+    See Also
+    --------
+    lower_ir_node
+    """
+    state: IOLowerIRState = {
+        "partition_info": partition_info,
+        "rank": comm.rank,
+        "nranks": comm.nranks,
+        "parquet_options": parquet_options,
+    }
+    mapper: IOLowerIRTransformer = CachingVisitor(io_lower_ir_node, state=state)
+    return mapper(ir)
+
+
+@singledispatch
+def io_lower_ir_node(
+    ir: IR, rec: IOLowerIRTransformer
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    """Lower the IR nodes for a given IR node."""
+    return ir, rec.state["partition_info"]
+
+
+@io_lower_ir_node.register(Scan)
+def _(
+    ir: Scan, rec: IOLowerIRTransformer
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    """
+    Lower the Scan node.
+
+    This expands SplitScan nodes into multiple Scan nodes.
+    """
+    plan = rec.state["partition_info"][ir].io_plan
+    assert plan is not None  # not great...
+
+    rank = rec.state["rank"]
+    nranks = rec.state["nranks"]
+    parquet_options = rec.state["parquet_options"]
+
+    scans: list[Scan | SplitScan] = []
+    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:
+            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,
+                        sindex,
+                        plan.factor,
+                        parquet_options,
+                    )
+                )
+                sindex += 1
+                splits_created += 1
+            sindex = 0
+
+    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:  # 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,
+                    )
+                )
+
+    # I have no idea if this is correct
+    new_ir = StreamingScan(scans, ir.schema)
+    rec.state["partition_info"][new_ir] = rec.state["partition_info"][ir]
+    return new_ir, rec.state["partition_info"]
 
 
 class Lineariser:
@@ -356,14 +526,11 @@ async def read_chunk(
 @define_actor()
 async def scan_node(
     context: Context,
-    comm: Communicator,
-    ir: Scan,
+    ir: StreamingScan,
     ir_context: IRExecutionContext,
     ch_out: Channel[TableChunk],
     *,
     num_producers: int,
-    plan: IOPartitionPlan,
-    parquet_options: ParquetOptions,
     estimated_chunk_bytes: int,
 ) -> None:
     """
@@ -373,8 +540,6 @@ async def scan_node(
     ----------
     context
         The rapidsmpf context.
-    comm
-        The communicator.
     ir
         The Scan node.
     ir_context
@@ -383,84 +548,15 @@ async def scan_node(
         The output Channel[TableChunk].
     num_producers
         The number of producers to use for the scan node.
-    plan
-        The partitioning plan.
-    parquet_options
-        The Parquet options.
     estimated_chunk_bytes
         Estimated size of each chunk in bytes. Used for memory reservation
         with block spilling to avoid thrashing.
     """
+    scans = ir.scans
+
     async with shutdown_on_error(
         context, ch_out, trace_ir=ir, ir_context=ir_context
     ) as tracer:
-        # Build a list of local Scan operations
-        scans: list[Scan | SplitScan] = []
-        if plan.flavor == IOPartitionFlavor.SPLIT_FILES:
-            count = plan.factor * len(ir.paths)
-            local_count = math.ceil(count / comm.nranks)
-            local_offset = local_count * comm.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:
-                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,
-                            sindex,
-                            plan.factor,
-                            parquet_options,
-                        )
-                    )
-                    sindex += 1
-                    splits_created += 1
-                sindex = 0
-
-        else:
-            count = math.ceil(len(ir.paths) / plan.factor)
-            local_count = math.ceil(count / comm.nranks)
-            local_offset = local_count * comm.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:  # 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,
-                        )
-                    )
-
         # Send basic metadata
         await send_metadata(
             ch_out,
@@ -628,7 +724,8 @@ def make_rapidsmpf_read_parquet_node(
         ) from e
 
 
-@generate_ir_sub_network.register(Scan)
+@generate_ir_sub_network.register(Scan)  # TODO: see if this even is hit?
+@generate_ir_sub_network.register(StreamingScan)
 def _(
     ir: Scan, rec: SubNetGenerator
 ) -> tuple[dict[IR, list[Any]], dict[IR, ChannelManager]]:
@@ -697,13 +794,10 @@ def _(
         nodes[ir] = [
             scan_node(
                 rec.state["context"],
-                rec.state["comm"],
                 ir,
                 rec.state["ir_context"],
                 ch_out,
                 num_producers=num_producers,
-                plan=plan,
-                parquet_options=parquet_options,
                 estimated_chunk_bytes=executor.target_partition_size,
             )
         ]
diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py
index 4472735ac356..b9501dd3e6b5 100644
--- a/python/cudf_polars/cudf_polars/streaming/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/io.py
@@ -291,6 +291,34 @@ def _(
         return ir, {ir: PartitionInfo(count=1, io_plan=plan)}
 
 
+class StreamingScan(IR):
+    """A streaming scan node."""
+
+    __slots__ = (
+        "scans",
+        "schema",
+    )
+    _non_child = (
+        "scans",
+        "schema",
+    )
+    _n_non_child_args = 2  # TODO: verify this
+    scans: list[Scan | SplitScan]
+    schema: Schema
+
+    def __init__(self, scans: list[Scan | SplitScan], schema: Schema):
+        self.scans = scans
+        self.schema = schema
+        self._non_child_args = (scans,)
+        self.children = ()
+
+    def get_hashable(self) -> Hashable:
+        """Hashable representation of the node."""
+        # We don't need to include schema, since it's in all the base scan nodes.
+        # TODO: Why do we have it in the first place?
+        return (type(self), *tuple(x.get_hashable() for x in self.scans))
+
+
 class StreamingSink(IR):
     """Sink a dataframe in streaming mode."""
 

From ece73249f97a61b3d6d8aebe1f03c71eb8253fd8 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Wed, 3 Jun 2026 04:48:04 -0700
Subject: [PATCH 02/20] fixes

---
 python/cudf_polars/cudf_polars/engine/core.py |  4 +-
 .../cudf_polars/streaming/actor_graph/io.py   | 51 +++++++++++++++----
 python/cudf_polars/tests/test_scan.py         | 21 ++++----
 3 files changed, 53 insertions(+), 23 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py
index 7d703b642339..4e416d999154 100644
--- a/python/cudf_polars/cudf_polars/engine/core.py
+++ b/python/cudf_polars/cudf_polars/engine/core.py
@@ -28,6 +28,7 @@
 from cudf_polars.streaming.actor_graph.collectives import ReserveOpIDs
 from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id
 from cudf_polars.streaming.actor_graph.core import generate_network
+from cudf_polars.streaming.actor_graph.io import io_lower_ir_graph
 from cudf_polars.streaming.actor_graph.tracing import log_query_plan
 from cudf_polars.streaming.actor_graph.utils import empty_table_chunk
 from cudf_polars.streaming.base import StatsCollector
@@ -692,12 +693,9 @@ def evaluate_on_rank(
         # so we only log it once.
         log_query_plan(ir, config_options)
 
-    from cudf_polars.streaming.actor_graph.io import io_lower_ir_graph
-
     ir, partition_info = io_lower_ir_graph(
         ir, partition_info, comm, config_options.parquet_options
     )
-    # TODO: log this query plan too?
 
     with ReserveOpIDs(ir, config_options) as collective_id_map:
         return execute_ir_on_rank(
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 c4a7ee0cbd90..c250bdbafcf1 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
@@ -7,8 +7,9 @@
 import asyncio
 import dataclasses
 import math
-from functools import singledispatch
-from typing import TYPE_CHECKING, Any
+import operator
+from functools import reduce, singledispatch
+from typing import TYPE_CHECKING, Any, TypeAlias, TypedDict
 
 from rapidsmpf.memory.memory_reservation import opaque_memory_usage
 from rapidsmpf.streaming.core.memory_reserve_or_wait import (
@@ -25,6 +26,7 @@
     DataFrameScan,
     Scan,
     Sink,
+    Union,
     _prepare_parquet_predicate,
 )
 from cudf_polars.dsl.to_ast import to_parquet_filter
@@ -47,6 +49,7 @@
 )
 from cudf_polars.streaming.base import (
     IOPartitionFlavor,
+    PartitionInfo,
 )
 from cudf_polars.streaming.io import (
     SplitScan,
@@ -69,15 +72,12 @@
     from cudf_polars.streaming.actor_graph.tracing import ActorTracer
     from cudf_polars.streaming.base import (
         IOPartitionPlan,
-        PartitionInfo,
         StatsCollector,
     )
     from cudf_polars.utils.config import (
         ParquetOptions,
     )
 
-from typing import TYPE_CHECKING, TypeAlias, TypedDict
-
 from cudf_polars.dsl.traversal import CachingVisitor
 from cudf_polars.typing import GenericTransformer
 
@@ -90,9 +90,6 @@ class IOLowerIRState(TypedDict):
     nranks: int
     parquet_options: ParquetOptions
 
-    # config_options: ConfigOptions[StreamingExecutor]
-    # stats: StatsCollector
-
 
 IOLowerIRTransformer: TypeAlias = GenericTransformer[
     "IR", "tuple[IR, MutableMapping[IR, PartitionInfo]]", IOLowerIRState
@@ -145,18 +142,50 @@ def io_lower_ir_graph(
     return mapper(ir)
 
 
+def _io_lower_pwise(
+    ir: IR, rec: IOLowerIRTransformer
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    """Lower children and preserve partitioning from a single child."""
+    children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
+    partition_info = reduce(operator.or_, _partition_info)
+
+    if len(children) == 1:
+        partition = partition_info[children[0]]
+    else:
+        partition = PartitionInfo(count=max(partition_info[c].count for c in children))
+
+    new_node = ir.reconstruct(children)
+    partition_info[new_node] = partition
+    return new_node, partition_info
+
+
 @singledispatch
 def io_lower_ir_node(
     ir: IR, rec: IOLowerIRTransformer
 ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
     """Lower the IR nodes for a given IR node."""
-    return ir, rec.state["partition_info"]
+    if not ir.children:
+        return ir, rec.state["partition_info"]
+    return _io_lower_pwise(ir, rec)
+
+
+@io_lower_ir_node.register(Union)
+def _(
+    ir: Union, rec: IOLowerIRTransformer
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    # TODO: Determine if we really need this. The only difference is `sum` instead of `max` for getting count.
+    children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
+    partition_info = reduce(operator.or_, _partition_info)
+    count = sum(partition_info[c].count for c in children)
+    new_node = ir.reconstruct(children)
+    partition_info[new_node] = PartitionInfo(count=count)
+    return new_node, partition_info
 
 
 @io_lower_ir_node.register(Scan)
 def _(
     ir: Scan, rec: IOLowerIRTransformer
-) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+) -> tuple[StreamingScan, MutableMapping[IR, PartitionInfo]]:
     """
     Lower the Scan node.
 
@@ -724,7 +753,7 @@ def make_rapidsmpf_read_parquet_node(
         ) from e
 
 
-@generate_ir_sub_network.register(Scan)  # TODO: see if this even is hit?
+# @generate_ir_sub_network.register(Scan)  # TODO: see if this even is hit?
 @generate_ir_sub_network.register(StreamingScan)
 def _(
     ir: Scan, rec: SubNetGenerator
diff --git a/python/cudf_polars/tests/test_scan.py b/python/cudf_polars/tests/test_scan.py
index 1dca93dd38ae..ef8f02e196b9 100644
--- a/python/cudf_polars/tests/test_scan.py
+++ b/python/cudf_polars/tests/test_scan.py
@@ -43,7 +43,7 @@
     params=[(None, None), ("row-index", 0), ("index", 10)],
     ids=["no_row_index", "zero_offset_row_index", "offset_row_index"],
 )
-def row_index(request):
+def row_index(request) -> tuple[str | None, int | None]:
     return request.param
 
 
@@ -474,13 +474,14 @@ def test_select_arbitrary_order_with_row_index_column(engine: pl.GPUEngine, tmp_
 )
 def test_scan_csv_with_and_without_header(
     engine: pl.GPUEngine,
-    df,
-    tmp_path,
-    has_header,
-    new_columns,
-    row_index,
-    columns,
-    zlice,
+    df: pl.DataFrame,
+    tmp_path: Path,
+    *,
+    has_header: bool,
+    new_columns: list[str] | None,
+    row_index: tuple[str | None, int | None],
+    columns: list[str] | None,
+    zlice: tuple[int, int] | None,
 ):
     path = tmp_path / "test.csv"
     make_partitioned_source(
@@ -489,12 +490,14 @@ def test_scan_csv_with_and_without_header(
 
     name, offset = row_index
 
+    # offset is only used if name is used, but that's not an overload
+    # in the type signature.
     q = pl.scan_csv(
         path,
         has_header=has_header,
         new_columns=new_columns,
         row_index_name=name,
-        row_index_offset=offset,
+        row_index_offset=offset,  # type: ignore[arg-type]
     )
 
     if zlice is not None:

From 38eacf1b84ce3a9381a441144d1647b64096c6d5 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Wed, 3 Jun 2026 06:15:24 -0700
Subject: [PATCH 03/20] fixes

---
 .../cudf_polars/streaming/actor_graph/io.py     | 17 ++++++++---------
 python/cudf_polars/cudf_polars/streaming/io.py  | 16 +++++++++-------
 2 files changed, 17 insertions(+), 16 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 c250bdbafcf1..beb3d65419e9 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
@@ -265,7 +265,7 @@ def _(
                 )
 
     # I have no idea if this is correct
-    new_ir = StreamingScan(scans, ir.schema)
+    new_ir = StreamingScan(scans, ir)
     rec.state["partition_info"][new_ir] = rec.state["partition_info"][ir]
     return new_ir, rec.state["partition_info"]
 
@@ -753,10 +753,9 @@ def make_rapidsmpf_read_parquet_node(
         ) from e
 
 
-# @generate_ir_sub_network.register(Scan)  # TODO: see if this even is hit?
 @generate_ir_sub_network.register(StreamingScan)
 def _(
-    ir: Scan, rec: SubNetGenerator
+    ir: StreamingScan, rec: SubNetGenerator
 ) -> tuple[dict[IR, list[Any]], dict[IR, ChannelManager]]:
     config_options = rec.state["config_options"]
     executor = config_options.executor
@@ -781,11 +780,11 @@ def _(
     if (
         parquet_options.use_rapidsmpf_native
         and (partition_info.count > 1 or _dynamic_planning_on(config_options))
-        and ir.typ == "parquet"
-        and ir.row_index is None
-        and ir.include_file_paths is None
-        and ir.n_rows == -1
-        and ir.skip_rows == 0
+        and ir.base_scan.typ == "parquet"
+        and ir.base_scan.row_index is None
+        and ir.base_scan.include_file_paths is None
+        and ir.base_scan.n_rows == -1
+        and ir.base_scan.skip_rows == 0
         and not distributed_split_files
     ):
         # Create new channel to so ch_out can be used to add metadata
@@ -793,7 +792,7 @@ def _(
         native_node = make_rapidsmpf_read_parquet_node(
             rec.state["context"],
             rec.state["comm"],
-            ir,
+            ir.base_scan,  # TODO: verify this...
             num_producers,
             ch_in,
             rec.state["stats"],
diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py
index b9501dd3e6b5..0ebbe29eeddb 100644
--- a/python/cudf_polars/cudf_polars/streaming/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/io.py
@@ -295,26 +295,28 @@ class StreamingScan(IR):
     """A streaming scan node."""
 
     __slots__ = (
+        "base_scan",
         "scans",
-        "schema",
     )
     _non_child = (
         "scans",
-        "schema",
+        "base_scan",
     )
     _n_non_child_args = 2  # TODO: verify this
     scans: list[Scan | SplitScan]
-    schema: Schema
+    # These are essentially the base scan properties, shared by all the scans
+    # schema: Schema
+    base_scan: Scan
 
-    def __init__(self, scans: list[Scan | SplitScan], schema: Schema):
+    def __init__(self, scans: list[Scan | SplitScan], base_scan: Scan):
         self.scans = scans
-        self.schema = schema
-        self._non_child_args = (scans,)
+        self.base_scan = base_scan
+        self._non_child_args = (scans, base_scan)
         self.children = ()
 
     def get_hashable(self) -> Hashable:
         """Hashable representation of the node."""
-        # We don't need to include schema, since it's in all the base scan nodes.
+        # We don't need to include base_scan, since it's in all the scan nodes.
         # TODO: Why do we have it in the first place?
         return (type(self), *tuple(x.get_hashable() for x in self.scans))
 

From 589bb6a047d2e7edb24aa740b3e20b1ce00c5ae2 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Wed, 3 Jun 2026 07:39:45 -0700
Subject: [PATCH 04/20] New tests

---
 .../cudf_polars/cudf_polars/streaming/io.py   |  5 +++-
 .../cudf_polars/tests/streaming/test_scan.py  | 27 +++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py
index 0ebbe29eeddb..05e81b5b277a 100644
--- a/python/cudf_polars/cudf_polars/streaming/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/io.py
@@ -297,10 +297,12 @@ class StreamingScan(IR):
     __slots__ = (
         "base_scan",
         "scans",
+        "schema",
     )
     _non_child = (
         "scans",
         "base_scan",
+        "schema",
     )
     _n_non_child_args = 2  # TODO: verify this
     scans: list[Scan | SplitScan]
@@ -311,12 +313,13 @@ class StreamingScan(IR):
     def __init__(self, scans: list[Scan | SplitScan], 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 = ()
 
     def get_hashable(self) -> Hashable:
         """Hashable representation of the node."""
-        # We don't need to include base_scan, since it's in all the scan nodes.
+        # We don't need to include base_scan / schema, since it's in all the scan nodes.
         # TODO: Why do we have it in the first place?
         return (type(self), *tuple(x.get_hashable() for x in self.scans))
 
diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py
index 11afaac838ac..9d08a2362883 100644
--- a/python/cudf_polars/tests/streaming/test_scan.py
+++ b/python/cudf_polars/tests/streaming/test_scan.py
@@ -19,6 +19,7 @@
 
 if TYPE_CHECKING:
     import concurrent.futures
+    from pathlib import Path
 
 
 @pytest.fixture(scope="module")
@@ -125,3 +126,29 @@ def test_target_partition_size(
         assert count > n_files
     else:
         assert count < n_files
+
+
+def test_scan_join(engine: pl.GPUEngine, tmp_path: Path) -> None:
+    # This test exercises some logic on nodes with multiple children (join)
+    # where one or more of the children are Scan nodes.
+    left = pl.DataFrame({"a": ["a", "b", "c", "d"], "b": [1, 2, 3, 4]})
+    right = pl.DataFrame({"a": ["a", "b", "c", "d"], "c": [10, 20, 30, 40]})
+
+    left.write_parquet(tmp_path / "left.parquet")
+    right.write_parquet(tmp_path / "right.parquet")
+
+    left_q = pl.scan_parquet(tmp_path / "left.parquet")
+    right_q = pl.scan_parquet(tmp_path / "right.parquet")
+    q = left_q.join(right_q, on="a", how="inner")
+    assert_gpu_result_equal(q, engine=engine)
+
+
+def test_scan_union(engine: pl.GPUEngine, tmp_path: Path) -> None:
+    # This test exercises some logic on nodes with a Union[Scan, ...]
+    df = pl.DataFrame({"a": ["a", "b", "c", "d"], "b": [1, 2, 3, 4]})
+    df.write_parquet(tmp_path / "data.parquet")
+
+    df_q = pl.scan_parquet(tmp_path / "data.parquet")
+
+    q = pl.concat([df_q, df_q])
+    assert_gpu_result_equal(q, engine=engine)

From dbe3446d98a34103c5a67f1b8a136819e03c2e88 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Wed, 3 Jun 2026 07:51:26 -0700
Subject: [PATCH 05/20] update io_plan handling

---
 python/cudf_polars/cudf_polars/streaming/actor_graph/io.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

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 beb3d65419e9..2ae44eab281f 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
@@ -192,7 +192,11 @@ def _(
     This expands SplitScan nodes into multiple Scan nodes.
     """
     plan = rec.state["partition_info"][ir].io_plan
-    assert plan is not None  # not great...
+    # rec.state["partition_info"] is just a mapping from IR nodes to PartitionInfo.
+    # We promise that PartitionInfo.io_plan is not None for Scan nodes,
+    # but don't currently enforce that promise in the type system.
+    if plan is None:
+        raise RuntimeError(f"Scan node must have a partition plan. node={ir}")
 
     rank = rec.state["rank"]
     nranks = rec.state["nranks"]

From e1a4b4cd7e064869ea6b970facb1a7f9ffc27981 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Wed, 3 Jun 2026 07:53:56 -0700
Subject: [PATCH 06/20] cleanup

---
 python/cudf_polars/cudf_polars/streaming/io.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py
index 05e81b5b277a..f810cc1c7319 100644
--- a/python/cudf_polars/cudf_polars/streaming/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/io.py
@@ -304,7 +304,7 @@ class StreamingScan(IR):
         "base_scan",
         "schema",
     )
-    _n_non_child_args = 2  # TODO: verify this
+    _n_non_child_args = 2
     scans: list[Scan | SplitScan]
     # These are essentially the base scan properties, shared by all the scans
     # schema: Schema

From 56de96ec0c4fe1634fdf7d0345f6809728832847 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Wed, 3 Jun 2026 08:17:07 -0700
Subject: [PATCH 07/20] inline io_lower_pwise

---
 .../cudf_polars/streaming/actor_graph/io.py   | 20 +++++++------------
 .../cudf_polars/cudf_polars/streaming/io.py   |  3 ---
 2 files changed, 7 insertions(+), 16 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 2ae44eab281f..f9975f1c0a16 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
@@ -142,10 +142,14 @@ def io_lower_ir_graph(
     return mapper(ir)
 
 
-def _io_lower_pwise(
+@singledispatch
+def io_lower_ir_node(
     ir: IR, rec: IOLowerIRTransformer
 ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
-    """Lower children and preserve partitioning from a single child."""
+    """Lower the IR nodes for a given IR node."""
+    if not ir.children:
+        return ir, rec.state["partition_info"]
+
     children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
     partition_info = reduce(operator.or_, _partition_info)
 
@@ -159,16 +163,6 @@ def _io_lower_pwise(
     return new_node, partition_info
 
 
-@singledispatch
-def io_lower_ir_node(
-    ir: IR, rec: IOLowerIRTransformer
-) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
-    """Lower the IR nodes for a given IR node."""
-    if not ir.children:
-        return ir, rec.state["partition_info"]
-    return _io_lower_pwise(ir, rec)
-
-
 @io_lower_ir_node.register(Union)
 def _(
     ir: Union, rec: IOLowerIRTransformer
@@ -268,7 +262,7 @@ def _(
                     )
                 )
 
-    # I have no idea if this is correct
+    # The PartitionInfo for the derived StreamingScan node is the same as the original Scan node.
     new_ir = StreamingScan(scans, ir)
     rec.state["partition_info"][new_ir] = rec.state["partition_info"][ir]
     return new_ir, rec.state["partition_info"]
diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py
index f810cc1c7319..215c0cdd3715 100644
--- a/python/cudf_polars/cudf_polars/streaming/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/io.py
@@ -306,8 +306,6 @@ class StreamingScan(IR):
     )
     _n_non_child_args = 2
     scans: list[Scan | SplitScan]
-    # These are essentially the base scan properties, shared by all the scans
-    # schema: Schema
     base_scan: Scan
 
     def __init__(self, scans: list[Scan | SplitScan], base_scan: Scan):
@@ -320,7 +318,6 @@ def __init__(self, scans: list[Scan | SplitScan], base_scan: 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.
-        # TODO: Why do we have it in the first place?
         return (type(self), *tuple(x.get_hashable() for x in self.scans))
 
 

From af40691e87800c1416672c12c1c7d9ec03e296f9 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Wed, 3 Jun 2026 10:30:18 -0700
Subject: [PATCH 08/20] Single-pass lowering

This refactors the refactor to just lower to this StreamingScan inside
the the initial lowering.
---
 python/cudf_polars/cudf_polars/engine/core.py |   9 +-
 .../cudf_polars/streaming/actor_graph/core.py |   3 +-
 .../cudf_polars/streaming/actor_graph/io.py   | 206 +-----------------
 .../cudf_polars/streaming/dispatch.py         |   7 +
 .../cudf_polars/streaming/explain.py          |  13 ++
 .../cudf_polars/cudf_polars/streaming/io.py   | 114 +++++++++-
 .../cudf_polars/streaming/parallel.py         |   9 +
 .../cudf_polars/streaming/select.py           |   6 +-
 .../cudf_polars/tests/streaming/test_scan.py  |  93 +++++++-
 9 files changed, 246 insertions(+), 214 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py
index 4e416d999154..79914d12fe8d 100644
--- a/python/cudf_polars/cudf_polars/engine/core.py
+++ b/python/cudf_polars/cudf_polars/engine/core.py
@@ -28,7 +28,6 @@
 from cudf_polars.streaming.actor_graph.collectives import ReserveOpIDs
 from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id
 from cudf_polars.streaming.actor_graph.core import generate_network
-from cudf_polars.streaming.actor_graph.io import io_lower_ir_graph
 from cudf_polars.streaming.actor_graph.tracing import log_query_plan
 from cudf_polars.streaming.actor_graph.utils import empty_table_chunk
 from cudf_polars.streaming.base import StatsCollector
@@ -686,17 +685,15 @@ def evaluate_on_rank(
         Collected channel metadata.
     """
     stats = allgather_stats(comm, ctx.br(), ir, config_options, py_executor)
-    ir, partition_info = lower_ir_graph(ir, config_options, stats)
+    ir, partition_info = lower_ir_graph(
+        ir, config_options, stats, rank=comm.rank, nranks=comm.nranks
+    )
 
     if comm.rank == 0:
         # At least for now, the query plan is identical on all ranks,
         # so we only log it once.
         log_query_plan(ir, config_options)
 
-    ir, partition_info = io_lower_ir_graph(
-        ir, partition_info, comm, config_options.parquet_options
-    )
-
     with ReserveOpIDs(ir, config_options) as collective_id_map:
         return execute_ir_on_rank(
             ctx,
diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py
index 837384a1561f..d09837be6593 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py
@@ -24,6 +24,7 @@
     generate_ir_sub_network_wrapper,
     metadata_drain_node,
 )
+from cudf_polars.streaming.io import StreamingScan
 from cudf_polars.streaming.over import Over
 from cudf_polars.utils.config import SPMDContext
 
@@ -246,7 +247,7 @@ def generate_network(
     num_io_nodes: int = 0
     ir_dep_count: defaultdict[IR, int] = defaultdict(int)
     for node in traversal([ir]):
-        if isinstance(node, (DataFrameScan, Scan)):
+        if isinstance(node, (DataFrameScan, Scan, StreamingScan)):
             num_io_nodes += 1
         for child in node.children:
             ir_dep_count[child] += 1
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 f9975f1c0a16..a6e7250e0404 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
@@ -7,9 +7,7 @@
 import asyncio
 import dataclasses
 import math
-import operator
-from functools import reduce, singledispatch
-from typing import TYPE_CHECKING, Any, TypeAlias, TypedDict
+from typing import TYPE_CHECKING, Any
 
 from rapidsmpf.memory.memory_reservation import opaque_memory_usage
 from rapidsmpf.streaming.core.memory_reserve_or_wait import (
@@ -24,9 +22,7 @@
 from cudf_polars.dsl.ir import (
     IR,
     DataFrameScan,
-    Scan,
     Sink,
-    Union,
     _prepare_parquet_predicate,
 )
 from cudf_polars.dsl.to_ast import to_parquet_filter
@@ -49,10 +45,8 @@
 )
 from cudf_polars.streaming.base import (
     IOPartitionFlavor,
-    PartitionInfo,
 )
 from cudf_polars.streaming.io import (
-    SplitScan,
     StreamingScan,
     StreamingSink,
     _prepare_sink_directory,
@@ -61,211 +55,19 @@
 from cudf_polars.streaming.utils import _dynamic_planning_on
 
 if TYPE_CHECKING:
-    from collections.abc import MutableMapping
-
     from rapidsmpf.communicator.communicator import Communicator
     from rapidsmpf.streaming.core.channel import Channel
     from rapidsmpf.streaming.core.context import Context
 
-    from cudf_polars.dsl.ir import IR, IRExecutionContext
+    from cudf_polars.dsl.ir import IR, IRExecutionContext, Scan
     from cudf_polars.streaming.actor_graph.core import SubNetGenerator
     from cudf_polars.streaming.actor_graph.tracing import ActorTracer
     from cudf_polars.streaming.base import (
         IOPartitionPlan,
+        PartitionInfo,
         StatsCollector,
     )
-    from cudf_polars.utils.config import (
-        ParquetOptions,
-    )
-
-from cudf_polars.dsl.traversal import CachingVisitor
-from cudf_polars.typing import GenericTransformer
-
-
-class IOLowerIRState(TypedDict):
-    """State used for lowering IR nodes."""
-
-    partition_info: MutableMapping[IR, PartitionInfo]
-    rank: int
-    nranks: int
-    parquet_options: ParquetOptions
-
-
-IOLowerIRTransformer: TypeAlias = GenericTransformer[
-    "IR", "tuple[IR, MutableMapping[IR, PartitionInfo]]", IOLowerIRState
-]
-
-
-def io_lower_ir_graph(
-    ir: IR,
-    partition_info: MutableMapping[IR, PartitionInfo],
-    comm: Communicator,
-    parquet_options: ParquetOptions,
-) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
-    """
-    Rewrite an IR graph and extract partitioning information.
-
-    Parameters
-    ----------
-    ir
-        Root of the graph to rewrite.
-    partition_info
-        Partition information for the graph. Mutate this in place for new IR nodes.
-    comm
-        Communicator for the graph. The output IR graph will contain nodes unique
-        to this rank.
-    parquet_options
-        Parquet reader options.
-
-    Returns
-    -------
-    new_ir, partition_info
-        The rewritten graph and a mapping from unique nodes
-        in the new graph to associated partitioning information.
-
-    Notes
-    -----
-    This function traverses the unique nodes of the graph with
-    root `ir`, and applies :func:`lower_ir_node` to each node.
-
-    See Also
-    --------
-    lower_ir_node
-    """
-    state: IOLowerIRState = {
-        "partition_info": partition_info,
-        "rank": comm.rank,
-        "nranks": comm.nranks,
-        "parquet_options": parquet_options,
-    }
-    mapper: IOLowerIRTransformer = CachingVisitor(io_lower_ir_node, state=state)
-    return mapper(ir)
-
-
-@singledispatch
-def io_lower_ir_node(
-    ir: IR, rec: IOLowerIRTransformer
-) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
-    """Lower the IR nodes for a given IR node."""
-    if not ir.children:
-        return ir, rec.state["partition_info"]
-
-    children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
-    partition_info = reduce(operator.or_, _partition_info)
-
-    if len(children) == 1:
-        partition = partition_info[children[0]]
-    else:
-        partition = PartitionInfo(count=max(partition_info[c].count for c in children))
-
-    new_node = ir.reconstruct(children)
-    partition_info[new_node] = partition
-    return new_node, partition_info
-
-
-@io_lower_ir_node.register(Union)
-def _(
-    ir: Union, rec: IOLowerIRTransformer
-) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
-    # TODO: Determine if we really need this. The only difference is `sum` instead of `max` for getting count.
-    children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
-    partition_info = reduce(operator.or_, _partition_info)
-    count = sum(partition_info[c].count for c in children)
-    new_node = ir.reconstruct(children)
-    partition_info[new_node] = PartitionInfo(count=count)
-    return new_node, partition_info
-
-
-@io_lower_ir_node.register(Scan)
-def _(
-    ir: Scan, rec: IOLowerIRTransformer
-) -> tuple[StreamingScan, MutableMapping[IR, PartitionInfo]]:
-    """
-    Lower the Scan node.
-
-    This expands SplitScan nodes into multiple Scan nodes.
-    """
-    plan = rec.state["partition_info"][ir].io_plan
-    # rec.state["partition_info"] is just a mapping from IR nodes to PartitionInfo.
-    # We promise that PartitionInfo.io_plan is not None for Scan nodes,
-    # but don't currently enforce that promise in the type system.
-    if plan is None:
-        raise RuntimeError(f"Scan node must have a partition plan. node={ir}")
-
-    rank = rec.state["rank"]
-    nranks = rec.state["nranks"]
-    parquet_options = rec.state["parquet_options"]
-
-    scans: list[Scan | SplitScan] = []
-    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:
-            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,
-                        sindex,
-                        plan.factor,
-                        parquet_options,
-                    )
-                )
-                sindex += 1
-                splits_created += 1
-            sindex = 0
-
-    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:  # 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,
-                    )
-                )
-
-    # The PartitionInfo for the derived StreamingScan node is the same as the original Scan node.
-    new_ir = StreamingScan(scans, ir)
-    rec.state["partition_info"][new_ir] = rec.state["partition_info"][ir]
-    return new_ir, rec.state["partition_info"]
+    from cudf_polars.streaming.io import SplitScan
 
 
 class Lineariser:
diff --git a/python/cudf_polars/cudf_polars/streaming/dispatch.py b/python/cudf_polars/cudf_polars/streaming/dispatch.py
index 81937e6ce6fd..3dd612f3de02 100644
--- a/python/cudf_polars/cudf_polars/streaming/dispatch.py
+++ b/python/cudf_polars/cudf_polars/streaming/dispatch.py
@@ -31,10 +31,17 @@ class State(TypedDict):
         GPUEngine configuration options.
     stats
         Statistics collector.
+    rank
+        Rank of the current worker for IO sharding. Always
+        0 for non-streaming engines.
+    nranks
+        Number of workers for IO sharding. Always 1 for non-streaming engines.
     """
 
     config_options: ConfigOptions[StreamingExecutor]
     stats: StatsCollector
+    rank: int
+    nranks: int
 
 
 LowerIRTransformer: TypeAlias = GenericTransformer[
diff --git a/python/cudf_polars/cudf_polars/streaming/explain.py b/python/cudf_polars/cudf_polars/streaming/explain.py
index 92a76fa55209..67cdf311ba16 100644
--- a/python/cudf_polars/cudf_polars/streaming/explain.py
+++ b/python/cudf_polars/cudf_polars/streaming/explain.py
@@ -28,6 +28,7 @@
 )
 from cudf_polars.dsl.translate import Translator
 from cudf_polars.dsl.traversal import traversal
+from cudf_polars.streaming.io import StreamingScan
 from cudf_polars.streaming.parallel import lower_ir_graph
 from cudf_polars.streaming.shuffle import Shuffle
 from cudf_polars.streaming.statistics import (
@@ -288,6 +289,18 @@ def _(ir: Scan) -> dict[str, Serializable]:
     }
 
 
+@_serialize_properties.register
+def _(ir: StreamingScan) -> dict[str, Serializable]:
+    return {
+        "typ": ir.base_scan.typ,
+        "scan_count": len(ir.scans),
+        "prefix": os.path.commonprefix(ir.base_scan.paths),
+        "predicate": (
+            _serialize_expr(ir.base_scan.predicate) if ir.base_scan.predicate else None
+        ),
+    }
+
+
 @_serialize_properties.register
 def _(ir: Join) -> dict[str, Serializable]:
     return {
diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py
index 215c0cdd3715..3e90b346fa94 100644
--- a/python/cudf_polars/cudf_polars/streaming/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/io.py
@@ -103,6 +103,105 @@ def scan_partition_plan(
     return IOPartitionPlan(1, IOPartitionFlavor.SINGLE_FILE)
 
 
+def expand_scan_for_rank(
+    ir: Scan,
+    plan: IOPartitionPlan,
+    *,
+    rank: int,
+    nranks: int,
+    parquet_options: ParquetOptions,
+) -> list[Scan | SplitScan]:
+    """
+    Expand a Scan node into rank-local Scan and SplitScan operations.
+
+    Parameters
+    ----------
+    ir
+        The Scan node to expand.
+    plan
+        The IO partitioning plan for the scan.
+    rank
+        Rank of the current worker.
+    nranks
+        Number of workers. Values less than 1 are treated as 1.
+    parquet_options
+        Parquet reader options.
+
+    Returns
+    -------
+    list[Scan | SplitScan]
+        Rank-local scan operations.
+    """
+    nranks = max(nranks, 1)
+    scans: list[Scan | SplitScan] = []
+    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:
+            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,
+                        sindex,
+                        plan.factor,
+                        parquet_options,
+                    )
+                )
+                sindex += 1
+                splits_created += 1
+            sindex = 0
+
+    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:  # 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,
+                    )
+                )
+
+    return scans
+
+
 class SplitScan(IR):
     """
     Input from a split file.
@@ -265,6 +364,7 @@ def _(
     ir: Scan, rec: LowerIRTransformer
 ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
     config_options = rec.state["config_options"]
+    parquet_options = config_options.parquet_options
     if (
         ir.typ in ("csv", "parquet", "ndjson")
         and ir.n_rows == -1
@@ -282,13 +382,21 @@ def _(
             count = plan.factor * len(paths)
         else:
             count = math.ceil(len(paths) / plan.factor)
-
-        return ir, {ir: PartitionInfo(count=count, io_plan=plan)}
     else:
         plan = IOPartitionPlan(
             flavor=IOPartitionFlavor.SINGLE_READ, factor=len(ir.paths)
         )
-        return ir, {ir: PartitionInfo(count=1, io_plan=plan)}
+        count = 1
+
+    scans = expand_scan_for_rank(
+        ir,
+        plan,
+        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)}
 
 
 class StreamingScan(IR):
diff --git a/python/cudf_polars/cudf_polars/streaming/parallel.py b/python/cudf_polars/cudf_polars/streaming/parallel.py
index 5396787d73d8..7372bec6f74f 100644
--- a/python/cudf_polars/cudf_polars/streaming/parallel.py
+++ b/python/cudf_polars/cudf_polars/streaming/parallel.py
@@ -68,6 +68,9 @@ def lower_ir_graph(
     ir: IR,
     config_options: ConfigOptions[StreamingExecutor],
     stats: StatsCollector,
+    *,
+    rank: int = 0,
+    nranks: int = 1,
 ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
     """
     Rewrite an IR graph and extract partitioning information.
@@ -80,6 +83,10 @@ def lower_ir_graph(
         GPUEngine configuration options.
     stats
         Pre-computed statistics collector.
+    rank
+        Rank of the current worker.
+    nranks
+        Number of workers in the current cluster.
 
     Returns
     -------
@@ -99,6 +106,8 @@ def lower_ir_graph(
     state: State = {
         "config_options": config_options,
         "stats": stats,
+        "rank": rank,
+        "nranks": nranks,
     }
     mapper: LowerIRTransformer = CachingVisitor(lower_ir_node, state=state)
     return mapper(ir)
diff --git a/python/cudf_polars/cudf_polars/streaming/select.py b/python/cudf_polars/cudf_polars/streaming/select.py
index e0a97fdbafbb..e1466f0b6610 100644
--- a/python/cudf_polars/cudf_polars/streaming/select.py
+++ b/python/cudf_polars/cudf_polars/streaming/select.py
@@ -20,6 +20,7 @@
     decompose_expr_graph,
     make_expr_decomposer,
 )
+from cudf_polars.streaming.io import StreamingScan
 from cudf_polars.streaming.over import _fuse_over_nodes
 from cudf_polars.streaming.repartition import Repartition
 from cudf_polars.streaming.utils import (
@@ -421,8 +422,11 @@ def _(
         ):
             # Task engine case
             scan_child = child.children[0]
-        elif isinstance(child, Scan):  # pragma: no cover; Requires rapidsmpf runtime
+        elif isinstance(child, StreamingScan):  # pragma: no cover
             # RapidsMPF case
+            scan_child = child.base_scan
+        elif isinstance(child, Scan):  # pragma: no cover; Requires rapidsmpf runtime
+            # Legacy task-engine case
             scan_child = child
 
     if scan_child and scan_child.predicate is None and scan_child.typ == "parquet":
diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py
index 9d08a2362883..34b3af0761c7 100644
--- a/python/cudf_polars/tests/streaming/test_scan.py
+++ b/python/cudf_polars/tests/streaming/test_scan.py
@@ -10,12 +10,16 @@
 import polars as pl
 
 from cudf_polars import Translator
+from cudf_polars.containers import DataType
+from cudf_polars.dsl.ir import Scan
 from cudf_polars.engine.options import StreamingOptions
+from cudf_polars.streaming.base import IOPartitionFlavor, IOPartitionPlan
+from cudf_polars.streaming.io import SplitScan, 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
 from cudf_polars.testing.io import make_partitioned_source
-from cudf_polars.utils.config import ConfigOptions
+from cudf_polars.utils.config import ConfigOptions, ParquetOptions
 
 if TYPE_CHECKING:
     import concurrent.futures
@@ -152,3 +156,90 @@ def test_scan_union(engine: pl.GPUEngine, tmp_path: Path) -> None:
 
     q = pl.concat([df_q, df_q])
     assert_gpu_result_equal(q, engine=engine)
+
+
+def _make_parquet_scan(paths: list[str]) -> Scan:
+    return Scan(
+        {"x": DataType(pl.Int64())},
+        "parquet",
+        {},
+        None,
+        paths,
+        None,
+        0,
+        -1,
+        None,
+        None,
+        None,
+        ParquetOptions(),
+    )
+
+
+@pytest.mark.parametrize(
+    "plan,paths,rank,nranks,expected_len",
+    [
+        (
+            IOPartitionPlan(2, IOPartitionFlavor.FUSED_FILES),
+            [f"f{i}" for i in range(6)],
+            0,
+            1,
+            3,
+        ),
+        (
+            IOPartitionPlan(2, IOPartitionFlavor.FUSED_FILES),
+            [f"f{i}" for i in range(6)],
+            0,
+            2,
+            2,
+        ),
+        (
+            IOPartitionPlan(2, IOPartitionFlavor.FUSED_FILES),
+            [f"f{i}" for i in range(6)],
+            1,
+            2,
+            1,
+        ),
+        (IOPartitionPlan(3, IOPartitionFlavor.SINGLE_READ), ["a", "b", "c"], 1, 2, 0),
+    ],
+)
+def test_expand_scan_for_rank_fused_and_single_read(
+    plan: IOPartitionPlan,
+    paths: list[str],
+    rank: int,
+    nranks: int,
+    expected_len: int,
+) -> None:
+    scans = expand_scan_for_rank(
+        _make_parquet_scan(paths),
+        plan,
+        rank=rank,
+        nranks=nranks,
+        parquet_options=ParquetOptions(),
+    )
+    assert len(scans) == expected_len
+    assert all(not isinstance(scan, SplitScan) for scan in scans)
+
+
+def test_expand_scan_for_rank_split_files() -> None:
+    plan = IOPartitionPlan(4, IOPartitionFlavor.SPLIT_FILES)
+    scans = expand_scan_for_rank(
+        _make_parquet_scan(["file.parquet"]),
+        plan,
+        rank=0,
+        nranks=2,
+        parquet_options=ParquetOptions(),
+    )
+    assert len(scans) == 2
+    assert all(isinstance(scan, SplitScan) for scan in scans)
+
+
+def test_expand_scan_for_rank_treats_zero_nranks_as_one() -> None:
+    plan = IOPartitionPlan(1, IOPartitionFlavor.FUSED_FILES)
+    scans = expand_scan_for_rank(
+        _make_parquet_scan(["a", "b"]),
+        plan,
+        rank=0,
+        nranks=0,
+        parquet_options=ParquetOptions(),
+    )
+    assert len(scans) == 2

From 3f9d869936a9aa79fd6baa04cef02a7bb3b537bd Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Wed, 3 Jun 2026 10:39:36 -0700
Subject: [PATCH 09/20] fixup

---
 python/cudf_polars/cudf_polars/streaming/io.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py
index 3e90b346fa94..5bcfaea76f7e 100644
--- a/python/cudf_polars/cudf_polars/streaming/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/io.py
@@ -132,7 +132,6 @@ def expand_scan_for_rank(
     list[Scan | SplitScan]
         Rank-local scan operations.
     """
-    nranks = max(nranks, 1)
     scans: list[Scan | SplitScan] = []
     if plan.flavor == IOPartitionFlavor.SPLIT_FILES:
         count = plan.factor * len(ir.paths)

From 4a9358c6d4afb1a5f044910b73d131d9131d0486 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Wed, 3 Jun 2026 13:43:50 -0700
Subject: [PATCH 10/20] Test fixes

---
 python/cudf_polars/tests/streaming/test_explain.py |  3 ++-
 python/cudf_polars/tests/streaming/test_scan.py    | 12 ------------
 2 files changed, 2 insertions(+), 13 deletions(-)

diff --git a/python/cudf_polars/tests/streaming/test_explain.py b/python/cudf_polars/tests/streaming/test_explain.py
index df61f366179a..3bbf9f238d82 100644
--- a/python/cudf_polars/tests/streaming/test_explain.py
+++ b/python/cudf_polars/tests/streaming/test_explain.py
@@ -521,6 +521,7 @@ def test_scan_properties(tmp_path: Path, predicate: pl.Expr | None):
         "prefix": f"{root}/",
         "typ": "parquet",
         "predicate": None,
+        "scan_count": 1,
     }
     if predicate is not None:
         q = q.filter(predicate)
@@ -537,7 +538,7 @@ def test_scan_properties(tmp_path: Path, predicate: pl.Expr | None):
     dag = serialize_query(q, engine)
 
     node = dag.nodes[dag.roots[0]]
-    assert node.type == "Scan"
+    assert node.type == "StreamingScan"
     assert node.properties == expected_properties
 
 
diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py
index 34b3af0761c7..88d3a40f7593 100644
--- a/python/cudf_polars/tests/streaming/test_scan.py
+++ b/python/cudf_polars/tests/streaming/test_scan.py
@@ -231,15 +231,3 @@ def test_expand_scan_for_rank_split_files() -> None:
     )
     assert len(scans) == 2
     assert all(isinstance(scan, SplitScan) for scan in scans)
-
-
-def test_expand_scan_for_rank_treats_zero_nranks_as_one() -> None:
-    plan = IOPartitionPlan(1, IOPartitionFlavor.FUSED_FILES)
-    scans = expand_scan_for_rank(
-        _make_parquet_scan(["a", "b"]),
-        plan,
-        rank=0,
-        nranks=0,
-        parquet_options=ParquetOptions(),
-    )
-    assert len(scans) == 2

From d39c724d972e8aeb690321e95e130a6fada3d5ca Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Wed, 3 Jun 2026 15:02:34 -0700
Subject: [PATCH 11/20] cleanup

---
 python/cudf_polars/cudf_polars/streaming/actor_graph/io.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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 a6e7250e0404..8aad70bbabf5 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
@@ -592,7 +592,7 @@ def _(
         native_node = make_rapidsmpf_read_parquet_node(
             rec.state["context"],
             rec.state["comm"],
-            ir.base_scan,  # TODO: verify this...
+            ir.base_scan,
             num_producers,
             ch_in,
             rec.state["stats"],

From b8ca40e5f0b57d7e1c690814019413ab3fa66273 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Wed, 3 Jun 2026 15:25:14 -0700
Subject: [PATCH 12/20] docs

---
 .../cudf_polars/streaming/actor_graph/io.py   | 32 +++-----
 .../cudf_polars/cudf_polars/streaming/io.py   | 79 +++++++++++++++++++
 2 files changed, 88 insertions(+), 23 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 8aad70bbabf5..2fbe609abfde 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
@@ -5,7 +5,6 @@
 from __future__ import annotations
 
 import asyncio
-import dataclasses
 import math
 from typing import TYPE_CHECKING, Any
 
@@ -43,16 +42,13 @@
     recv_metadata,
     send_metadata,
 )
-from cudf_polars.streaming.base import (
-    IOPartitionFlavor,
-)
 from cudf_polars.streaming.io import (
     StreamingScan,
     StreamingSink,
     _prepare_sink_directory,
     _sink_to_file,
+    determine_non_native_fallback,
 )
-from cudf_polars.streaming.utils import _dynamic_planning_on
 
 if TYPE_CHECKING:
     from rapidsmpf.communicator.communicator import Communicator
@@ -567,25 +563,19 @@ def _(
     assert partition_info.io_plan is not None, "Scan node must have a partition plan"
     plan: IOPartitionPlan = partition_info.io_plan
 
-    # Native node cannot split large files in distributed mode yet
-    distributed_split_files = (
-        plan.flavor == IOPartitionFlavor.SPLIT_FILES and rec.state["comm"].nranks > 1
-    )
-
     # Use rapidsmpf native read_parquet node if possible
     ch_in: Channel[TableChunk] | None = None
     ch_out = channels[ir].reserve_input_slot()
     nodes: dict[IR, list[Any]] = {}
     native_node: Any = None
-    if (
-        parquet_options.use_rapidsmpf_native
-        and (partition_info.count > 1 or _dynamic_planning_on(config_options))
-        and ir.base_scan.typ == "parquet"
-        and ir.base_scan.row_index is None
-        and ir.base_scan.include_file_paths is None
-        and ir.base_scan.n_rows == -1
-        and ir.base_scan.skip_rows == 0
-        and not distributed_split_files
+
+    if determine_non_native_fallback(
+        ir.base_scan,
+        plan=plan,
+        count=partition_info.count,
+        nranks=rec.state["comm"].nranks,
+        parquet_options=parquet_options,
+        config_options=config_options,
     ):
         # Create new channel to so ch_out can be used to add metadata
         ch_in = rec.state["context"].create_channel()
@@ -599,7 +589,6 @@ def _(
             partition_info,
         )
 
-    if native_node is not None and ch_in is not None:
         # Need metadata node, because the native read_parquet
         # node does not send metadata.
         metadata_node = metadata_feeder_node(
@@ -616,9 +605,6 @@ def _(
         )
         nodes[ir] = [native_node, metadata_node]
     else:
-        # Fall back to scan_node (predicate not convertible, or other constraint)
-        parquet_options = dataclasses.replace(parquet_options, chunked=False)
-
         nodes[ir] = [
             scan_node(
                 rec.state["context"],
diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py
index 5bcfaea76f7e..e38e6aa36edb 100644
--- a/python/cudf_polars/cudf_polars/streaming/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/io.py
@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import dataclasses
 import functools
 import itertools
 import math
@@ -31,6 +32,7 @@
     SerializedDataSourceInfo,
 )
 from cudf_polars.streaming.dispatch import lower_ir_node
+from cudf_polars.streaming.utils import _dynamic_planning_on
 from cudf_polars.utils.config import Cluster
 from cudf_polars.utils.cuda_stream import get_cuda_stream
 from cudf_polars.utils.versions import POLARS_VERSION_LT_137
@@ -358,6 +360,67 @@ def _(
     return ir, {ir: PartitionInfo(count=1)}  # pragma: no cover
 
 
+def determine_non_native_fallback(
+    ir: Scan,
+    *,
+    plan: IOPartitionPlan,
+    count: int,
+    nranks: int,
+    parquet_options: ParquetOptions,
+    config_options: ConfigOptions[StreamingExecutor],
+) -> bool:
+    """
+    Determine whether we will use the cudf-polars (non-native) parquet reader.
+
+    Parameters
+    ----------
+    ir
+        The Scan node that might need to fall back.
+    plan
+        The IO partitioning plan.
+    count
+        The number of partitions associated with this Scan node.
+    nranks
+        The number of ranks.
+    parquet_options
+        The parquet options.
+    config_options
+        The configuration options.
+
+    Returns
+    -------
+    bool
+        Whether to use the cudf-polars (non-native) parquet reader.
+
+    Notes
+    -----
+    cudf-polars current falls back under the following conditions:
+
+    - Our plan indicates we should split the file into multiple partitions
+    - We have more than one rank
+    - There's more than one partition or dynamic planning is enabled
+    - The file type is parquet
+    - The row index is not set
+    - File paths are not included
+    - The number of rows is not set
+    - The skip rows is not set
+    """
+    distributed_split_files = (
+        plan.flavor == IOPartitionFlavor.SPLIT_FILES and nranks > 1
+    )
+
+    return not (
+        parquet_options.use_rapidsmpf_native
+        and (count > 1 or _dynamic_planning_on(config_options))
+        and ir.typ == "parquet"
+        and ir.row_index is None
+        and ir.include_file_paths is None
+        and ir.n_rows == -1
+        and ir.skip_rows == 0
+        and not distributed_split_files
+    )
+
+
 @lower_ir_node.register(Scan)
 def _(
     ir: Scan, rec: LowerIRTransformer
@@ -387,6 +450,22 @@ def _(
         )
         count = 1
 
+    # In `generate_ir_sub_network` for `StreamingScan`, we have this big condition
+    # for whether we're going to actually use rapidsmpf's native parquet reader
+    # or fall back to a Scan node. When we do fall back, we use the non-chunked
+    # reader.
+    # This
+
+    if determine_non_native_fallback(
+        ir,
+        plan=plan,
+        count=count,
+        nranks=rec.state["nranks"],
+        parquet_options=parquet_options,
+        config_options=config_options,
+    ):
+        parquet_options = dataclasses.replace(parquet_options, chunked=False)
+
     scans = expand_scan_for_rank(
         ir,
         plan,

From b6fc80ca75337eae647413bd049d7c725be1a4b1 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Wed, 3 Jun 2026 19:33:16 -0700
Subject: [PATCH 13/20] Fixed fallback name, usage

---
 .../cudf_polars/streaming/actor_graph/io.py    |  7 ++++---
 python/cudf_polars/cudf_polars/streaming/io.py | 18 ++++++------------
 2 files changed, 10 insertions(+), 15 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 2fbe609abfde..edb3e049b646 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
@@ -47,7 +47,7 @@
     StreamingSink,
     _prepare_sink_directory,
     _sink_to_file,
-    determine_non_native_fallback,
+    should_use_native_parquet_node,
 )
 
 if TYPE_CHECKING:
@@ -569,14 +569,15 @@ def _(
     nodes: dict[IR, list[Any]] = {}
     native_node: Any = None
 
-    if determine_non_native_fallback(
+    use_native = should_use_native_parquet_node(
         ir.base_scan,
         plan=plan,
         count=partition_info.count,
         nranks=rec.state["comm"].nranks,
         parquet_options=parquet_options,
         config_options=config_options,
-    ):
+    )
+    if use_native:
         # Create new channel to so ch_out can be used to add metadata
         ch_in = rec.state["context"].create_channel()
         native_node = make_rapidsmpf_read_parquet_node(
diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py
index e38e6aa36edb..d7f0d2190eca 100644
--- a/python/cudf_polars/cudf_polars/streaming/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/io.py
@@ -360,7 +360,7 @@ def _(
     return ir, {ir: PartitionInfo(count=1)}  # pragma: no cover
 
 
-def determine_non_native_fallback(
+def should_use_native_parquet_node(
     ir: Scan,
     *,
     plan: IOPartitionPlan,
@@ -370,7 +370,7 @@ def determine_non_native_fallback(
     config_options: ConfigOptions[StreamingExecutor],
 ) -> bool:
     """
-    Determine whether we will use the cudf-polars (non-native) parquet reader.
+    Determine whether we should use rapidsmpf's native parquet node.
 
     Parameters
     ----------
@@ -390,11 +390,11 @@ def determine_non_native_fallback(
     Returns
     -------
     bool
-        Whether to use the cudf-polars (non-native) parquet reader.
+        Whether to use rapidsmpf's native parquet node.
 
     Notes
     -----
-    cudf-polars current falls back under the following conditions:
+    Native parquet node is used under the following conditions:
 
     - Our plan indicates we should split the file into multiple partitions
     - We have more than one rank
@@ -409,7 +409,7 @@ def determine_non_native_fallback(
         plan.flavor == IOPartitionFlavor.SPLIT_FILES and nranks > 1
     )
 
-    return not (
+    return (
         parquet_options.use_rapidsmpf_native
         and (count > 1 or _dynamic_planning_on(config_options))
         and ir.typ == "parquet"
@@ -450,13 +450,7 @@ def _(
         )
         count = 1
 
-    # In `generate_ir_sub_network` for `StreamingScan`, we have this big condition
-    # for whether we're going to actually use rapidsmpf's native parquet reader
-    # or fall back to a Scan node. When we do fall back, we use the non-chunked
-    # reader.
-    # This
-
-    if determine_non_native_fallback(
+    if not should_use_native_parquet_node(
         ir,
         plan=plan,
         count=count,

From 83c19c057b514104b640b3e176fd350a9cd41ac3 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Thu, 4 Jun 2026 04:05:11 -0700
Subject: [PATCH 14/20] fixup

---
 python/cudf_polars/cudf_polars/streaming/actor_graph/core.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py
index d09837be6593..fb7cc4e1439a 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py
@@ -15,7 +15,6 @@
 from cudf_polars.dsl.ir import (
     DataFrameScan,
     Join,
-    Scan,
     Union,
 )
 from cudf_polars.dsl.traversal import CachingVisitor, traversal
@@ -247,7 +246,7 @@ def generate_network(
     num_io_nodes: int = 0
     ir_dep_count: defaultdict[IR, int] = defaultdict(int)
     for node in traversal([ir]):
-        if isinstance(node, (DataFrameScan, Scan, StreamingScan)):
+        if isinstance(node, (DataFrameScan, StreamingScan)):
             num_io_nodes += 1
         for child in node.children:
             ir_dep_count[child] += 1

From 889a19a4e5207245a7a640c215a49e3784723ddb Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Thu, 4 Jun 2026 04:36:43 -0700
Subject: [PATCH 15/20] Remove schema from non_child

---
 python/cudf_polars/cudf_polars/streaming/io.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py
index d7f0d2190eca..12cb6fd13a3e 100644
--- a/python/cudf_polars/cudf_polars/streaming/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/io.py
@@ -482,7 +482,6 @@ class StreamingScan(IR):
     _non_child = (
         "scans",
         "base_scan",
-        "schema",
     )
     _n_non_child_args = 2
     scans: list[Scan | SplitScan]

From 84599b2ad06fad4bef9475564ae9422562d7a48b Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Thu, 4 Jun 2026 05:55:25 -0700
Subject: [PATCH 16/20] StreamingScan.do_evaluate raises

---
 python/cudf_polars/cudf_polars/streaming/io.py  | 13 +++++++++++++
 python/cudf_polars/tests/streaming/test_scan.py | 12 ++++++++++--
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py
index 12cb6fd13a3e..d61ccc4acaf6 100644
--- a/python/cudf_polars/cudf_polars/streaming/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/io.py
@@ -499,6 +499,19 @@ def get_hashable(self) -> Hashable:
         # 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))
 
+    @classmethod
+    def do_evaluate(
+        cls,
+        scans: list[Scan | SplitScan],
+        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."
+        )
+
 
 class StreamingSink(IR):
     """Sink a dataframe in streaming mode."""
diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py
index 88d3a40f7593..cf7ff3fb50ed 100644
--- a/python/cudf_polars/tests/streaming/test_scan.py
+++ b/python/cudf_polars/tests/streaming/test_scan.py
@@ -11,10 +11,10 @@
 
 from cudf_polars import Translator
 from cudf_polars.containers import DataType
-from cudf_polars.dsl.ir import Scan
+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, expand_scan_for_rank
+from cudf_polars.streaming.io import 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
@@ -231,3 +231,11 @@ def test_expand_scan_for_rank_split_files() -> None:
     )
     assert len(scans) == 2
     assert all(isinstance(scan, SplitScan) for scan in scans)
+
+
+def test_streaming_scan_raises():
+    # This isn't reachable by normal cudf-polars usage.
+    scan = _make_parquet_scan(["file.parquet"])
+    ctx = IRExecutionContext()
+    with pytest.raises(NotImplementedError):
+        StreamingScan.do_evaluate([scan], scan, context=ctx)

From 51fa4bd402488b7a4ef56f4f4f39e064620ee462 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Thu, 4 Jun 2026 06:02:00 -0700
Subject: [PATCH 17/20] expand scan assertions

---
 .../cudf_polars/tests/streaming/test_scan.py  | 39 +++++++++++++------
 1 file changed, 27 insertions(+), 12 deletions(-)

diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py
index cf7ff3fb50ed..58beee9aa5aa 100644
--- a/python/cudf_polars/tests/streaming/test_scan.py
+++ b/python/cudf_polars/tests/streaming/test_scan.py
@@ -176,30 +176,30 @@ def _make_parquet_scan(paths: list[str]) -> Scan:
 
 
 @pytest.mark.parametrize(
-    "plan,paths,rank,nranks,expected_len",
+    "plan,paths,rank,nranks,expected_path_groups",
     [
         (
             IOPartitionPlan(2, IOPartitionFlavor.FUSED_FILES),
             [f"f{i}" for i in range(6)],
             0,
             1,
-            3,
+            [["f0", "f1"], ["f2", "f3"], ["f4", "f5"]],
         ),
         (
             IOPartitionPlan(2, IOPartitionFlavor.FUSED_FILES),
             [f"f{i}" for i in range(6)],
             0,
             2,
-            2,
+            [["f0", "f1"], ["f2", "f3"]],
         ),
         (
             IOPartitionPlan(2, IOPartitionFlavor.FUSED_FILES),
             [f"f{i}" for i in range(6)],
             1,
             2,
-            1,
+            [["f4", "f5"]],
         ),
-        (IOPartitionPlan(3, IOPartitionFlavor.SINGLE_READ), ["a", "b", "c"], 1, 2, 0),
+        (IOPartitionPlan(3, IOPartitionFlavor.SINGLE_READ), ["a", "b", "c"], 1, 2, []),
     ],
 )
 def test_expand_scan_for_rank_fused_and_single_read(
@@ -207,7 +207,7 @@ def test_expand_scan_for_rank_fused_and_single_read(
     paths: list[str],
     rank: int,
     nranks: int,
-    expected_len: int,
+    expected_path_groups: list[list[str]],
 ) -> None:
     scans = expand_scan_for_rank(
         _make_parquet_scan(paths),
@@ -216,21 +216,36 @@ def test_expand_scan_for_rank_fused_and_single_read(
         nranks=nranks,
         parquet_options=ParquetOptions(),
     )
-    assert len(scans) == expected_len
-    assert all(not isinstance(scan, SplitScan) for scan in scans)
+    for scan, expected_paths in zip(scans, expected_path_groups, strict=True):
+        assert isinstance(scan, Scan)
+        assert scan.paths == expected_paths
 
 
-def test_expand_scan_for_rank_split_files() -> None:
+@pytest.mark.parametrize(
+    "rank,expected_splits",
+    [
+        (0, [(0, 4), (1, 4)]),
+        (1, [(2, 4), (3, 4)]),
+    ],
+)
+def test_expand_scan_for_rank_split_files(
+    rank: int,
+    expected_splits: list[tuple[int, int]],
+) -> None:
     plan = IOPartitionPlan(4, IOPartitionFlavor.SPLIT_FILES)
     scans = expand_scan_for_rank(
         _make_parquet_scan(["file.parquet"]),
         plan,
-        rank=0,
+        rank=rank,
         nranks=2,
         parquet_options=ParquetOptions(),
     )
-    assert len(scans) == 2
-    assert all(isinstance(scan, SplitScan) for scan in scans)
+    assert len(scans) == len(expected_splits)
+    for scan, (split_index, total_splits) in zip(scans, expected_splits, strict=True):
+        assert isinstance(scan, SplitScan)
+        assert scan.split_index == split_index
+        assert scan.total_splits == total_splits
+        assert scan.base_scan.paths == ["file.parquet"]
 
 
 def test_streaming_scan_raises():

From b7b00e1bf69e3952a505fef74c2c8974bde18f28 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Thu, 4 Jun 2026 08:08:42 -0700
Subject: [PATCH 18/20] expand scan assertions

---
 python/cudf_polars/tests/streaming/test_scan.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py
index 58beee9aa5aa..c9ccb13202dc 100644
--- a/python/cudf_polars/tests/streaming/test_scan.py
+++ b/python/cudf_polars/tests/streaming/test_scan.py
@@ -248,9 +248,9 @@ def test_expand_scan_for_rank_split_files(
         assert scan.base_scan.paths == ["file.parquet"]
 
 
-def test_streaming_scan_raises():
+def test_streaming_scan_raises() -> None:
     # This isn't reachable by normal cudf-polars usage.
     scan = _make_parquet_scan(["file.parquet"])
     ctx = IRExecutionContext()
-    with pytest.raises(NotImplementedError):
+    with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"):
         StreamingScan.do_evaluate([scan], scan, context=ctx)

From 5a5e68019e3fec57797533f643cdc8fdb18b9933 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Thu, 4 Jun 2026 08:15:18 -0700
Subject: [PATCH 19/20] can -> should

---
 python/cudf_polars/cudf_polars/streaming/actor_graph/io.py | 4 ++--
 python/cudf_polars/cudf_polars/streaming/io.py             | 4 ++--
 2 files changed, 4 insertions(+), 4 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 edb3e049b646..7f4db5e38bb3 100644
--- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
@@ -47,7 +47,7 @@
     StreamingSink,
     _prepare_sink_directory,
     _sink_to_file,
-    should_use_native_parquet_node,
+    can_use_native_parquet_node,
 )
 
 if TYPE_CHECKING:
@@ -569,7 +569,7 @@ def _(
     nodes: dict[IR, list[Any]] = {}
     native_node: Any = None
 
-    use_native = should_use_native_parquet_node(
+    use_native = can_use_native_parquet_node(
         ir.base_scan,
         plan=plan,
         count=partition_info.count,
diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py
index d61ccc4acaf6..00b4de8944e3 100644
--- a/python/cudf_polars/cudf_polars/streaming/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/io.py
@@ -360,7 +360,7 @@ def _(
     return ir, {ir: PartitionInfo(count=1)}  # pragma: no cover
 
 
-def should_use_native_parquet_node(
+def can_use_native_parquet_node(
     ir: Scan,
     *,
     plan: IOPartitionPlan,
@@ -450,7 +450,7 @@ def _(
         )
         count = 1
 
-    if not should_use_native_parquet_node(
+    if not can_use_native_parquet_node(
         ir,
         plan=plan,
         count=count,

From 3a84c646f19f0a8fc182e9c13c61eed2bc2eec90 Mon Sep 17 00:00:00 2001
From: Tom Augspurger <toaugspurger@nvidia.com>
Date: Thu, 4 Jun 2026 08:20:13 -0700
Subject: [PATCH 20/20] Remove stale comment

---
 python/cudf_polars/cudf_polars/streaming/io.py |  2 +-
 python/cudf_polars/tests/test_scan.py          | 10 ++++------
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py
index 00b4de8944e3..d56c31ce379b 100644
--- a/python/cudf_polars/cudf_polars/streaming/io.py
+++ b/python/cudf_polars/cudf_polars/streaming/io.py
@@ -125,7 +125,7 @@ def expand_scan_for_rank(
     rank
         Rank of the current worker.
     nranks
-        Number of workers. Values less than 1 are treated as 1.
+        Number of workers.
     parquet_options
         Parquet reader options.
 
diff --git a/python/cudf_polars/tests/test_scan.py b/python/cudf_polars/tests/test_scan.py
index 3880d8909119..7a9a4f2bb10f 100644
--- a/python/cudf_polars/tests/test_scan.py
+++ b/python/cudf_polars/tests/test_scan.py
@@ -41,10 +41,10 @@
 
 
 @pytest.fixture(
-    params=[(None, None), ("row-index", 0), ("index", 10)],
+    params=[(None, 0), ("row-index", 0), ("index", 10)],
     ids=["no_row_index", "zero_offset_row_index", "offset_row_index"],
 )
-def row_index(request) -> tuple[str | None, int | None]:
+def row_index(request) -> tuple[str | None, int]:
     return request.param
 
 
@@ -480,7 +480,7 @@ def test_scan_csv_with_and_without_header(
     *,
     has_header: bool,
     new_columns: list[str] | None,
-    row_index: tuple[str | None, int | None],
+    row_index: tuple[str | None, int],
     columns: list[str] | None,
     zlice: tuple[int, int] | None,
 ):
@@ -491,14 +491,12 @@ def test_scan_csv_with_and_without_header(
 
     name, offset = row_index
 
-    # offset is only used if name is used, but that's not an overload
-    # in the type signature.
     q = pl.scan_csv(
         path,
         has_header=has_header,
         new_columns=new_columns,
         row_index_name=name,
-        row_index_offset=offset,  # type: ignore[arg-type]
+        row_index_offset=offset,
     )
 
     if zlice is not None: