Add ParquetScanTask for all parquet scan tasks

python/cudf_polars/cudf_polars/dsl/utils/io.py

@@ -15,7 +15,11 @@
 
 from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars
 from cudf_polars.dsl.traversal import traversal
-from cudf_polars.streaming.io import Scan, StreamingScan
+from cudf_polars.streaming.io import (
+    ParquetSourceInfo,
+    Scan,
+    StreamingScan,
+)
 
 if TYPE_CHECKING:
     from cudf_polars.dsl.ir import IR
@@ -181,14 +185,12 @@ def prefetch_parquet_file_metadata_for_ir(
     -------
     A dictionary mapping each individual path to its cached parquet metadata.
     """
-    from cudf_polars.streaming.io import ParquetSourceInfo, StreamingScan
-
     all_paths: set[str] = set()
 
     for node in traversal([root]):
         if isinstance(node, StreamingScan) and node.base_scan.typ == "parquet":
-            for scan in node.scans:
-                for path in scan.paths:
+            for task in node.tasks:
+                for path in task.paths:
                     all_paths.add(path)
         elif isinstance(node, Scan) and node.typ == "parquet":  # pragma: no cover
             raise RuntimeError("Unexpected parquet 'Scan' node in lowered IR graph.")
@@ -241,7 +243,7 @@ def attach_cached_parquet_metadata(
     cached_parquet_info_map: dict[str, CachedParquetInfo],
 ) -> None:
     """
-    Attach prefetched metadata to scan nodes.
+    Attach prefetched metadata to parquet scan tasks.
 
     This is an optimization only and does not affect IR identity.
 
@@ -254,10 +256,15 @@ def attach_cached_parquet_metadata(
     """
     for node in traversal([root]):
         if isinstance(node, StreamingScan) and node.base_scan.typ == "parquet":
-            for scan in node.scans:
-                if not all(path in cached_parquet_info_map for path in scan.paths):
-                    continue
-                cached = [cached_parquet_info_map[path] for path in scan.paths]
-                Scan._validate_cached_parquet_info(scan.paths, cached)
-                scan.cached_parquet_info = cached
-                scan._non_child_args = (*scan._non_child_args[:-1], cached)
+            base_scan = node.base_scan
+            task_paths = {path for task in node.tasks for path in task.paths}
+            cached_paths = [
+                path
+                for path in base_scan.paths
+                if path in task_paths and path in cached_parquet_info_map
+            ]
+            cached = [cached_parquet_info_map[path] for path in cached_paths]
+            if not cached:
+                continue
+            Scan._validate_cached_parquet_info(cached_paths, cached)
+            base_scan.cached_parquet_info = cached

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

@@ -39,6 +39,7 @@
     send_metadata,
 )
 from cudf_polars.streaming.io import (
+    ScanTask,
     StreamingScan,
     StreamingSink,
     _prepare_sink_directory,
@@ -60,7 +61,6 @@
         IOPartitionPlan,
         PartitionInfo,
     )
-    from cudf_polars.streaming.io import FusedScan, SplitScan
     from cudf_polars.utils.config import MaxConcurrentIOTasks
 
 
@@ -516,7 +516,7 @@ def _(
 
 async def read_chunk(
     context: Context,
-    scan: IR,
+    task: IR,
     seq_num: int,
     ch_out: Channel[TableChunk],
     ir_context: IRExecutionContext,
@@ -530,8 +530,8 @@ async def read_chunk(
     ----------
     context
         The rapidsmpf context.
-    scan
-        The Scan or DataFrameScan node.
+    task
+        The scan task to evaluate.
     seq_num
         The sequence number.
     ch_out
@@ -546,7 +546,7 @@ async def read_chunk(
     """
     reservation_bytes = (
         estimated_chunk_bytes
-        if isinstance(scan, DataFrameScan)
+        if isinstance(task, DataFrameScan)
         else 2 * estimated_chunk_bytes
     )
     start = time.monotonic_ns()
@@ -558,8 +558,8 @@ async def read_chunk(
     admitted = time.monotonic_ns()
     with opaque_memory_usage(reservation):
         df = await ir_context.to_thread(
-            scan.do_evaluate,
-            *scan._non_child_args,
+            task.do_evaluate,
+            *task._non_child_args,
             context=ir_context,
         )
         chunk = TableChunk.from_pylibcudf_table(
@@ -569,14 +569,17 @@ async def read_chunk(
             br=context.br(),
         )
     stop = time.monotonic_ns()
+    ir_type = (
+        task.trace_ir_type() if isinstance(task, ScanTask) else type(task).__name__
+    )
     log(
         "IO Task",
         scope=Scope.IO_TASK.value,
         start=start,
         admitted=admitted,
         stop=stop,
-        ir_id=scan.get_stable_id(),
-        ir_type=type(scan).__name__,
+        ir_id=task.get_stable_id(),
+        ir_type=ir_type,
         sequence_number=seq_num,
         estimated_output_bytes=estimated_chunk_bytes,
         reservation_bytes=reservation_bytes,
@@ -613,7 +616,7 @@ async def scan_node(
         Estimated retained output size of each chunk in bytes. Used to estimate
         peak memory for admission before launching each read.
     """
-    scans: Sequence[SplitScan] | Sequence[FusedScan] = ir.scans
+    tasks: Sequence[ScanTask] = ir.tasks
 
     async with shutdown_on_error(
         context, ch_out, trace_ir=ir, ir_context=ir_context
@@ -622,21 +625,21 @@ async def scan_node(
         await send_metadata(
             ch_out,
             context,
-            ChannelMetadata(local_count=len(scans)),
+            ChannelMetadata(local_count=len(tasks)),
         )
 
         # If there is nothing to scan, drain the channel and return
-        if len(scans) == 0:
+        if len(tasks) == 0:
             await ch_out.drain(context)
             return
 
-        # If there is only one scan or one producer, we can
+        # If there is only one task or one producer, we can
         # skip the lineariser and read the chunks directly
-        if len(scans) == 1 or num_producers == 1:
-            for seq_num, scan in enumerate(scans):
+        if len(tasks) == 1 or num_producers == 1:
+            for seq_num, task in enumerate(tasks):
                 await read_chunk(
                     context,
-                    scan,
+                    task,
                     seq_num,
                     ch_out,
                     ir_context,
@@ -647,23 +650,22 @@ async def scan_node(
             return
 
         # Use Lineariser to ensure ordered delivery
-        num_producers = min(num_producers, len(scans))
+        num_producers = min(num_producers, len(tasks))
         lineariser = Lineariser(context, ch_out, num_producers)
 
         # Assign tasks to producers using round-robin
-        producer_tasks: list[list[tuple[int, SplitScan | FusedScan]]] = [
+        producer_tasks: list[list[tuple[int, ScanTask]]] = [
             [] for _ in range(num_producers)
         ]
-        for task_idx, scan in enumerate(scans):
+        for task_idx, task in enumerate(tasks):
             producer_id = task_idx % num_producers
-            # mypy resolves __iter__ on union-of-sequences to the common base (IR)
-            producer_tasks[producer_id].append((task_idx, scan))  # type: ignore[arg-type]
+            producer_tasks[producer_id].append((task_idx, task))
 
         async def _producer(producer_id: int, ch_out: Channel) -> None:
-            for task_idx, scan in producer_tasks[producer_id]:
+            for task_idx, task in producer_tasks[producer_id]:
                 await read_chunk(
                     context,
-                    scan,
+                    task,
                     task_idx,
                     ch_out,
                     ir_context,

python/cudf_polars/cudf_polars/streaming/explain.py

@@ -563,7 +563,7 @@ def _(ir: Scan) -> dict[str, Serializable]:
 def _(ir: StreamingScan) -> dict[str, Serializable]:
     return {
         "typ": ir.base_scan.typ,
-        "scan_count": len(ir.scans),
+        "task_count": len(ir.tasks),
         "prefix": os.path.commonprefix(ir.base_scan.paths),
         "predicate": (
             _serialize_expr(ir.base_scan.predicate) if ir.base_scan.predicate else None