Add IR-fusion optimization pass to streaming cudf-polars executor

python/cudf_polars/cudf_polars/experimental/explain.py

@@ -17,6 +17,7 @@
     Sort,
 )
 from cudf_polars.dsl.translate import Translator
+from cudf_polars.experimental.fusion import Fused
 from cudf_polars.experimental.parallel import lower_ir_graph
 from cudf_polars.utils.config import ConfigOptions
 
@@ -30,7 +31,10 @@
 
 
 def explain_query(
-    q: pl.LazyFrame, engine: pl.GPUEngine, *, physical: bool = True
+    q: pl.LazyFrame,
+    engine: pl.GPUEngine,
+    *,
+    physical: bool = True,
 ) -> str:
     """
     Return a formatted string representation of the IR plan.
@@ -125,3 +129,11 @@ def _(ir: Sort, *, offset: str = "") -> str:
 def _(ir: Scan, *, offset: str = "") -> str:
     label = f"SCAN {ir.typ.upper()}"
     return _repr_header(offset, label, ir.schema)
+
+
+@_repr_ir.register
+def _(ir: Fused, *, offset: str = "") -> str:
+    nodes = tuple(type(node).__name__.upper() for node in ir.subnodes)
+    if ir.fused_io:
+        nodes = (type(ir.fused_io.children[0]).__name__.upper(), *nodes)
+    return _repr_header(offset, f"FUSED {nodes}", ir.schema)

python/cudf_polars/cudf_polars/experimental/fusion.py

@@ -0,0 +1,190 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""IR node fusion."""
+
+from __future__ import annotations
+
+import operator
+from collections import defaultdict
+from functools import reduce
+from typing import TYPE_CHECKING, Any
+
+from cudf_polars.dsl.ir import (
+    IR,
+    Cache,
+    Filter,
+    GroupBy,
+    HStack,
+    MapFunction,
+    Projection,
+    Select,
+    Sort,
+    Union,
+)
+from cudf_polars.dsl.traversal import CachingVisitor, traversal
+from cudf_polars.experimental.dispatch import generate_ir_tasks
+from cudf_polars.experimental.io import Scan, SplitScan
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, MutableMapping, Sequence
+    from typing import Any
+
+    from cudf_polars.containers import DataFrame
+    from cudf_polars.experimental.base import PartitionInfo
+    from cudf_polars.experimental.dispatch import LowerIRTransformer
+    from cudf_polars.typing import Schema
+    from cudf_polars.utils.config import ConfigOptions
+
+
+class Fused(IR):
+    """Fused IR node."""
+
+    __slots__ = ("fused_io", "subnodes")
+    _non_child = ("schema", "subnodes", "fused_io")
+    subnodes: tuple[IR, ...]
+    """Fused sub-nodes."""
+    fused_io: Union | None
+    """Fused IO node."""
+
+    def __init__(
+        self,
+        schema: Schema,
+        subnodes: tuple[IR, ...],
+        fused_io: Union | None,
+        *children: IR,
+    ):
+        self.schema = schema
+        self.subnodes = subnodes
+        self.fused_io = fused_io
+        self._non_child_args = (
+            [node.do_evaluate for node in subnodes],
+            [tuple(node._non_child_args) for node in self.subnodes],
+            [],
+            [],
+        )
+        self.children = children
+
+    @classmethod
+    def do_evaluate(
+        cls,
+        funcs: Sequence[Callable],
+        subargs: Sequence[Sequence[Any]],
+        io_funcs: Sequence[Callable],
+        io_args: Sequence[Any],
+        *children: DataFrame,
+    ) -> DataFrame:
+        """Evaluate and return a dataframe."""
+        if io_funcs:
+            children = (io_funcs[0](*io_args),)
+        for func, args in zip(funcs, subargs, strict=False):
+            children = (func(*args, *children),)
+        return children[0]
+
+
+def _fuse_ir_node(
+    ir: IR, rec: LowerIRTransformer
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    partition_info: MutableMapping[IR, PartitionInfo] = {}
+    children = ()
+    if ir.children:
+        children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
+        partition_info = reduce(operator.or_, _partition_info)
+
+    new_node: IR
+    if ir in rec.state["fusable"]:
+        if isinstance(ir, Union):
+            new_node = Fused(ir.schema, (), ir)
+        elif len(children) == 1 and isinstance(children[0], Fused):
+            (child,) = children
+            grandchildren = child.children
+            new_node = Fused(
+                ir.schema, (*child.subnodes, ir), child.fused_io, *grandchildren
+            )
+        else:
+            new_node = Fused(ir.schema, (ir,), None, *children)
+    else:
+        new_node = ir.reconstruct(children)
+
+    partition_info[new_node] = rec.state["partition_info"][ir]
+    return new_node, partition_info
+
+
+def _fusable_ir_type(ir: IR) -> bool:
+    return (
+        # Basic fusion
+        isinstance(
+            ir, (Cache, Filter, GroupBy, HStack, MapFunction, Projection, Select, Sort)
+        )
+        or (
+            # IO Fusion
+            isinstance(ir, Union)
+            and all(isinstance(n, (Scan, SplitScan)) for n in ir.children)
+        )
+    )
+
+
+@generate_ir_tasks.register(Fused)
+def _(
+    ir: Fused, partition_info: MutableMapping[IR, PartitionInfo]
+) -> MutableMapping[Any, Any]:
+    from cudf_polars.experimental.parallel import _default_generate_ir_tasks
+
+    if ir.fused_io is None:
+        return _default_generate_ir_tasks(ir, partition_info)
+
+    assert isinstance(ir.fused_io, Union)
+    graph: MutableMapping[Any, Any] = {}
+    for i, key in enumerate(partition_info[ir].keys(ir)):
+        io = ir.fused_io.children[i]
+        assert isinstance(io, (Scan, SplitScan))
+        graph[key] = (
+            ir.do_evaluate,
+            ir._non_child_args[0],
+            ir._non_child_args[1],
+            [io.do_evaluate],
+            list(io._non_child_args),
+        )
+    return graph
+
+
+def fuse_ir_graph(
+    ir: IR,
+    partition_info: MutableMapping[IR, PartitionInfo],
+    config_options: ConfigOptions,
+) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
+    """
+    Rewrite an IR graph with fused nodes.
+
+    Parameters
+    ----------
+    ir
+        Root of the graph to rewrite.
+    partition_info
+        Initial partitioning information.
+    config_options
+        GPUEngine configuration options.
+
+    Returns
+    -------
+    new_ir, partition_info
+        The rewritten graph, and a mapping from unique nodes
+        in the new graph to associated partitioning information.
+    """
+    assert config_options.executor.name == "streaming", (
+        "'in-memory' executor not supported in 'fuse_ir_graph'"
+    )
+    if not config_options.executor.task_fusion:
+        return ir, partition_info
+
+    parents: defaultdict[IR, int] = defaultdict(int)
+    if _fusable_ir_type(ir):
+        parents[ir] = 1
+    for node in traversal([ir]):
+        for child in node.children:
+            if _fusable_ir_type(child):
+                # Basic fusion
+                parents[child] += 1
+    fusable = {node for node, count in parents.items() if count == 1}
+    state = {"fusable": fusable, "partition_info": partition_info}
+    mapper = CachingVisitor(_fuse_ir_node, state=state)
+    return mapper(ir)

python/cudf_polars/cudf_polars/experimental/io.py

@@ -7,7 +7,6 @@
 import dataclasses
 import enum
 import math
-import random
 from enum import IntEnum
 from typing import TYPE_CHECKING, Any, TypeVar
 
@@ -263,9 +262,10 @@ def _sample_pq_statistics(ir: Scan) -> dict[str, np.floating[T]]:
     import numpy as np
 
     # Use average total_uncompressed_size of three files
-    n_sample = min(3, len(ir.paths))
+    n_sample = 5  # TODO: Make this configurable
+    stride = max(1, int(len(ir.paths) / n_sample))
     metadata = plc.io.parquet_metadata.read_parquet_metadata(
-        plc.io.SourceInfo(random.sample(ir.paths, n_sample))
+        plc.io.SourceInfo(ir.paths[: stride * n_sample : stride])
     )
     column_sizes = {}
     rowgroup_offsets_per_file = np.insert(

python/cudf_polars/cudf_polars/experimental/parallel.py

@@ -32,6 +32,7 @@
     generate_ir_tasks,
     lower_ir_node,
 )
+from cudf_polars.experimental.fusion import fuse_ir_graph
 from cudf_polars.experimental.utils import _concat, _lower_ir_fallback
 
 if TYPE_CHECKING:
@@ -82,7 +83,8 @@ def lower_ir_graph(
     lower_ir_node
     """
     mapper = CachingVisitor(lower_ir_node, state={"config_options": config_options})
-    return mapper(ir)
+    lowered_ir, partition_info = mapper(ir)
+    return fuse_ir_graph(lowered_ir, partition_info, config_options)
 
 
 def task_graph(
@@ -219,7 +221,7 @@ def evaluate_streaming(ir: IR, config_options: ConfigOptions) -> DataFrame:
 
 
 @generate_ir_tasks.register(IR)
-def _(
+def _default_generate_ir_tasks(
     ir: IR, partition_info: MutableMapping[IR, PartitionInfo]
 ) -> MutableMapping[Any, Any]:
     # Generate pointwise (embarrassingly-parallel) tasks by default
@@ -319,8 +321,8 @@ def _lower_ir_pwise(
 
 
 _lower_ir_pwise_preserve = partial(_lower_ir_pwise, preserve_partitioning=True)
-lower_ir_node.register(Projection, _lower_ir_pwise_preserve)
 lower_ir_node.register(Filter, _lower_ir_pwise_preserve)
+lower_ir_node.register(Projection, _lower_ir_pwise_preserve)
 lower_ir_node.register(Cache, _lower_ir_pwise)
 lower_ir_node.register(HStack, _lower_ir_pwise)
 lower_ir_node.register(HConcat, _lower_ir_pwise)

python/cudf_polars/cudf_polars/utils/config.py

@@ -191,6 +191,9 @@ class StreamingExecutor:
     rapidsmpf_spill
         Whether to wrap task arguments and output in objects that are
         spillable by 'rapidsmpf'.
+    task_fusion
+        Whether to fuse sequential partition-wise IR nodes into
+        a single task per partition.
     """
 
     name: Literal["streaming"] = dataclasses.field(default="streaming", init=False)
@@ -203,6 +206,7 @@ class StreamingExecutor:
     broadcast_join_limit: int = 0
     shuffle_method: ShuffleMethod | None = None
     rapidsmpf_spill: bool = False
+    task_fusion: bool = True
 
     def __post_init__(self) -> None:
         if self.scheduler == "synchronous" and self.shuffle_method == "rapidsmpf":
@@ -244,6 +248,8 @@ def __post_init__(self) -> None:
             raise TypeError("broadcast_join_limit must be an int")
         if not isinstance(self.rapidsmpf_spill, bool):
             raise TypeError("rapidsmpf_spill must be bool")
+        if not isinstance(self.task_fusion, bool):
+            raise TypeError("task_fusion must be bool")
 
     def __hash__(self) -> int:
         # cardinality factory, a dict, isn't natively hashable. We'll dump it

python/cudf_polars/tests/experimental/test_explain.py

@@ -63,6 +63,7 @@ def test_explain_physical_plan(tmp_path, df):
         executor_options={
             "target_partition_size": 10_000,
             "scheduler": DEFAULT_SCHEDULER,
+            "task_fusion": False,
         },
     )
 
@@ -73,7 +74,8 @@ def test_explain_physical_plan(tmp_path, df):
     assert "SELECT ('sum', 'y')" in plan or "PROJECTION ('sum', 'y')" in plan
 
 
-def test_explain_physical_plan_with_groupby(tmp_path, df):
+@pytest.mark.parametrize("task_fusion", [True, False])
+def test_explain_physical_plan_with_groupby(tmp_path, df, task_fusion):
     make_partitioned_source(df, tmp_path, fmt="parquet", n_files=1)
 
     q = (
@@ -89,12 +91,17 @@ def test_explain_physical_plan_with_groupby(tmp_path, df):
         executor_options={
             "target_partition_size": 10_000,
             "scheduler": DEFAULT_SCHEDULER,
+            "task_fusion": task_fusion,
         },
     )
 
-    plan = explain_query(q, engine, physical=True)
-
-    assert "GROUPBY ('g',)" in plan
+    if task_fusion:
+        plan = explain_query(q, engine, physical=True)
+        assert "FUSED ('GROUPBY'" in plan
+        assert "FUSED ('SPLITSCAN'" in plan
+    else:
+        plan = explain_query(q, engine, physical=True)
+        assert "GROUPBY ('g',)" in plan
 
 
 def test_explain_logical_plan_with_join(tmp_path, df):

python/cudf_polars/tests/experimental/test_scan.py

@@ -61,7 +61,7 @@ def test_target_partition_size(tmp_path, df, blocksize, n_files):
 
     # Check partitioning
     qir = Translator(q._ldf.visit(), engine).translate_ir()
-    ir, info = lower_ir_graph(qir, ConfigOptions(engine.config))
+    ir, info = lower_ir_graph(qir, ConfigOptions.from_polars_engine(engine))
     count = info[ir].count
     if blocksize <= 12_000:
         assert count > n_files