From 247eaaf0efb41010f87d807b34ef56907392aa89 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 13 May 2025 09:18:18 -0700 Subject: [PATCH 01/28] initial pass at IR fusion --- .../cudf_polars/experimental/explain.py | 19 ++- .../cudf_polars/experimental/parallel.py | 118 +++++++++++++++++- .../tests/experimental/test_explain.py | 2 +- 3 files changed, 133 insertions(+), 6 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/explain.py b/python/cudf_polars/cudf_polars/experimental/explain.py index 5e8def363be5..3581c6f858d8 100644 --- a/python/cudf_polars/cudf_polars/experimental/explain.py +++ b/python/cudf_polars/cudf_polars/experimental/explain.py @@ -17,7 +17,7 @@ Sort, ) from cudf_polars.dsl.translate import Translator -from cudf_polars.experimental.parallel import lower_ir_graph +from cudf_polars.experimental.parallel import Fused, fuse_ir_graph, lower_ir_graph from cudf_polars.utils.config import ConfigOptions if TYPE_CHECKING: @@ -30,7 +30,11 @@ def explain_query( - q: pl.LazyFrame, engine: pl.GPUEngine, *, physical: bool = True + q: pl.LazyFrame, + engine: pl.GPUEngine, + *, + physical: bool = True, + fuse: bool = True, ) -> str: """ Return a formatted string representation of the IR plan. @@ -44,6 +48,9 @@ def explain_query( physical : bool, default True If True, show the physical (lowered) plan. If False, show the logical (pre-lowering) plan. + fuse : bool, default True + Whether to show the fused physical plan. + Ignored if physical is False. Returns ------- @@ -55,6 +62,8 @@ def explain_query( if physical: lowered_ir, partition_info = lower_ir_graph(ir, config) + if fuse: + lowered_ir, partition_info = fuse_ir_graph(lowered_ir, partition_info) return _repr_ir_tree(lowered_ir, partition_info) else: return _repr_ir_tree(ir) @@ -125,3 +134,9 @@ 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) + return _repr_header(offset, f"FUSED {nodes}", ir.schema) diff --git a/python/cudf_polars/cudf_polars/experimental/parallel.py b/python/cudf_polars/cudf_polars/experimental/parallel.py index 0ceaf409c473..1462d16fe287 100644 --- a/python/cudf_polars/cudf_polars/experimental/parallel.py +++ b/python/cudf_polars/cudf_polars/experimental/parallel.py @@ -6,6 +6,7 @@ import itertools import operator +from collections import defaultdict from functools import partial, reduce from typing import TYPE_CHECKING, Any @@ -20,10 +21,13 @@ IR, Cache, Filter, + GroupBy, HConcat, HStack, MapFunction, Projection, + Select, + Sort, Union, ) from cudf_polars.dsl.traversal import CachingVisitor, traversal @@ -35,11 +39,12 @@ from cudf_polars.experimental.utils import _concat, _lower_ir_fallback if TYPE_CHECKING: - from collections.abc import MutableMapping + from collections.abc import Callable, MutableMapping, Sequence from typing import Any from cudf_polars.containers import DataFrame from cudf_polars.experimental.dispatch import LowerIRTransformer + from cudf_polars.typing import Schema from cudf_polars.utils.config import ConfigOptions @@ -53,6 +58,96 @@ def _( ) +class Fused(IR): + """Fused IR node.""" + + __slots__ = ("subnodes",) + _non_child = ("schema", "subnodes") + subnodes: tuple[IR, ...] + """Fused sub-nodes.""" + + def __init__( + self, + schema: Schema, + subnodes: tuple[IR, ...], + *children: IR, + ): + self.schema = schema + self.subnodes = subnodes + 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]], + *children: DataFrame, + ) -> DataFrame: + """Evaluate and return a dataframe.""" + 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 len(children) == 1 and isinstance(children[0], Fused): + (child,) = children + grandchildren = child.children + new_node = Fused(ir.schema, (*child.subnodes, ir), *grandchildren) + else: + new_node = Fused(ir.schema, (ir,), *children) + else: + new_node = ir.reconstruct(children) + + partition_info[new_node] = rec.state["partition_info"][ir] + return new_node, partition_info + + +def fuse_ir_graph( + ir: IR, partition_info: MutableMapping[IR, PartitionInfo] +) -> 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. + + Returns + ------- + new_ir, partition_info + The rewritten graph, and a mapping from unique nodes + in the new graph to associated partitioning information. + """ + parents: defaultdict[IR, int] = defaultdict(int) + for node in traversal([ir]): + for child in node.children: + if isinstance(child, (Select, GroupBy, Projection, Filter, Sort)): + 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) + + def lower_ir_graph( ir: IR, config_options: ConfigOptions ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: @@ -211,6 +306,8 @@ def evaluate_streaming(ir: IR, config_options: ConfigOptions) -> DataFrame: """ ir, partition_info = lower_ir_graph(ir, config_options) + ir, partition_info = fuse_ir_graph(ir, partition_info) + graph, key = task_graph(ir, partition_info) graph = post_process_task_graph(graph, key, config_options) @@ -319,8 +416,23 @@ 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(Cache, _lower_ir_pwise) lower_ir_node.register(HStack, _lower_ir_pwise) lower_ir_node.register(HConcat, _lower_ir_pwise) + + +@lower_ir_node.register(Cache) +def _( + ir: Cache, rec: LowerIRTransformer +) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: + return rec(ir.children[0]) + + +@lower_ir_node.register(Projection) +def _( + ir: Projection, rec: LowerIRTransformer +) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: + child, pi = rec(ir.children[0]) + if ir.schema == child.schema: + return child, pi + return _lower_ir_pwise_preserve(ir, rec) diff --git a/python/cudf_polars/tests/experimental/test_explain.py b/python/cudf_polars/tests/experimental/test_explain.py index d251860002dd..ee0d64e57e50 100644 --- a/python/cudf_polars/tests/experimental/test_explain.py +++ b/python/cudf_polars/tests/experimental/test_explain.py @@ -92,7 +92,7 @@ def test_explain_physical_plan_with_groupby(tmp_path, df): }, ) - plan = explain_query(q, engine, physical=True) + plan = explain_query(q, engine, physical=True, fuse=False) assert "GROUPBY ('g',)" in plan From 2c882d72affc3c9ee323f17b452e6fb174249369 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 13 May 2025 10:23:45 -0700 Subject: [PATCH 02/28] add general task-fusion support --- .../cudf_polars/experimental/explain.py | 2 + .../cudf_polars/experimental/parallel.py | 68 ++++++++++++++++--- 2 files changed, 61 insertions(+), 9 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/explain.py b/python/cudf_polars/cudf_polars/experimental/explain.py index 3581c6f858d8..58b12265491e 100644 --- a/python/cudf_polars/cudf_polars/experimental/explain.py +++ b/python/cudf_polars/cudf_polars/experimental/explain.py @@ -139,4 +139,6 @@ def _(ir: Scan, *, offset: str = "") -> str: @_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) diff --git a/python/cudf_polars/cudf_polars/experimental/parallel.py b/python/cudf_polars/cudf_polars/experimental/parallel.py index 1462d16fe287..1980f4101cb8 100644 --- a/python/cudf_polars/cudf_polars/experimental/parallel.py +++ b/python/cudf_polars/cudf_polars/experimental/parallel.py @@ -12,7 +12,6 @@ import cudf_polars.experimental.distinct import cudf_polars.experimental.groupby -import cudf_polars.experimental.io import cudf_polars.experimental.join import cudf_polars.experimental.select import cudf_polars.experimental.shuffle @@ -36,6 +35,7 @@ generate_ir_tasks, lower_ir_node, ) +from cudf_polars.experimental.io import Scan, SplitScan from cudf_polars.experimental.utils import _concat, _lower_ir_fallback if TYPE_CHECKING: @@ -61,22 +61,28 @@ def _( class Fused(IR): """Fused IR node.""" - __slots__ = ("subnodes",) - _non_child = ("schema", "subnodes") + __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], + None, + (), ) self.children = children @@ -85,9 +91,13 @@ def do_evaluate( cls, funcs: Sequence[Callable], subargs: Sequence[Sequence[Any]], + io_func: Callable | None, + io_args: Sequence[Any], *children: DataFrame, ) -> DataFrame: """Evaluate and return a dataframe.""" + if io_func: + children = (io_func(*io_args),) for func, args in zip(funcs, subargs, strict=False): children = (func(*args, *children),) return children[0] @@ -104,12 +114,16 @@ def _fuse_ir_node( new_node: IR if ir in rec.state["fusable"]: - if len(children) == 1 and isinstance(children[0], Fused): + 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), *grandchildren) + new_node = Fused( + ir.schema, (*child.subnodes, ir), child.fused_io, *grandchildren + ) else: - new_node = Fused(ir.schema, (ir,), *children) + new_node = Fused(ir.schema, (ir,), None, *children) else: new_node = ir.reconstruct(children) @@ -117,6 +131,18 @@ def _fuse_ir_node( return new_node, partition_info +def _fusable_ir_type(ir: IR) -> bool: + return ( + # Basic fusion + isinstance(ir, (Filter, GroupBy, HStack, MapFunction, Projection, Select, Sort)) + or ( + # IO Fusion + isinstance(ir, Union) + and all(isinstance(n, (Scan, SplitScan)) for n in ir.children) + ) + ) + + def fuse_ir_graph( ir: IR, partition_info: MutableMapping[IR, PartitionInfo] ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: @@ -137,11 +163,13 @@ def fuse_ir_graph( in the new graph to associated partitioning information. """ 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 isinstance(child, (Select, GroupBy, Projection, Filter, Sort)): + 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) @@ -316,7 +344,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 @@ -335,6 +363,28 @@ def _( } +@generate_ir_tasks.register(Fused) +def _( + ir: Fused, partition_info: MutableMapping[IR, PartitionInfo] +) -> MutableMapping[Any, Any]: + 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, + io._non_child_args, + ) + return graph + + @lower_ir_node.register(Union) def _( ir: Union, rec: LowerIRTransformer From 65bff6862c45d9500d491fcc0b1912d7a0baef96 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 13 May 2025 10:26:37 -0700 Subject: [PATCH 03/28] experimental tests passing --- python/cudf_polars/tests/experimental/test_explain.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf_polars/tests/experimental/test_explain.py b/python/cudf_polars/tests/experimental/test_explain.py index ee0d64e57e50..8696f7259815 100644 --- a/python/cudf_polars/tests/experimental/test_explain.py +++ b/python/cudf_polars/tests/experimental/test_explain.py @@ -66,7 +66,7 @@ def test_explain_physical_plan(tmp_path, df): }, ) - plan = explain_query(q, engine) + plan = explain_query(q, engine, fuse=False) assert "UNION" in plan assert "SPLITSCAN" in plan From ecf0db52eb004431f229db878d7505f32d309626 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 13 May 2025 10:41:57 -0700 Subject: [PATCH 04/28] deterministic partitioning --- python/cudf_polars/cudf_polars/experimental/io.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/experimental/io.py b/python/cudf_polars/cudf_polars/experimental/io.py index 216d5b7ae662..34424d7f018f 100644 --- a/python/cudf_polars/cudf_polars/experimental/io.py +++ b/python/cudf_polars/cudf_polars/experimental/io.py @@ -263,7 +263,8 @@ 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 = min(5, len(ir.paths)) + random.seed(42) # Deterministic partitioning metadata = plc.io.parquet_metadata.read_parquet_metadata( plc.io.SourceInfo(random.sample(ir.paths, n_sample)) ) From 3da460ce214b9610f926c980c9f15ea3d54b5c4f Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 13 May 2025 13:02:20 -0700 Subject: [PATCH 05/28] fix shuffle bug --- python/cudf_polars/cudf_polars/experimental/shuffle.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/shuffle.py b/python/cudf_polars/cudf_polars/experimental/shuffle.py index 136c930fc6c5..475415057f50 100644 --- a/python/cudf_polars/cudf_polars/experimental/shuffle.py +++ b/python/cudf_polars/cudf_polars/experimental/shuffle.py @@ -57,9 +57,10 @@ def insert_partition( on = options["on"] assert not other, f"Unexpected arguments: {other}" - columns_to_hash = tuple(df.column_names.index(val) for val in on) + column_names = options["column_names"] + columns_to_hash = tuple(column_names.index(val) for val in on) packed_inputs = partition_and_pack( - df.table, + df.select(column_names).table, columns_to_hash=columns_to_hash, num_partitions=partition_count, stream=DEFAULT_STREAM, From 08e8205f0c1904ee571a488c49f076e7d09529ca Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 13 May 2025 13:27:01 -0700 Subject: [PATCH 06/28] roll back breaking changes --- .../cudf_polars/experimental/parallel.py | 23 ++++--------------- .../cudf_polars/experimental/shuffle.py | 5 ++-- 2 files changed, 7 insertions(+), 21 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/parallel.py b/python/cudf_polars/cudf_polars/experimental/parallel.py index 1980f4101cb8..0ff1b8656bf5 100644 --- a/python/cudf_polars/cudf_polars/experimental/parallel.py +++ b/python/cudf_polars/cudf_polars/experimental/parallel.py @@ -134,7 +134,9 @@ def _fuse_ir_node( def _fusable_ir_type(ir: IR) -> bool: return ( # Basic fusion - isinstance(ir, (Filter, GroupBy, HStack, MapFunction, Projection, Select, Sort)) + isinstance( + ir, (Cache, Filter, GroupBy, HStack, MapFunction, Projection, Select, Sort) + ) or ( # IO Fusion isinstance(ir, Union) @@ -467,22 +469,7 @@ def _lower_ir_pwise( _lower_ir_pwise_preserve = partial(_lower_ir_pwise, preserve_partitioning=True) 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) - - -@lower_ir_node.register(Cache) -def _( - ir: Cache, rec: LowerIRTransformer -) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: - return rec(ir.children[0]) - - -@lower_ir_node.register(Projection) -def _( - ir: Projection, rec: LowerIRTransformer -) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: - child, pi = rec(ir.children[0]) - if ir.schema == child.schema: - return child, pi - return _lower_ir_pwise_preserve(ir, rec) diff --git a/python/cudf_polars/cudf_polars/experimental/shuffle.py b/python/cudf_polars/cudf_polars/experimental/shuffle.py index 475415057f50..136c930fc6c5 100644 --- a/python/cudf_polars/cudf_polars/experimental/shuffle.py +++ b/python/cudf_polars/cudf_polars/experimental/shuffle.py @@ -57,10 +57,9 @@ def insert_partition( on = options["on"] assert not other, f"Unexpected arguments: {other}" - column_names = options["column_names"] - columns_to_hash = tuple(column_names.index(val) for val in on) + columns_to_hash = tuple(df.column_names.index(val) for val in on) packed_inputs = partition_and_pack( - df.select(column_names).table, + df.table, columns_to_hash=columns_to_hash, num_partitions=partition_count, stream=DEFAULT_STREAM, From ddab70092f5267b3b106a6baa4f6d7365e811f17 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 13 May 2025 14:05:56 -0700 Subject: [PATCH 07/28] fix spilling --- .../cudf_polars/experimental/parallel.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/parallel.py b/python/cudf_polars/cudf_polars/experimental/parallel.py index 0ff1b8656bf5..9100505aaaec 100644 --- a/python/cudf_polars/cudf_polars/experimental/parallel.py +++ b/python/cudf_polars/cudf_polars/experimental/parallel.py @@ -81,8 +81,8 @@ def __init__( self._non_child_args = ( [node.do_evaluate for node in subnodes], [tuple(node._non_child_args) for node in self.subnodes], - None, - (), + [], + [], ) self.children = children @@ -91,13 +91,13 @@ def do_evaluate( cls, funcs: Sequence[Callable], subargs: Sequence[Sequence[Any]], - io_func: Callable | None, + io_funcs: Sequence[Callable], io_args: Sequence[Any], *children: DataFrame, ) -> DataFrame: """Evaluate and return a dataframe.""" - if io_func: - children = (io_func(*io_args),) + 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] @@ -381,8 +381,8 @@ def _( ir.do_evaluate, ir._non_child_args[0], ir._non_child_args[1], - io.do_evaluate, - io._non_child_args, + [io.do_evaluate], + list(io._non_child_args), ) return graph From ccb909f372441fbb544acebf7ea6e4ba5c2b38a9 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 14 May 2025 06:02:47 -0700 Subject: [PATCH 08/28] avoid random in metadata sampling --- python/cudf_polars/cudf_polars/experimental/io.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/io.py b/python/cudf_polars/cudf_polars/experimental/io.py index 34424d7f018f..c85ffe3abe1c 100644 --- a/python/cudf_polars/cudf_polars/experimental/io.py +++ b/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,10 +262,9 @@ 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(5, len(ir.paths)) - random.seed(42) # Deterministic partitioning metadata = plc.io.parquet_metadata.read_parquet_metadata( - plc.io.SourceInfo(random.sample(ir.paths, n_sample)) + # TODO: Add configuration for random sampling + plc.io.SourceInfo(ir.paths[:5]) ) column_sizes = {} rowgroup_offsets_per_file = np.insert( From fddf0934a5f840e163249a21438099805baf0916 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 14 May 2025 06:03:18 -0700 Subject: [PATCH 09/28] avoid random in metadata sampling --- python/cudf_polars/cudf_polars/experimental/io.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/io.py b/python/cudf_polars/cudf_polars/experimental/io.py index c85ffe3abe1c..256681aee62d 100644 --- a/python/cudf_polars/cudf_polars/experimental/io.py +++ b/python/cudf_polars/cudf_polars/experimental/io.py @@ -262,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 = 5 # TODO: Make this configurable + stride = max(1, int(len(ir.paths) / n_sample)) metadata = plc.io.parquet_metadata.read_parquet_metadata( - # TODO: Add configuration for random sampling - plc.io.SourceInfo(ir.paths[:5]) + plc.io.SourceInfo(ir.paths[: stride * n_sample : stride]) ) column_sizes = {} rowgroup_offsets_per_file = np.insert( From 6db02ec79e5a87aa7530f91896f73c99c6261e77 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 14 May 2025 06:19:50 -0700 Subject: [PATCH 10/28] add test coverage for explain --- python/cudf_polars/tests/experimental/test_explain.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/tests/experimental/test_explain.py b/python/cudf_polars/tests/experimental/test_explain.py index 8696f7259815..37a269433ffd 100644 --- a/python/cudf_polars/tests/experimental/test_explain.py +++ b/python/cudf_polars/tests/experimental/test_explain.py @@ -93,9 +93,12 @@ def test_explain_physical_plan_with_groupby(tmp_path, df): ) plan = explain_query(q, engine, physical=True, fuse=False) - assert "GROUPBY ('g',)" in plan + plan = explain_query(q, engine, physical=True, fuse=True) + assert "FUSED ('GROUPBY'" in plan + assert "FUSED ('SPLITSCAN'" in plan + def test_explain_logical_plan_with_join(tmp_path, df): make_partitioned_source(df, tmp_path, fmt="parquet", n_files=2) From a857850b9f3a3e65d63fedce2622c63666e0b3af Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 14 May 2025 06:54:24 -0700 Subject: [PATCH 11/28] refactor --- .../cudf_polars/experimental/explain.py | 9 +- .../cudf_polars/experimental/fusion.py | 190 ++++++++++++++++++ .../cudf_polars/experimental/parallel.py | 157 +-------------- .../cudf_polars/cudf_polars/utils/config.py | 6 + .../tests/experimental/test_explain.py | 20 +- .../tests/experimental/test_scan.py | 2 +- 6 files changed, 216 insertions(+), 168 deletions(-) create mode 100644 python/cudf_polars/cudf_polars/experimental/fusion.py diff --git a/python/cudf_polars/cudf_polars/experimental/explain.py b/python/cudf_polars/cudf_polars/experimental/explain.py index 58b12265491e..dd693425cdbf 100644 --- a/python/cudf_polars/cudf_polars/experimental/explain.py +++ b/python/cudf_polars/cudf_polars/experimental/explain.py @@ -17,7 +17,8 @@ Sort, ) from cudf_polars.dsl.translate import Translator -from cudf_polars.experimental.parallel import Fused, fuse_ir_graph, lower_ir_graph +from cudf_polars.experimental.fusion import Fused +from cudf_polars.experimental.parallel import lower_ir_graph from cudf_polars.utils.config import ConfigOptions if TYPE_CHECKING: @@ -34,7 +35,6 @@ def explain_query( engine: pl.GPUEngine, *, physical: bool = True, - fuse: bool = True, ) -> str: """ Return a formatted string representation of the IR plan. @@ -48,9 +48,6 @@ def explain_query( physical : bool, default True If True, show the physical (lowered) plan. If False, show the logical (pre-lowering) plan. - fuse : bool, default True - Whether to show the fused physical plan. - Ignored if physical is False. Returns ------- @@ -62,8 +59,6 @@ def explain_query( if physical: lowered_ir, partition_info = lower_ir_graph(ir, config) - if fuse: - lowered_ir, partition_info = fuse_ir_graph(lowered_ir, partition_info) return _repr_ir_tree(lowered_ir, partition_info) else: return _repr_ir_tree(ir) diff --git a/python/cudf_polars/cudf_polars/experimental/fusion.py b/python/cudf_polars/cudf_polars/experimental/fusion.py new file mode 100644 index 000000000000..d037762b51ff --- /dev/null +++ b/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) diff --git a/python/cudf_polars/cudf_polars/experimental/parallel.py b/python/cudf_polars/cudf_polars/experimental/parallel.py index 9100505aaaec..b7ab235507d9 100644 --- a/python/cudf_polars/cudf_polars/experimental/parallel.py +++ b/python/cudf_polars/cudf_polars/experimental/parallel.py @@ -6,12 +6,12 @@ import itertools import operator -from collections import defaultdict from functools import partial, reduce from typing import TYPE_CHECKING, Any import cudf_polars.experimental.distinct import cudf_polars.experimental.groupby +import cudf_polars.experimental.io import cudf_polars.experimental.join import cudf_polars.experimental.select import cudf_polars.experimental.shuffle @@ -20,13 +20,10 @@ IR, Cache, Filter, - GroupBy, HConcat, HStack, MapFunction, Projection, - Select, - Sort, Union, ) from cudf_polars.dsl.traversal import CachingVisitor, traversal @@ -35,16 +32,15 @@ generate_ir_tasks, lower_ir_node, ) -from cudf_polars.experimental.io import Scan, SplitScan +from cudf_polars.experimental.fusion import fuse_ir_graph from cudf_polars.experimental.utils import _concat, _lower_ir_fallback if TYPE_CHECKING: - from collections.abc import Callable, MutableMapping, Sequence + from collections.abc import MutableMapping from typing import Any from cudf_polars.containers import DataFrame from cudf_polars.experimental.dispatch import LowerIRTransformer - from cudf_polars.typing import Schema from cudf_polars.utils.config import ConfigOptions @@ -58,126 +54,6 @@ def _( ) -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) - ) - ) - - -def fuse_ir_graph( - ir: IR, partition_info: MutableMapping[IR, PartitionInfo] -) -> 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. - - Returns - ------- - new_ir, partition_info - The rewritten graph, and a mapping from unique nodes - in the new graph to associated partitioning information. - """ - 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) - - def lower_ir_graph( ir: IR, config_options: ConfigOptions ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: @@ -207,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( @@ -336,8 +213,6 @@ def evaluate_streaming(ir: IR, config_options: ConfigOptions) -> DataFrame: """ ir, partition_info = lower_ir_graph(ir, config_options) - ir, partition_info = fuse_ir_graph(ir, partition_info) - graph, key = task_graph(ir, partition_info) graph = post_process_task_graph(graph, key, config_options) @@ -365,28 +240,6 @@ def _default_generate_ir_tasks( } -@generate_ir_tasks.register(Fused) -def _( - ir: Fused, partition_info: MutableMapping[IR, PartitionInfo] -) -> MutableMapping[Any, Any]: - 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 - - @lower_ir_node.register(Union) def _( ir: Union, rec: LowerIRTransformer diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 7f29415457f3..0b70bec6c14e 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/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 diff --git a/python/cudf_polars/tests/experimental/test_explain.py b/python/cudf_polars/tests/experimental/test_explain.py index 37a269433ffd..0963297e76ad 100644 --- a/python/cudf_polars/tests/experimental/test_explain.py +++ b/python/cudf_polars/tests/experimental/test_explain.py @@ -63,17 +63,19 @@ def test_explain_physical_plan(tmp_path, df): executor_options={ "target_partition_size": 10_000, "scheduler": DEFAULT_SCHEDULER, + "task_fusion": False, }, ) - plan = explain_query(q, engine, fuse=False) + plan = explain_query(q, engine) assert "UNION" in plan assert "SPLITSCAN" in plan 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,15 +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, fuse=False) - assert "GROUPBY ('g',)" in plan - - plan = explain_query(q, engine, physical=True, fuse=True) - assert "FUSED ('GROUPBY'" in plan - assert "FUSED ('SPLITSCAN'" 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): diff --git a/python/cudf_polars/tests/experimental/test_scan.py b/python/cudf_polars/tests/experimental/test_scan.py index df968f6a6872..ca8383842129 100644 --- a/python/cudf_polars/tests/experimental/test_scan.py +++ b/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 From c129bc3cfb7159bbd10e8792e3fbc6a91f3d45ca Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 14 May 2025 08:32:59 -0700 Subject: [PATCH 12/28] avoid calling fuse_ir_graph unless task_fusion is True --- python/cudf_polars/cudf_polars/experimental/fusion.py | 10 ---------- .../cudf_polars/cudf_polars/experimental/parallel.py | 8 +++++++- 2 files changed, 7 insertions(+), 11 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/fusion.py b/python/cudf_polars/cudf_polars/experimental/fusion.py index d037762b51ff..04032a3f3484 100644 --- a/python/cudf_polars/cudf_polars/experimental/fusion.py +++ b/python/cudf_polars/cudf_polars/experimental/fusion.py @@ -33,7 +33,6 @@ 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): @@ -150,7 +149,6 @@ def _( 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. @@ -161,8 +159,6 @@ def fuse_ir_graph( Root of the graph to rewrite. partition_info Initial partitioning information. - config_options - GPUEngine configuration options. Returns ------- @@ -170,12 +166,6 @@ def fuse_ir_graph( 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 diff --git a/python/cudf_polars/cudf_polars/experimental/parallel.py b/python/cudf_polars/cudf_polars/experimental/parallel.py index b7ab235507d9..ff296c867007 100644 --- a/python/cudf_polars/cudf_polars/experimental/parallel.py +++ b/python/cudf_polars/cudf_polars/experimental/parallel.py @@ -82,9 +82,15 @@ def lower_ir_graph( -------- lower_ir_node """ + assert config_options.executor.name == "streaming", ( + "'in-memory' executor not supported in 'lower_ir_graph'" + ) + mapper = CachingVisitor(lower_ir_node, state={"config_options": config_options}) lowered_ir, partition_info = mapper(ir) - return fuse_ir_graph(lowered_ir, partition_info, config_options) + if config_options.executor.task_fusion: + return fuse_ir_graph(lowered_ir, partition_info) + return lowered_ir, partition_info def task_graph( From f9056ed63879e4bbe391bcfa9cacccdd227bb1c8 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 14 May 2025 08:37:04 -0700 Subject: [PATCH 13/28] bump coverage --- python/cudf_polars/tests/test_config.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index aa9597347ee6..1a0c526f6002 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -226,6 +226,7 @@ def test_validate_shuffle_method() -> None: "groupby_n_ary", "broadcast_join_limit", "rapidsmpf_spill", + "task_fusion", ], ) def test_validate_max_rows_per_partition(option: str) -> None: From d03fc038b3a4e0f6e04611ff42ff73d1e5e8c426 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 14 May 2025 08:40:08 -0700 Subject: [PATCH 14/28] drop unnecessary assertion --- python/cudf_polars/cudf_polars/experimental/fusion.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/experimental/fusion.py b/python/cudf_polars/cudf_polars/experimental/fusion.py index 04032a3f3484..831cee583206 100644 --- a/python/cudf_polars/cudf_polars/experimental/fusion.py +++ b/python/cudf_polars/cudf_polars/experimental/fusion.py @@ -131,7 +131,6 @@ def _( 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] From 5e014db8340904f558d9b63e742f4b901a069e6d Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 14 May 2025 08:41:19 -0700 Subject: [PATCH 15/28] fix strict=False added by pre-commit --- python/cudf_polars/cudf_polars/experimental/fusion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/experimental/fusion.py b/python/cudf_polars/cudf_polars/experimental/fusion.py index 831cee583206..d9b784c3e8c7 100644 --- a/python/cudf_polars/cudf_polars/experimental/fusion.py +++ b/python/cudf_polars/cudf_polars/experimental/fusion.py @@ -75,7 +75,7 @@ def do_evaluate( """Evaluate and return a dataframe.""" if io_funcs: children = (io_funcs[0](*io_args),) - for func, args in zip(funcs, subargs, strict=False): + for func, args in zip(funcs, subargs, strict=True): children = (func(*args, *children),) return children[0] From c4751d2a08c0725b9fa4dec253045d72eecd77c1 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 14 May 2025 08:44:05 -0700 Subject: [PATCH 16/28] address review --- .../cudf_polars/cudf_polars/experimental/fusion.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/fusion.py b/python/cudf_polars/cudf_polars/experimental/fusion.py index d9b784c3e8c7..2b5859e2f4d0 100644 --- a/python/cudf_polars/cudf_polars/experimental/fusion.py +++ b/python/cudf_polars/cudf_polars/experimental/fusion.py @@ -90,7 +90,7 @@ def _fuse_ir_node( partition_info = reduce(operator.or_, _partition_info) new_node: IR - if ir in rec.state["fusable"]: + if ir in rec.state["fusible"]: if isinstance(ir, Union): new_node = Fused(ir.schema, (), ir) elif len(children) == 1 and isinstance(children[0], Fused): @@ -108,7 +108,7 @@ def _fuse_ir_node( return new_node, partition_info -def _fusable_ir_type(ir: IR) -> bool: +def _is_fusible_ir_type(ir: IR) -> bool: return ( # Basic fusion isinstance( @@ -155,7 +155,7 @@ def fuse_ir_graph( Parameters ---------- ir - Root of the graph to rewrite. + Root of the lowered-IR graph to rewrite. partition_info Initial partitioning information. @@ -166,14 +166,14 @@ def fuse_ir_graph( in the new graph to associated partitioning information. """ parents: defaultdict[IR, int] = defaultdict(int) - if _fusable_ir_type(ir): + if _is_fusible_ir_type(ir): parents[ir] = 1 for node in traversal([ir]): for child in node.children: - if _fusable_ir_type(child): + if _is_fusible_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} + fusible = {node for node, count in parents.items() if count == 1} + state = {"fusible": fusible, "partition_info": partition_info} mapper = CachingVisitor(_fuse_ir_node, state=state) return mapper(ir) From deb549b4be3080d40e47ec8c30cd4c4c939489fd Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 14 May 2025 09:00:49 -0700 Subject: [PATCH 17/28] improve comments --- .../cudf_polars/experimental/fusion.py | 20 +++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/fusion.py b/python/cudf_polars/cudf_polars/experimental/fusion.py index 2b5859e2f4d0..7142cd8ded05 100644 --- a/python/cudf_polars/cudf_polars/experimental/fusion.py +++ b/python/cudf_polars/cudf_polars/experimental/fusion.py @@ -108,16 +108,25 @@ def _fuse_ir_node( return new_node, partition_info -def _is_fusible_ir_type(ir: IR) -> bool: +def _is_fusible_ir_node( + ir: IR, partition_info: MutableMapping[IR, PartitionInfo] +) -> bool: return ( - # Basic fusion + # Basic fusion. + # These are partition-wise operations if they exist + # in the graph after lowering (the data-movement + # has been encoded with Shuffle/Reduction/etc). isinstance( ir, (Cache, Filter, GroupBy, HStack, MapFunction, Projection, Select, Sort) ) or ( - # IO Fusion + # IO Fusion. + # Multi-partition IO is always implemented as a Union + # over Scan/SplitScan nodes. The number of children + # must correspond to the number of partitions. isinstance(ir, Union) and all(isinstance(n, (Scan, SplitScan)) for n in ir.children) + and partition_info[ir].count == len(ir.children) ) ) @@ -166,12 +175,11 @@ def fuse_ir_graph( in the new graph to associated partitioning information. """ parents: defaultdict[IR, int] = defaultdict(int) - if _is_fusible_ir_type(ir): + if _is_fusible_ir_node(ir, partition_info): parents[ir] = 1 for node in traversal([ir]): for child in node.children: - if _is_fusible_ir_type(child): - # Basic fusion + if _is_fusible_ir_node(child, partition_info): parents[child] += 1 fusible = {node for node, count in parents.items() if count == 1} state = {"fusible": fusible, "partition_info": partition_info} From d07c3501acef981fc775d727050d5d99ca79ed40 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 14 May 2025 13:40:33 -0700 Subject: [PATCH 18/28] avoid fusing Sort nodes for now --- python/cudf_polars/cudf_polars/experimental/fusion.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/fusion.py b/python/cudf_polars/cudf_polars/experimental/fusion.py index 7142cd8ded05..d14b0f70751d 100644 --- a/python/cudf_polars/cudf_polars/experimental/fusion.py +++ b/python/cudf_polars/cudf_polars/experimental/fusion.py @@ -18,7 +18,6 @@ MapFunction, Projection, Select, - Sort, Union, ) from cudf_polars.dsl.traversal import CachingVisitor, traversal @@ -117,7 +116,7 @@ def _is_fusible_ir_node( # in the graph after lowering (the data-movement # has been encoded with Shuffle/Reduction/etc). isinstance( - ir, (Cache, Filter, GroupBy, HStack, MapFunction, Projection, Select, Sort) + ir, (Cache, Filter, GroupBy, HStack, MapFunction, Projection, Select) ) or ( # IO Fusion. @@ -143,7 +142,11 @@ def _( 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)) + try: + assert isinstance(io, (Scan, SplitScan)) + except Exception as err: + import pdb; pdb.set_trace() + pass graph[key] = ( ir.do_evaluate, ir._non_child_args[0], From 00e80981ef0b0d5db12363ecff9855ead8551255 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 14 May 2025 13:47:27 -0700 Subject: [PATCH 19/28] pre-commit --- python/cudf_polars/cudf_polars/experimental/fusion.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/fusion.py b/python/cudf_polars/cudf_polars/experimental/fusion.py index d14b0f70751d..10680c0698a6 100644 --- a/python/cudf_polars/cudf_polars/experimental/fusion.py +++ b/python/cudf_polars/cudf_polars/experimental/fusion.py @@ -116,7 +116,8 @@ def _is_fusible_ir_node( # in the graph after lowering (the data-movement # has been encoded with Shuffle/Reduction/etc). isinstance( - ir, (Cache, Filter, GroupBy, HStack, MapFunction, Projection, Select) + ir, + (Cache, Filter, GroupBy, HStack, MapFunction, Projection, Select), ) or ( # IO Fusion. @@ -142,11 +143,7 @@ def _( graph: MutableMapping[Any, Any] = {} for i, key in enumerate(partition_info[ir].keys(ir)): io = ir.fused_io.children[i] - try: - assert isinstance(io, (Scan, SplitScan)) - except Exception as err: - import pdb; pdb.set_trace() - pass + assert isinstance(io, (Scan, SplitScan)) graph[key] = ( ir.do_evaluate, ir._non_child_args[0], From 8c6efcf25f2f2ade0590095d77dc0495be684b5f Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 15 May 2025 10:59:19 -0700 Subject: [PATCH 20/28] cleanup --- .../cudf_polars/experimental/explain.py | 4 +- .../cudf_polars/experimental/fusion.py | 198 ++++++++++++------ .../cudf_polars/experimental/parallel.py | 2 +- 3 files changed, 133 insertions(+), 71 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/explain.py b/python/cudf_polars/cudf_polars/experimental/explain.py index dd693425cdbf..b275c8bcec5a 100644 --- a/python/cudf_polars/cudf_polars/experimental/explain.py +++ b/python/cudf_polars/cudf_polars/experimental/explain.py @@ -17,7 +17,7 @@ Sort, ) from cudf_polars.dsl.translate import Translator -from cudf_polars.experimental.fusion import Fused +from cudf_polars.experimental.fusion import Fused, FusedIO from cudf_polars.experimental.parallel import lower_ir_graph from cudf_polars.utils.config import ConfigOptions @@ -134,6 +134,6 @@ def _(ir: Scan, *, offset: str = "") -> str: @_repr_ir.register def _(ir: Fused, *, offset: str = "") -> str: nodes = tuple(type(node).__name__.upper() for node in ir.subnodes) - if ir.fused_io: + if isinstance(ir, FusedIO): nodes = (type(ir.fused_io.children[0]).__name__.upper(), *nodes) return _repr_header(offset, f"FUSED {nodes}", ir.schema) diff --git a/python/cudf_polars/cudf_polars/experimental/fusion.py b/python/cudf_polars/cudf_polars/experimental/fusion.py index 10680c0698a6..857ea90056e4 100644 --- a/python/cudf_polars/cudf_polars/experimental/fusion.py +++ b/python/cudf_polars/cudf_polars/experimental/fusion.py @@ -7,7 +7,7 @@ import operator from collections import defaultdict from functools import reduce -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Self from cudf_polars.dsl.ir import ( IR, @@ -37,48 +37,147 @@ class Fused(IR): """Fused IR node.""" - __slots__ = ("fused_io", "subnodes") - _non_child = ("schema", "subnodes", "fused_io") + __slots__ = ("fusible", "subnodes") + _non_child = ("schema", "subnodes", "fusible") subnodes: tuple[IR, ...] """Fused sub-nodes.""" - fused_io: Union | None - """Fused IO node.""" + fusible: bool + """Whether a parent may fuse to this node.""" def __init__( self, schema: Schema, subnodes: tuple[IR, ...], - fused_io: Union | None, + fusible: bool, # noqa: FBT001 *children: IR, ): self.schema = schema self.subnodes = subnodes - self.fused_io = fused_io + self.fusible = fusible self._non_child_args = ( [node.do_evaluate for node in subnodes], [tuple(node._non_child_args) for node in self.subnodes], - [], - [], ) self.children = children + def fuse_parent(self, parent: IR, *, fusible: bool = True) -> Self: + """Return a new IR node with a fused parent.""" + assert self.fusible, f"Fused node {self} is not fusible!" + return type(self)( + parent.schema, (*self.subnodes, parent), fusible, *self.children + ) + + @classmethod + def from_ir(cls, ir: IR, *, fusible: bool = True) -> Self: + "Construct from a single IR node." + return cls(ir.schema, (ir,), fusible, *ir.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=True): children = (func(*args, *children),) return children[0] +class FusedIO(IR): + """Fused-IO IR node.""" + + __slots__ = ("fused_io", "fusible", "subnodes") + _non_child = ("schema", "subnodes", "fused_io", "fusible") + subnodes: tuple[IR, ...] + """Fused sub-nodes.""" + fused_io: Union + """Fused IO node.""" + fusible: bool + """Whether a parent may fuse to this node.""" + + def __init__( + self, + schema: Schema, + subnodes: tuple[IR, ...], + fused_io: Union, + fusible: bool, # noqa: FBT001 + ): + self.schema = schema + self.subnodes = subnodes + self.fused_io = fused_io + self.fusible = fusible + self._non_child_args = ( + [node.do_evaluate for node in subnodes], + [tuple(node._non_child_args) for node in self.subnodes], + ) + self.children = () + + def fuse_parent(self, parent: IR, *, fusible: bool = True) -> Self: + """Return a new IR node with a fused parent.""" + assert self.fusible, f"FusedIO node {self} is not fusible!" + return type(self)( + parent.schema, + (*self.subnodes, parent), + self.fused_io, + fusible, + ) + + @classmethod + def from_ir(cls, ir: Union, *, fusible: bool = True) -> Self: + "Construct from a single IR node." + return cls(ir.schema, (), ir, fusible) + + @classmethod + def do_evaluate( + cls, + funcs: Sequence[Callable], + subargs: Sequence[Sequence[Any]], + io_cls: type[Scan | SplitScan], + io_args: Sequence[Any], + ) -> DataFrame: + """Evaluate and return a dataframe.""" + return Fused.do_evaluate(funcs, subargs, io_cls.do_evaluate(*io_args)) + + +def _maybe_make_fused( + ir: IR, + children: Sequence[IR], + partition_info: MutableMapping[IR, PartitionInfo], + *, + fusible: bool = True, +) -> IR: + """Return a Fused node if fusion is allowed.""" + new_node: IR + if ( + isinstance(ir, Union) + and all(isinstance(n, (Scan, SplitScan)) for n in ir.children) + and partition_info[ir].count == len(ir.children) + ): + # IO Fusion. + # Multi-partition IO is always implemented as a Union + # over Scan/SplitScan nodes. The number of children + # must correspond to the number of partitions. + new_node = FusedIO.from_ir(ir, fusible=fusible) + elif isinstance( + ir, + (Cache, Filter, GroupBy, HStack, MapFunction, Projection, Select), + ): + # Basic fusion. + # These are partition-wise operations if they exist + # in the graph after lowering (the data-movement + # has been encoded with Shuffle/Reduction/etc). + if isinstance(children[0], (Fused, FusedIO)) and children[0].fusible: + new_node = children[0].fuse_parent(ir, fusible=fusible) + else: + new_node = Fused.from_ir(ir, fusible=fusible).reconstruct(children) + else: + new_node = ir.reconstruct(children) + + return new_node + + def _fuse_ir_node( ir: IR, rec: LowerIRTransformer ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: @@ -88,67 +187,30 @@ def _fuse_ir_node( 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["fusible"]: - 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) - + new_node = _maybe_make_fused( + ir, + children, + rec.state["partition_info"], + # The new Fused node may only be fused by its + # parent if it corresponds to a unique sub-plan. + fusible=ir in rec.state["unique_subplans"], + ) partition_info[new_node] = rec.state["partition_info"][ir] return new_node, partition_info -def _is_fusible_ir_node( - ir: IR, partition_info: MutableMapping[IR, PartitionInfo] -) -> bool: - return ( - # Basic fusion. - # These are partition-wise operations if they exist - # in the graph after lowering (the data-movement - # has been encoded with Shuffle/Reduction/etc). - isinstance( - ir, - (Cache, Filter, GroupBy, HStack, MapFunction, Projection, Select), - ) - or ( - # IO Fusion. - # Multi-partition IO is always implemented as a Union - # over Scan/SplitScan nodes. The number of children - # must correspond to the number of partitions. - isinstance(ir, Union) - and all(isinstance(n, (Scan, SplitScan)) for n in ir.children) - and partition_info[ir].count == len(ir.children) - ) - ) - - -@generate_ir_tasks.register(Fused) +@generate_ir_tasks.register(FusedIO) def _( - ir: Fused, partition_info: MutableMapping[IR, PartitionInfo] + ir: FusedIO, 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) - 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], + *ir._non_child_args, + type(io), list(io._non_child_args), ) return graph @@ -175,13 +237,13 @@ def fuse_ir_graph( in the new graph to associated partitioning information. """ parents: defaultdict[IR, int] = defaultdict(int) - if _is_fusible_ir_node(ir, partition_info): - parents[ir] = 1 + # if _is_fusible_ir_node(ir, partition_info): + parents[ir] = 1 for node in traversal([ir]): for child in node.children: - if _is_fusible_ir_node(child, partition_info): - parents[child] += 1 - fusible = {node for node, count in parents.items() if count == 1} - state = {"fusible": fusible, "partition_info": partition_info} + # if _is_fusible_ir_node(child, partition_info): + parents[child] += 1 + unique_subplans = {node for node, count in parents.items() if count == 1} + state = {"unique_subplans": unique_subplans, "partition_info": partition_info} mapper = CachingVisitor(_fuse_ir_node, state=state) return mapper(ir) diff --git a/python/cudf_polars/cudf_polars/experimental/parallel.py b/python/cudf_polars/cudf_polars/experimental/parallel.py index ff296c867007..139e2ded74bc 100644 --- a/python/cudf_polars/cudf_polars/experimental/parallel.py +++ b/python/cudf_polars/cudf_polars/experimental/parallel.py @@ -227,7 +227,7 @@ def evaluate_streaming(ir: IR, config_options: ConfigOptions) -> DataFrame: @generate_ir_tasks.register(IR) -def _default_generate_ir_tasks( +def _( ir: IR, partition_info: MutableMapping[IR, PartitionInfo] ) -> MutableMapping[Any, Any]: # Generate pointwise (embarrassingly-parallel) tasks by default From 4fa356c7624bff40ca60f419beb25ef34d2accd9 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 15 May 2025 11:34:28 -0700 Subject: [PATCH 21/28] fix explain --- python/cudf_polars/cudf_polars/experimental/explain.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/explain.py b/python/cudf_polars/cudf_polars/experimental/explain.py index b275c8bcec5a..da4d3a41be1a 100644 --- a/python/cudf_polars/cudf_polars/experimental/explain.py +++ b/python/cudf_polars/cudf_polars/experimental/explain.py @@ -134,6 +134,11 @@ def _(ir: Scan, *, offset: str = "") -> str: @_repr_ir.register def _(ir: Fused, *, offset: str = "") -> str: nodes = tuple(type(node).__name__.upper() for node in ir.subnodes) - if isinstance(ir, FusedIO): - nodes = (type(ir.fused_io.children[0]).__name__.upper(), *nodes) return _repr_header(offset, f"FUSED {nodes}", ir.schema) + + +@_repr_ir.register +def _(ir: FusedIO, *, offset: str = "") -> str: + nodes = tuple(type(node).__name__.upper() for node in ir.subnodes) + nodes = (type(ir.fused_io.children[0]).__name__.upper(), *nodes) + return _repr_header(offset, f"FUSEDIO {nodes}", ir.schema) From b5384f714d1e381b450b6aaa321eb10643dd0464 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 15 May 2025 11:37:52 -0700 Subject: [PATCH 22/28] fix explain test --- python/cudf_polars/tests/experimental/test_explain.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf_polars/tests/experimental/test_explain.py b/python/cudf_polars/tests/experimental/test_explain.py index 0963297e76ad..fab7e95b869a 100644 --- a/python/cudf_polars/tests/experimental/test_explain.py +++ b/python/cudf_polars/tests/experimental/test_explain.py @@ -98,7 +98,7 @@ def test_explain_physical_plan_with_groupby(tmp_path, df, task_fusion): if task_fusion: plan = explain_query(q, engine, physical=True) assert "FUSED ('GROUPBY'" in plan - assert "FUSED ('SPLITSCAN'" in plan + assert "FUSEDIO ('SPLITSCAN'" in plan else: plan = explain_query(q, engine, physical=True) assert "GROUPBY ('g',)" in plan From 15f23327992b5ad1925b519e16319a7dda229c8c Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 15 May 2025 12:44:23 -0700 Subject: [PATCH 23/28] fix typing import --- python/cudf_polars/cudf_polars/experimental/fusion.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/experimental/fusion.py b/python/cudf_polars/cudf_polars/experimental/fusion.py index 857ea90056e4..9487f165c004 100644 --- a/python/cudf_polars/cudf_polars/experimental/fusion.py +++ b/python/cudf_polars/cudf_polars/experimental/fusion.py @@ -7,7 +7,7 @@ import operator from collections import defaultdict from functools import reduce -from typing import TYPE_CHECKING, Any, Self +from typing import TYPE_CHECKING from cudf_polars.dsl.ir import ( IR, @@ -28,6 +28,8 @@ from collections.abc import Callable, MutableMapping, Sequence from typing import Any + from typing_extensions import Self + from cudf_polars.containers import DataFrame from cudf_polars.experimental.base import PartitionInfo from cudf_polars.experimental.dispatch import LowerIRTransformer From 772f58419531c36f7423d94f9054e47c3de4da94 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 16 May 2025 06:40:14 -0700 Subject: [PATCH 24/28] add flag to pdsh.py --- .../cudf_polars/experimental/benchmarks/pdsh.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py index 95a5e96f4075..45eb9e40a539 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py @@ -145,6 +145,7 @@ class RunConfig: ) hardware: HardwareInfo = dataclasses.field(default_factory=HardwareInfo.collect) rapidsmpf_spill: bool + no_fuse: bool spill_device: float @classmethod @@ -170,6 +171,7 @@ def from_args(cls, args: argparse.Namespace) -> RunConfig: suffix=args.suffix, spill_device=args.spill_device, rapidsmpf_spill=args.rapidsmpf_spill, + no_fuse=args.no_fuse, ) def serialize(self) -> dict: @@ -195,6 +197,7 @@ def summarize(self) -> None: print(f"threads: {self.threads}") print(f"spill_device: {self.spill_device}") print(f"rapidsmpf_spill: {self.rapidsmpf_spill}") + print(f"task-fusion: {not self.no_fuse}") print(f"iterations: {self.iterations}") print("---------------------------------------") print(f"min time : {min([record.duration for record in records]):0.4f}") @@ -1120,6 +1123,12 @@ def _query_type(query: int | str) -> list[int]: help="Print an outline of the logical plan", default=False, ) +parser.add_argument( + "--no-fuse", + action=argparse.BooleanOptionalAction, + default=False, + help="Disable task fusion.", +) args = parser.parse_args() @@ -1187,6 +1196,8 @@ def run(args: argparse.Namespace) -> None: ) if run_config.rapidsmpf_spill: executor_options["rapidsmpf_spill"] = run_config.rapidsmpf_spill + if not run_config.no_fuse: + executor_options["task_fusion"] = False if run_config.scheduler == "distributed": executor_options["scheduler"] = "distributed" From b68a241729a9264ecc0610feeaf21eed426e5833 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 16 May 2025 06:46:04 -0700 Subject: [PATCH 25/28] fix obvious mistake --- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py index 45eb9e40a539..422685be5284 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py @@ -1196,7 +1196,7 @@ def run(args: argparse.Namespace) -> None: ) if run_config.rapidsmpf_spill: executor_options["rapidsmpf_spill"] = run_config.rapidsmpf_spill - if not run_config.no_fuse: + if run_config.no_fuse: executor_options["task_fusion"] = False if run_config.scheduler == "distributed": executor_options["scheduler"] = "distributed" From d21da69e7f76b0f6884b04f6cb943d0f0b2c8736 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 16 May 2025 07:05:58 -0700 Subject: [PATCH 26/28] fix misunderstanding --- .../cudf_polars/experimental/benchmarks/pdsh.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py index 422685be5284..82770c213ff5 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py @@ -145,7 +145,7 @@ class RunConfig: ) hardware: HardwareInfo = dataclasses.field(default_factory=HardwareInfo.collect) rapidsmpf_spill: bool - no_fuse: bool + task_fusion: bool spill_device: float @classmethod @@ -171,7 +171,7 @@ def from_args(cls, args: argparse.Namespace) -> RunConfig: suffix=args.suffix, spill_device=args.spill_device, rapidsmpf_spill=args.rapidsmpf_spill, - no_fuse=args.no_fuse, + task_fusion=args.task_fusion, ) def serialize(self) -> dict: @@ -197,7 +197,7 @@ def summarize(self) -> None: print(f"threads: {self.threads}") print(f"spill_device: {self.spill_device}") print(f"rapidsmpf_spill: {self.rapidsmpf_spill}") - print(f"task-fusion: {not self.no_fuse}") + print(f"task-fusion: {not self.task_fusion}") print(f"iterations: {self.iterations}") print("---------------------------------------") print(f"min time : {min([record.duration for record in records]):0.4f}") @@ -1124,10 +1124,10 @@ def _query_type(query: int | str) -> list[int]: default=False, ) parser.add_argument( - "--no-fuse", + "--task-fusion", action=argparse.BooleanOptionalAction, - default=False, - help="Disable task fusion.", + default=True, + help="Enable task fusion.", ) args = parser.parse_args() @@ -1196,10 +1196,9 @@ def run(args: argparse.Namespace) -> None: ) if run_config.rapidsmpf_spill: executor_options["rapidsmpf_spill"] = run_config.rapidsmpf_spill - if run_config.no_fuse: - executor_options["task_fusion"] = False if run_config.scheduler == "distributed": executor_options["scheduler"] = "distributed" + executor_options["task_fusion"] = run_config.task_fusion engine = pl.GPUEngine( raise_on_fail=True, From 4da62d4e5cdb0caadba272f4dc46d9d6c20df48c Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 28 May 2025 09:46:39 -0700 Subject: [PATCH 27/28] fix pdsh.py --- python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py index 49c7787c6e44..bdb4055ee6cc 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py @@ -207,7 +207,7 @@ def summarize(self) -> None: print(f"blocksize: {self.blocksize}") print(f"shuffle_method: {self.shuffle}") print(f"broadcast_join_limit: {self.broadcast_join_limit}") - print(f"task-fusion: {not self.task_fusion}") + print(f"task-fusion: {self.task_fusion}") if self.scheduler == "distributed": print(f"n_workers: {self.n_workers}") print(f"threads: {self.threads}") From f60282fde003234e9e38f3e7111832bf3af829f0 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 30 May 2025 13:35:36 -0700 Subject: [PATCH 28/28] address minor code-review comments --- python/cudf_polars/cudf_polars/experimental/fusion.py | 2 -- python/cudf_polars/cudf_polars/experimental/io.py | 2 +- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/experimental/fusion.py b/python/cudf_polars/cudf_polars/experimental/fusion.py index 9487f165c004..dad4a2b2dc71 100644 --- a/python/cudf_polars/cudf_polars/experimental/fusion.py +++ b/python/cudf_polars/cudf_polars/experimental/fusion.py @@ -239,11 +239,9 @@ def fuse_ir_graph( in the new graph to associated partitioning information. """ parents: defaultdict[IR, int] = defaultdict(int) - # if _is_fusible_ir_node(ir, partition_info): parents[ir] = 1 for node in traversal([ir]): for child in node.children: - # if _is_fusible_ir_node(child, partition_info): parents[child] += 1 unique_subplans = {node for node, count in parents.items() if count == 1} state = {"unique_subplans": unique_subplans, "partition_info": partition_info} diff --git a/python/cudf_polars/cudf_polars/experimental/io.py b/python/cudf_polars/cudf_polars/experimental/io.py index 256681aee62d..13242512be43 100644 --- a/python/cudf_polars/cudf_polars/experimental/io.py +++ b/python/cudf_polars/cudf_polars/experimental/io.py @@ -262,7 +262,7 @@ 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 = 5 # TODO: Make this configurable + n_sample = 3 # TODO: Make this configurable stride = max(1, int(len(ir.paths) / n_sample)) metadata = plc.io.parquet_metadata.read_parquet_metadata( plc.io.SourceInfo(ir.paths[: stride * n_sample : stride])