Refactor dynamic Scan node lowering

python/cudf_polars/cudf_polars/engine/core.py

@@ -685,7 +685,9 @@ 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,

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
@@ -24,6 +23,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 +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)):
+        if isinstance(node, (DataFrameScan, StreamingScan)):
             num_io_nodes += 1
         for child in node.children:
             ir_dep_count[child] += 1

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
 
@@ -22,7 +21,6 @@
 from cudf_polars.dsl.ir import (
     IR,
     DataFrameScan,
-    Scan,
     Sink,
     _prepare_parquet_predicate,
 )
@@ -44,31 +42,28 @@
     recv_metadata,
     send_metadata,
 )
-from cudf_polars.streaming.base import (
-    IOPartitionFlavor,
-)
 from cudf_polars.streaming.io import (
-    SplitScan,
+    StreamingScan,
     StreamingSink,
     _prepare_sink_directory,
     _sink_to_file,
+    can_use_native_parquet_node,
 )
-from cudf_polars.streaming.utils import _dynamic_planning_on
 
 if TYPE_CHECKING:
     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.streaming.io import SplitScan
 
 
 class Lineariser:
@@ -356,14 +351,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 +365,6 @@ async def scan_node(
     ----------
     context
         The rapidsmpf context.
-    comm
-        The communicator.
     ir
         The Scan node.
     ir_context
@@ -383,84 +373,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,9 +549,9 @@ def make_rapidsmpf_read_parquet_node(
         ) from e
 
 
-@generate_ir_sub_network.register(Scan)
+@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
@@ -642,39 +563,33 @@ 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.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
-    ):
+
+    use_native = can_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(
             rec.state["context"],
             rec.state["comm"],
-            ir,
+            ir.base_scan,
             num_producers,
             ch_in,
             rec.state["stats"],
             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(
@@ -691,19 +606,13 @@ 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"],
-                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,
             )
         ]

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[

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 {