From 7d0a0e55cdd7a6dec442a690c150a5fa81368fca Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 18 Aug 2026 12:59:27 -0700 Subject: [PATCH 01/12] add initial partitioning_hints infrastructure --- .../cudf_polars/streaming/actor_graph/core.py | 5 + .../streaming/actor_graph/dispatch.py | 6 +- .../streaming/partitioning_hints.py | 329 ++++++++++++++++++ .../streaming/test_partitioning_hints.py | 255 ++++++++++++++ 4 files changed, 594 insertions(+), 1 deletion(-) create mode 100644 python/cudf_polars/cudf_polars/streaming/partitioning_hints.py create mode 100644 python/cudf_polars/tests/streaming/test_partitioning_hints.py diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py index fbec424d3104..603d7891d113 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py @@ -23,6 +23,7 @@ metadata_drain_node, ) from cudf_polars.streaming.over import Over +from cudf_polars.streaming.partitioning_hints import collect_partitioning_hints from cudf_polars.utils.config import SPMDContext if TYPE_CHECKING: @@ -256,6 +257,9 @@ def generate_network( # Determine which nodes need fanout fanout_nodes = determine_fanout_nodes(ir, partition_info, ir_dep_count) + partitioning_hints = collect_partitioning_hints(ir, partition_info) + # import pdb; pdb.set_trace() + # pass # Generate the network state: GenState = { @@ -264,6 +268,7 @@ def generate_network( "config_options": config_options, "partition_info": partition_info, "fanout_nodes": fanout_nodes, + "partitioning_hints": partitioning_hints, "ir_context": ir_context, "max_concurrent_io_tasks": config_options.executor.max_concurrent_io_tasks, "stats": stats, diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py index 88d6784b0940..0dcdbe58673c 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py @@ -10,7 +10,7 @@ from cudf_polars.typing import GenericTransformer if TYPE_CHECKING: - from collections.abc import MutableMapping + from collections.abc import Mapping, MutableMapping from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.core.context import Context @@ -21,6 +21,7 @@ PartitionInfo, StatsCollector, ) + from cudf_polars.streaming.partitioning_hints import PartitioningHint from cudf_polars.utils.config import ConfigOptions, StreamingExecutor @@ -49,6 +50,8 @@ class GenState(TypedDict): Partition information. fanout_nodes Dictionary mapping IR nodes to fanout information. + partitioning_hints + Physical-layout hints for IR outputs. ir_context The execution context for the IR node. max_concurrent_io_tasks @@ -64,6 +67,7 @@ class GenState(TypedDict): config_options: ConfigOptions[StreamingExecutor] partition_info: MutableMapping[IR, PartitionInfo] fanout_nodes: dict[IR, FanoutInfo] + partitioning_hints: Mapping[IR, PartitioningHint] ir_context: IRExecutionContext max_concurrent_io_tasks: int stats: StatsCollector diff --git a/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py b/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py new file mode 100644 index 000000000000..f3d50823a090 --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py @@ -0,0 +1,329 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Partitioning hints for streaming actor-graph construction.""" + +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING, TypeAlias + +import pylibcudf as plc + +from cudf_polars.dsl import expr +from cudf_polars.dsl.ir import ( + Filter, + GroupBy, + Join, + MapFunction, + Projection, + Select, + Slice, + Sort, +) +from cudf_polars.dsl.traversal import post_traversal +from cudf_polars.dsl.utils.column_domain import column_domain_bindings +from cudf_polars.streaming.repartition import Repartition + +if TYPE_CHECKING: + from collections.abc import Mapping + + from cudf_polars.dsl.ir import IR + from cudf_polars.streaming.base import PartitionInfo + + +@dataclasses.dataclass(frozen=True) +class NamedOrderKey: + """Named sort key with logical Polars ordering options.""" + + name: str + descending: bool + nulls_last: bool + + +@dataclasses.dataclass(frozen=True) +class HashPartitioningHint: + """ + Hint that rows should be co-located by equality keys. + + ``peer_partition_count`` is a planning-time estimate of the partition + count that may align with a downstream peer input. + """ + + keys: tuple[str, ...] + peer_partition_count: int | None = None + + +@dataclasses.dataclass(frozen=True) +class OrderPartitioningHint: + """ + Hint that rows should be ordered by a key sequence. + + ``peer_partition_count`` is a planning-time estimate of the partition + count that may align with a downstream peer input. + """ + + keys: tuple[NamedOrderKey, ...] + peer_partition_count: int | None = None + + +PartitioningHint: TypeAlias = HashPartitioningHint | OrderPartitioningHint +_MaybeHint: TypeAlias = PartitioningHint | None + +__all__ = [ + "HashPartitioningHint", + "NamedOrderKey", + "OrderPartitioningHint", + "PartitioningHint", + "collect_partitioning_hints", +] + + +def collect_partitioning_hints( + ir: IR, + partition_info: Mapping[IR, PartitionInfo], +) -> dict[IR, PartitioningHint]: + """ + Collect physical-layout hints for each IR node. + + Parameters + ---------- + ir + The IR to determine the preferred partitioning for. + partition_info : Mapping[IR, PartitionInfo] + The partition information for each IR node. + + Returns + ------- + The physical-layout hints for each IR node. + + Notes + ----- + The returned hints describe layouts that downstream consumers would prefer, + if practical. These hints are not a strict runtime contract. + """ + hints: dict[IR, _MaybeHint] = {} + for node in reversed(list(post_traversal([ir]))): + for child, hint in _direct_child_hints(node, partition_info): + _record_hint(hints, child, hint) + + if (current_hint := hints.get(node)) is not None: + for child, child_hint in _propagate_hint(node, current_hint): + _record_hint(hints, child, child_hint) + + return {node: hint for node, hint in hints.items() if hint is not None} + + +def _direct_child_hints( + ir: IR, + partition_info: Mapping[IR, PartitionInfo], +) -> tuple[tuple[IR, PartitioningHint], ...]: + """Return hints created directly by *ir* for its children.""" + if isinstance(ir, Sort): + hint = _sort_hint(ir) + return () if hint is None else ((ir.children[0], hint),) + + if isinstance(ir, Join) and ir.options[0] != "Cross": + left_keys = _column_names(ir.left_on) + right_keys = _column_names(ir.right_on) + if left_keys is None or right_keys is None: + return () + peer_partition_count = max( + partition_info[ir.children[0]].count, + partition_info[ir.children[1]].count, + ) + return ( + (ir.children[0], HashPartitioningHint(left_keys, peer_partition_count)), + (ir.children[1], HashPartitioningHint(right_keys, peer_partition_count)), + ) + + if isinstance(ir, GroupBy) and not ir.maintain_order: + keys = _column_names(ir.keys) + return () if keys is None else ((ir.children[0], HashPartitioningHint(keys)),) + + return () + + +def _propagate_hint( + ir: IR, hint: PartitioningHint +) -> tuple[tuple[IR, PartitioningHint], ...]: + """Propagate *hint* through simple row-order-preserving IR nodes.""" + if len(ir.children) != 1: + return () + + (child,) = ir.children + if isinstance(ir, (Filter, GroupBy, Projection, Select, Slice)): + remapped = _remap_hint(hint, _child_zero_remapping(ir)) + elif isinstance(ir, Repartition) or ( + isinstance(ir, MapFunction) and ir.name in {"hint_sorted", "rechunk"} + ): + remapped = _remap_hint(hint, _identity_remapping(ir)) + elif isinstance(ir, MapFunction) and ir.name == "row_index": + remapped = _remap_hint(hint, _identity_remapping(child)) + else: + remapped = None + + return () if remapped is None else ((child, remapped),) + + +def _sort_hint(ir: Sort) -> OrderPartitioningHint | None: + names = _column_names(ir.by) + if names is None: + return None + return OrderPartitioningHint( + tuple( + NamedOrderKey( + name, + order == plc.types.Order.DESCENDING, + _nulls_last(order, null_order), + ) + for name, order, null_order in zip( + names, ir.order, ir.null_order, strict=True + ) + ) + ) + + +def _nulls_last(order: plc.types.Order, null_order: plc.types.NullOrder) -> bool: + """Return the Polars ``nulls_last`` value represented by cudf sort options.""" + return (order == plc.types.Order.ASCENDING) == ( + null_order == plc.types.NullOrder.AFTER + ) + + +def _column_names(named_exprs: tuple[expr.NamedExpr, ...]) -> tuple[str, ...] | None: + names = [] + for named_expr in named_exprs: + if not isinstance(named_expr.value, expr.Col): + return None + names.append(named_expr.value.name) + return tuple(names) + + +def _remap_hint( + hint: PartitioningHint, remapping: Mapping[str, str] +) -> PartitioningHint | None: + if isinstance(hint, HashPartitioningHint): + names = _remap_names(hint.keys, remapping) + return None if names is None else dataclasses.replace(hint, keys=names) + + keys = [] + for key in hint.keys: + if (name := remapping.get(key.name)) is None: + return None + keys.append(dataclasses.replace(key, name=name)) + return dataclasses.replace(hint, keys=tuple(keys)) + + +def _child_zero_remapping(ir: IR) -> dict[str, str]: + return { + output_name: binding.name + for output_name, binding in column_domain_bindings(ir).items() + if binding.child_index == 0 + } + + +def _identity_remapping(ir: IR) -> dict[str, str]: + return {name: name for name in ir.schema} + + +def _remap_names( + names: tuple[str, ...], remapping: Mapping[str, str] +) -> tuple[str, ...] | None: + remapped = [] + for name in names: + if (new_name := remapping.get(name)) is None: + return None + remapped.append(new_name) + return tuple(remapped) + + +def _record_hint( + hints: dict[IR, _MaybeHint], + node: IR, + hint: PartitioningHint, +) -> None: + if node not in hints: + hints[node] = hint + elif (current := hints[node]) is not None: + hints[node] = _merge_hints(current, hint) + + +def _merge_hints( + left: PartitioningHint, right: PartitioningHint +) -> PartitioningHint | None: + peer_partition_count = _merge_peer_partition_count(left, right) + if isinstance(left, HashPartitioningHint) and isinstance( + right, HashPartitioningHint + ): + hash_keys = _shortest_common_prefix(left.keys, right.keys) + return ( + None + if hash_keys is None + else HashPartitioningHint( + hash_keys, peer_partition_count=peer_partition_count + ) + ) + + if isinstance(left, OrderPartitioningHint) and isinstance( + right, OrderPartitioningHint + ): + order_keys = _longest_common_extension(left.keys, right.keys) + return ( + None + if order_keys is None + else OrderPartitioningHint( + order_keys, peer_partition_count=peer_partition_count + ) + ) + + if isinstance(left, OrderPartitioningHint): + order_hint = left + assert isinstance(right, HashPartitioningHint) + hash_hint = right + else: + assert isinstance(right, OrderPartitioningHint) + order_hint = right + hash_hint = left + order_names = tuple(key.name for key in order_hint.keys) + if _is_prefix(hash_hint.keys, order_names) or _is_prefix( + order_names, hash_hint.keys + ): + return dataclasses.replace( + order_hint, peer_partition_count=peer_partition_count + ) + return None + + +def _merge_peer_partition_count( + left: PartitioningHint, right: PartitioningHint +) -> int | None: + counts = [ + count + for count in (left.peer_partition_count, right.peer_partition_count) + if count is not None + ] + return max(counts) if counts else None + + +def _shortest_common_prefix( + left: tuple[str, ...], right: tuple[str, ...] +) -> tuple[str, ...] | None: + if _is_prefix(left, right): + return left + if _is_prefix(right, left): + return right + return None + + +def _longest_common_extension( + left: tuple[NamedOrderKey, ...], right: tuple[NamedOrderKey, ...] +) -> tuple[NamedOrderKey, ...] | None: + if _is_prefix(left, right): + return right + if _is_prefix(right, left): + return left + return None + + +def _is_prefix(left: tuple[object, ...], right: tuple[object, ...]) -> bool: + return len(left) <= len(right) and left == right[: len(left)] diff --git a/python/cudf_polars/tests/streaming/test_partitioning_hints.py b/python/cudf_polars/tests/streaming/test_partitioning_hints.py new file mode 100644 index 000000000000..6ae51619b0fa --- /dev/null +++ b/python/cudf_polars/tests/streaming/test_partitioning_hints.py @@ -0,0 +1,255 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import polars as pl + +from cudf_polars.containers import DataType +from cudf_polars.dsl import expr +from cudf_polars.dsl.ir import DataFrameScan, GroupBy, Join, Select, Sort, Union +from cudf_polars.streaming.base import PartitionInfo +from cudf_polars.streaming.partitioning_hints import ( + HashPartitioningHint, + NamedOrderKey, + OrderPartitioningHint, + collect_partitioning_hints, +) +from cudf_polars.utils.sorting import sort_order + +if TYPE_CHECKING: + from cudf_polars.dsl.ir import IR + +I64 = DataType(pl.Int64()) + + +def make_scan(*names: str) -> DataFrameScan: + frame = pl.DataFrame({name: [1] for name in names}) + return DataFrameScan(dict.fromkeys(names, I64), frame._df, None) + + +def named_col(name: str) -> expr.NamedExpr: + return expr.NamedExpr(name, expr.Col(I64, name)) + + +def make_sort( + child: IR, + *names: str, + descending: tuple[bool, ...] | None = None, + nulls_last: tuple[bool, ...] | None = None, +) -> Sort: + if descending is None: + descending = (False,) * len(names) + if nulls_last is None: + nulls_last = (False,) * len(names) + order, null_order = sort_order( + descending, nulls_last=nulls_last, num_keys=len(names) + ) + return Sort( + child.schema, + tuple(named_col(name) for name in names), + order, + null_order, + stable=False, + zlice=None, + df=child, + ) + + +def test_sort_creates_order_partition_hint() -> None: + scan = make_scan("a", "b") + sort = make_sort( + scan, + "a", + "b", + descending=(False, True), + nulls_last=(True, False), + ) + + hints = collect_partitioning_hints( + sort, + {sort: PartitionInfo(2), scan: PartitionInfo(2)}, + ) + + assert hints[scan] == OrderPartitioningHint( + ( + NamedOrderKey("a", descending=False, nulls_last=True), + NamedOrderKey("b", descending=True, nulls_last=False), + ) + ) + + +def test_join_creates_hash_partitioning_hints() -> None: + left = make_scan("k", "left_value") + right = make_scan("k", "right_value") + join = Join( + {"k": I64, "left_value": I64, "right_value": I64}, + (named_col("k"),), + (named_col("k"),), + ("Inner", False, None, "_right", True, "none"), + left, + right, + ) + + hints = collect_partitioning_hints( + join, + { + join: PartitionInfo(7), + left: PartitionInfo(3), + right: PartitionInfo(7), + }, + ) + + assert hints[left] == HashPartitioningHint(("k",), peer_partition_count=7) + assert hints[right] == HashPartitioningHint(("k",), peer_partition_count=7) + + +def test_groupby_creates_hash_partition_hint() -> None: + scan = make_scan("a", "b") + groupby = GroupBy( + {"a": I64}, + (named_col("a"),), + (), + maintain_order=False, + zlice=None, + df=scan, + ) + + hints = collect_partitioning_hints( + groupby, + {groupby: PartitionInfo(2), scan: PartitionInfo(2)}, + ) + + assert hints[scan] == HashPartitioningHint(("a",)) + + +def test_select_remaps_order_partition_hint() -> None: + scan = make_scan("a", "b") + select = Select( + {"x": I64, "b": I64}, + (expr.NamedExpr("x", expr.Col(I64, "a")), named_col("b")), + should_broadcast=False, + df=scan, + ) + sort = make_sort(select, "x") + + hints = collect_partitioning_hints( + sort, + { + sort: PartitionInfo(2), + select: PartitionInfo(2), + scan: PartitionInfo(2), + }, + ) + + assert hints[scan] == OrderPartitioningHint( + (NamedOrderKey("a", descending=False, nulls_last=False),) + ) + + +def test_groupby_remaps_order_partition_hint() -> None: + scan = make_scan("a", "b") + groupby = GroupBy( + {"key": I64}, + (expr.NamedExpr("key", expr.Col(I64, "a")),), + (), + maintain_order=False, + zlice=None, + df=scan, + ) + sort = make_sort(groupby, "key") + + hints = collect_partitioning_hints( + sort, + { + sort: PartitionInfo(2), + groupby: PartitionInfo(2), + scan: PartitionInfo(2), + }, + ) + + assert hints[scan] == OrderPartitioningHint( + (NamedOrderKey("a", descending=False, nulls_last=False),) + ) + + +def test_fanout_keeps_more_specific_compatible_order_hint() -> None: + scan = make_scan("a", "b") + root = Union( + scan.schema, + None, + False, # noqa: FBT003 + make_sort(scan, "a"), + make_sort(scan, "a", "b"), + ) + + hints = collect_partitioning_hints( + root, + {root: PartitionInfo(2), scan: PartitionInfo(2)}, + ) + + assert hints[scan] == OrderPartitioningHint( + ( + NamedOrderKey("a", descending=False, nulls_last=False), + NamedOrderKey("b", descending=False, nulls_last=False), + ) + ) + + +def test_fanout_merges_compatible_hash_hint_into_order_hint() -> None: + scan = make_scan("a", "b") + right = make_scan("a", "right_value") + join = Join( + {"a": I64, "b": I64, "right_value": I64}, + (named_col("a"),), + (named_col("a"),), + ("Inner", False, None, "_right", True, "none"), + scan, + right, + ) + root = Union( + scan.schema, + None, + False, # noqa: FBT003 + make_sort(scan, "a", "b"), + join, + ) + + hints = collect_partitioning_hints( + root, + { + root: PartitionInfo(5), + join: PartitionInfo(5), + scan: PartitionInfo(3), + right: PartitionInfo(5), + }, + ) + + assert hints[scan] == OrderPartitioningHint( + ( + NamedOrderKey("a", descending=False, nulls_last=False), + NamedOrderKey("b", descending=False, nulls_last=False), + ), + peer_partition_count=5, + ) + assert hints[right] == HashPartitioningHint(("a",), peer_partition_count=5) + + +def test_conflicting_fanout_drops_hint() -> None: + scan = make_scan("a", "b") + root = Union( + scan.schema, + None, + False, # noqa: FBT003 + make_sort(scan, "a"), + make_sort(scan, "b"), + ) + + hints = collect_partitioning_hints( + root, + {root: PartitionInfo(2), scan: PartitionInfo(2)}, + ) + + assert scan not in hints From c675252f386ae2518cc1abec2b1d5dcf7ce4ec32 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 18 Aug 2026 13:24:36 -0700 Subject: [PATCH 02/12] cleanup --- .../streaming/partitioning_hints.py | 57 +++++++------------ 1 file changed, 22 insertions(+), 35 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py b/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py index f3d50823a090..4bf08339e7e3 100644 --- a/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py +++ b/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py @@ -88,7 +88,7 @@ def collect_partitioning_hints( Parameters ---------- ir - The IR to determine the preferred partitioning for. + The IR to collect partitioning hints for. partition_info : Mapping[IR, PartitionInfo] The partition information for each IR node. @@ -146,19 +146,22 @@ def _direct_child_hints( def _propagate_hint( ir: IR, hint: PartitioningHint ) -> tuple[tuple[IR, PartitioningHint], ...]: - """Propagate *hint* through simple row-order-preserving IR nodes.""" + """Propagate *hint* through nodes whose children may benefit from it.""" if len(ir.children) != 1: return () (child,) = ir.children - if isinstance(ir, (Filter, GroupBy, Projection, Select, Slice)): + if isinstance(ir, (Filter, Projection, Select, Slice)): + remapped = _remap_hint(hint, _child_zero_remapping(ir)) + elif isinstance(ir, GroupBy): + # Only group-key outputs remap to the child; aggregate outputs stop here. remapped = _remap_hint(hint, _child_zero_remapping(ir)) elif isinstance(ir, Repartition) or ( isinstance(ir, MapFunction) and ir.name in {"hint_sorted", "rechunk"} ): - remapped = _remap_hint(hint, _identity_remapping(ir)) + remapped = _remap_hint(hint, {name: name for name in ir.schema}) elif isinstance(ir, MapFunction) and ir.name == "row_index": - remapped = _remap_hint(hint, _identity_remapping(child)) + remapped = _remap_hint(hint, {name: name for name in child.schema}) else: remapped = None @@ -174,7 +177,8 @@ def _sort_hint(ir: Sort) -> OrderPartitioningHint | None: NamedOrderKey( name, order == plc.types.Order.DESCENDING, - _nulls_last(order, null_order), + (order == plc.types.Order.ASCENDING) + == (null_order == plc.types.NullOrder.AFTER), ) for name, order, null_order in zip( names, ir.order, ir.null_order, strict=True @@ -183,13 +187,6 @@ def _sort_hint(ir: Sort) -> OrderPartitioningHint | None: ) -def _nulls_last(order: plc.types.Order, null_order: plc.types.NullOrder) -> bool: - """Return the Polars ``nulls_last`` value represented by cudf sort options.""" - return (order == plc.types.Order.ASCENDING) == ( - null_order == plc.types.NullOrder.AFTER - ) - - def _column_names(named_exprs: tuple[expr.NamedExpr, ...]) -> tuple[str, ...] | None: names = [] for named_expr in named_exprs: @@ -203,15 +200,20 @@ def _remap_hint( hint: PartitioningHint, remapping: Mapping[str, str] ) -> PartitioningHint | None: if isinstance(hint, HashPartitioningHint): - names = _remap_names(hint.keys, remapping) - return None if names is None else dataclasses.replace(hint, keys=names) - - keys = [] + remapped_names = [] + for name in hint.keys: + if (new_name := remapping.get(name)) is None: + return None + remapped_names.append(new_name) + return dataclasses.replace(hint, keys=tuple(remapped_names)) + + remapped_keys = [] for key in hint.keys: - if (name := remapping.get(key.name)) is None: + new_name = remapping.get(key.name) + if new_name is None: return None - keys.append(dataclasses.replace(key, name=name)) - return dataclasses.replace(hint, keys=tuple(keys)) + remapped_keys.append(dataclasses.replace(key, name=new_name)) + return dataclasses.replace(hint, keys=tuple(remapped_keys)) def _child_zero_remapping(ir: IR) -> dict[str, str]: @@ -222,21 +224,6 @@ def _child_zero_remapping(ir: IR) -> dict[str, str]: } -def _identity_remapping(ir: IR) -> dict[str, str]: - return {name: name for name in ir.schema} - - -def _remap_names( - names: tuple[str, ...], remapping: Mapping[str, str] -) -> tuple[str, ...] | None: - remapped = [] - for name in names: - if (new_name := remapping.get(name)) is None: - return None - remapped.append(new_name) - return tuple(remapped) - - def _record_hint( hints: dict[IR, _MaybeHint], node: IR, From 1c26f03904018a112af57f8b7c32ebd074fe33e4 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 18 Aug 2026 14:39:38 -0700 Subject: [PATCH 03/12] remove leftover comment --- python/cudf_polars/cudf_polars/streaming/actor_graph/core.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py index 603d7891d113..5fe73807bc7b 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py @@ -258,8 +258,6 @@ def generate_network( # Determine which nodes need fanout fanout_nodes = determine_fanout_nodes(ir, partition_info, ir_dep_count) partitioning_hints = collect_partitioning_hints(ir, partition_info) - # import pdb; pdb.set_trace() - # pass # Generate the network state: GenState = { From b0ea471de2f9b5f06d00dab3713572d203d3a201 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 19 Aug 2026 09:04:07 -0700 Subject: [PATCH 04/12] save work - mostly unsuccessful experiment --- .../cudf_polars/streaming/actor_graph/core.py | 11 +- .../streaming/actor_graph/groupby.py | 32 +++++ .../cudf_polars/streaming/actor_graph/join.py | 24 +++- .../streaming/partitioning_hints.py | 46 ++++++- .../cudf_polars/cudf_polars/utils/config.py | 69 +++++++++++ .../tests/streaming/test_groupby.py | 73 +++++++++++ .../streaming/test_partitioning_hints.py | 71 ++++++++++- python/cudf_polars/tests/test_config.py | 115 ++++++++++++++++++ 8 files changed, 433 insertions(+), 8 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py index 5fe73807bc7b..67ecb68bfd6d 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py @@ -257,7 +257,16 @@ def generate_network( # Determine which nodes need fanout fanout_nodes = determine_fanout_nodes(ir, partition_info, ir_dep_count) - partitioning_hints = collect_partitioning_hints(ir, partition_info) + dynamic_planning = config_options.executor.dynamic_planning + hint_options = ( + None if dynamic_planning is None else dynamic_planning.partitioning_hints + ) + partitioning_hints = ( + collect_partitioning_hints(ir, partition_info) + if hint_options is not None + and (hint_options.use_partition_counts or hint_options.use_ordering) + else {} + ) # Generate the network state: GenState = { diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py index e1c93620d1f0..bfb1a2016f9b 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py @@ -45,6 +45,7 @@ shutdown_on_error, ) from cudf_polars.streaming.groupby import _has_stable_sorted_agg, combine, decompose +from cudf_polars.streaming.partitioning_hints import apply_peer_partition_count_hint from cudf_polars.streaming.repartition import Repartition if TYPE_CHECKING: @@ -55,6 +56,7 @@ from cudf_polars.dsl.ir import IRExecutionContext from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator from cudf_polars.streaming.actor_graph.tracing import ActorTracer + from cudf_polars.streaming.partitioning_hints import PartitioningHint from cudf_polars.typing import Schema @@ -515,8 +517,10 @@ async def _choose_strategy( input_drained: bool, # noqa: FBT001 collective_ids: list[int], target_partition_size: int, + broadcast_limit: int, skip_global_comm: bool, # noqa: FBT001 maintain_order: bool, # noqa: FBT001 + partitioning_hint: PartitioningHint | None, tracer: ActorTracer | None, ) -> int: """ @@ -540,10 +544,14 @@ async def _choose_strategy( The collective IDs. target_partition_size The target partition size. + broadcast_limit + The maximum size for broadcast join inputs. skip_global_comm Whether to skip the global communication. maintain_order Whether the operation should maintain input ordering semantics. + partitioning_hint + Optional downstream physical-layout hint. tracer Optional tracer for runtime metrics. @@ -602,6 +610,14 @@ async def _choose_strategy( output_count = min(ideal_count, output_count_limit) if not use_tree: output_count = max(2, output_count, min_row_limit_count) + broadcastable_output = ( + total_estimated_size < broadcast_limit + and total_estimated_rows < MAX_ROWS_PER_PARTITION + ) + if not broadcastable_output: + output_count = apply_peer_partition_count_hint( + output_count, partitioning_hint + ) if tracer is not None: tracer.decision = ( "tree_local" @@ -625,7 +641,9 @@ async def groupby_actor( ch_out: Channel[TableChunk], ch_in: Channel[TableChunk], target_partition_size: int, + broadcast_limit: int, collective_ids: list[int], + partitioning_hint: PartitioningHint | None, ) -> None: """ Dynamic GroupBy or Distinct actor that selects the best strategy at runtime. @@ -651,8 +669,12 @@ async def groupby_actor( The input channel. target_partition_size The target partition size. + broadcast_limit + The maximum size for broadcast join inputs. collective_ids The collective IDs. + partitioning_hint + Optional downstream physical-layout hint. """ async with shutdown_on_error( context, ch_in, ch_out, trace_ir=ir, ir_context=ir_context @@ -723,8 +745,10 @@ async def groupby_actor( input_drained, collective_ids, target_partition_size, + broadcast_limit, skip_global_comm, maintain_order, + partitioning_hint, tracer, ) @@ -779,6 +803,12 @@ def _( assert len(collective_ids) == 2, ( f"{type(ir).__name__} requires 2 collective IDs, got {len(collective_ids)}" ) + hint_options = config_options.executor.dynamic_planning.partitioning_hints + partitioning_hint = ( + rec.state["partitioning_hints"].get(ir) + if hint_options is not None and hint_options.use_partition_counts + else None + ) actors[ir] = [ groupby_actor( rec.state["context"], @@ -788,7 +818,9 @@ def _( channels[ir].reserve_input_slot(), channels[ir.children[0]].reserve_output_slot(), config_options.executor.target_partition_size, + config_options.executor.broadcast_limit, collective_ids, + partitioning_hint, ) ] diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index d944847e621b..a21bef35515a 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -56,6 +56,7 @@ send_metadata, shutdown_on_error, ) +from cudf_polars.streaming.partitioning_hints import apply_peer_partition_count_hint from cudf_polars.streaming.repartition import Repartition from cudf_polars.streaming.utils import _concat @@ -71,6 +72,7 @@ from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator from cudf_polars.streaming.actor_graph.tracing import ActorTracer from cudf_polars.streaming.base import PartitionInfo + from cudf_polars.streaming.partitioning_hints import PartitioningHint from cudf_polars.utils.config import StreamingExecutor @@ -714,6 +716,7 @@ def _choose_strategy_from_samples( left_sample: TableSizeStats, right_sample: TableSizeStats, chunkwise: bool, + partitioning_hint: PartitioningHint | None, tracer: ActorTracer | None, ) -> JoinStrategy: """Choose potential broadcast side and minimum shuffle modulus.""" @@ -795,6 +798,9 @@ def _choose_strategy_from_samples( estimated_rows_count + MAX_ROWS_PER_PARTITION - 1 ) // MAX_ROWS_PER_PARTITION min_shuffle_modulus = max(min_shuffle_modulus, min_partitions_for_row_limit) + min_shuffle_modulus = apply_peer_partition_count_hint( + min_shuffle_modulus, partitioning_hint + ) shuffle_modulus = _choose_shuffle_modulus( comm, @@ -859,6 +865,7 @@ async def _choose_strategy( right_metadata: ChannelMetadata, executor: StreamingExecutor, collective_ids: list[int], + partitioning_hint: PartitioningHint | None, *, tracer: ActorTracer | None, ) -> tuple[TableSizeStats, TableSizeStats, JoinStrategy]: @@ -878,9 +885,10 @@ async def _choose_strategy( hash_chunkwise = isinstance( left_partitioning.inter_rank_scheme, HashScheme ) and isinstance(right_partitioning.inter_rank_scheme, HashScheme) - if hash_chunkwise and left_partitioning.is_aligned_with( + aligned = hash_chunkwise and left_partitioning.is_aligned_with( right_partitioning, context.br() - ): + ) + if aligned: # We can use a chunkwise join chunkwise = True left_sample = TableSizeStats( @@ -932,6 +940,7 @@ async def _choose_strategy( left_sample=left_sample, right_sample=right_sample, chunkwise=chunkwise, + partitioning_hint=partitioning_hint, tracer=tracer, ) @@ -949,6 +958,7 @@ async def join_actor( ch_right: Channel[TableChunk], executor: StreamingExecutor, collective_ids: list[int], + partitioning_hint: PartitioningHint | None, ) -> None: """ Dynamic Join actor that selects the best strategy at runtime. @@ -977,6 +987,8 @@ async def join_actor( Streaming executor configuration. collective_ids List of collective IDs for shuffle/broadcast; consumed as needed. + partitioning_hint + Optional downstream physical-layout hint. """ async with shutdown_on_error( context, @@ -1001,6 +1013,7 @@ async def join_actor( right_metadata, executor, collective_ids, + partitioning_hint, tracer=tracer, ) ch_left_replay = context.create_channel() @@ -1150,6 +1163,12 @@ def _( f"{len(collective_ids)} for this Join. " "Ensure ReserveOpIDs is run with dynamic_planning enabled." ) + hint_options = executor.dynamic_planning.partitioning_hints + partitioning_hint = ( + rec.state["partitioning_hints"].get(ir) + if hint_options is not None and hint_options.use_partition_counts + else None + ) actors[ir] = [ join_actor( rec.state["context"], @@ -1161,6 +1180,7 @@ def _( channels[right].reserve_output_slot(), executor, collective_ids, + partitioning_hint, ) ] return actors, channels diff --git a/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py b/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py index 4bf08339e7e3..e36ac9dbee1f 100644 --- a/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py +++ b/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py @@ -58,11 +58,14 @@ class OrderPartitioningHint: """ Hint that rows should be ordered by a key sequence. + ``strict_key_count`` means a downstream consumer would benefit from strict + partitioning on the first N ordering keys. ``peer_partition_count`` is a planning-time estimate of the partition count that may align with a downstream peer input. """ keys: tuple[NamedOrderKey, ...] + strict_key_count: int | None = None peer_partition_count: int | None = None @@ -74,6 +77,7 @@ class OrderPartitioningHint: "NamedOrderKey", "OrderPartitioningHint", "PartitioningHint", + "apply_peer_partition_count_hint", "collect_partitioning_hints", ] @@ -113,6 +117,16 @@ def collect_partitioning_hints( return {node: hint for node, hint in hints.items() if hint is not None} +def apply_peer_partition_count_hint( + partition_count: int, + hint: PartitioningHint | None, +) -> int: + """Raise *partition_count* to *hint*'s peer count when present.""" + if hint is None or hint.peer_partition_count is None: + return partition_count + return max(partition_count, hint.peer_partition_count) + + def _direct_child_hints( ir: IR, partition_info: Mapping[IR, PartitionInfo], @@ -151,8 +165,12 @@ def _propagate_hint( return () (child,) = ir.children - if isinstance(ir, (Filter, Projection, Select, Slice)): + if isinstance(ir, (Projection, Select)): remapped = _remap_hint(hint, _child_zero_remapping(ir)) + elif isinstance(ir, (Filter, Slice)): + remapped = _clear_peer_partition_count( + _remap_hint(hint, _child_zero_remapping(ir)) + ) elif isinstance(ir, GroupBy): # Only group-key outputs remap to the child; aggregate outputs stop here. remapped = _remap_hint(hint, _child_zero_remapping(ir)) @@ -216,6 +234,14 @@ def _remap_hint( return dataclasses.replace(hint, keys=tuple(remapped_keys)) +def _clear_peer_partition_count( + hint: PartitioningHint | None, +) -> PartitioningHint | None: + if hint is None or hint.peer_partition_count is None: + return hint + return dataclasses.replace(hint, peer_partition_count=None) + + def _child_zero_remapping(ir: IR) -> dict[str, str]: return { output_name: binding.name @@ -259,7 +285,11 @@ def _merge_hints( None if order_keys is None else OrderPartitioningHint( - order_keys, peer_partition_count=peer_partition_count + order_keys, + strict_key_count=_merge_strict_key_count( + left.strict_key_count, right.strict_key_count + ), + peer_partition_count=peer_partition_count, ) ) @@ -275,8 +305,13 @@ def _merge_hints( if _is_prefix(hash_hint.keys, order_names) or _is_prefix( order_names, hash_hint.keys ): + strict_key_count = min(len(hash_hint.keys), len(order_names)) return dataclasses.replace( - order_hint, peer_partition_count=peer_partition_count + order_hint, + strict_key_count=_merge_strict_key_count( + order_hint.strict_key_count, strict_key_count + ), + peer_partition_count=peer_partition_count, ) return None @@ -292,6 +327,11 @@ def _merge_peer_partition_count( return max(counts) if counts else None +def _merge_strict_key_count(*counts: int | None) -> int | None: + count = max((count for count in counts if count is not None), default=0) + return count or None + + def _shortest_common_prefix( left: tuple[str, ...], right: tuple[str, ...] ) -> tuple[str, ...] | None: diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 315aa0c8f8cd..225785a5eefa 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -57,6 +57,7 @@ "InMemoryExecutor", "JoinFilterPushdownOptions", "ParquetOptions", + "PartitioningHintOptions", "RayContext", "SPMDContext", "StreamingExecutor", @@ -367,6 +368,48 @@ def default_broadcast_limit(min_device_size: int | None) -> int: return min(max(int(min_device_size * 0.15), 1), _DEFAULT_BROADCAST_LIMIT) +@dataclasses.dataclass(frozen=True) +class PartitioningHintOptions: + """ + Configuration for dynamic partitioning-hint usage. + + Parameters + ---------- + use_partition_counts + Whether shuffle actors may use compatible downstream partition-count + hints. Default is True. + use_ordering + Whether ordered actors may use compatible downstream ordering hints. + Default is True. + """ + + _env_prefix = "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__PARTITIONING_HINTS" + + use_partition_counts: bool = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__USE_PARTITION_COUNTS", _bool_converter, default=True + ) + ) + use_ordering: bool = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__USE_ORDERING", _bool_converter, default=True + ) + ) + + def __post_init__(self) -> None: # noqa: D105 + if not isinstance(self.use_partition_counts, bool): + raise TypeError("use_partition_counts must be a bool") + if not isinstance(self.use_ordering, bool): + raise TypeError("use_ordering must be a bool") + + +def _default_partitioning_hints() -> PartitioningHintOptions | None: + enabled = os.environ.get(PartitioningHintOptions._env_prefix) + if enabled is not None and not _bool_converter(enabled): + return None + return PartitioningHintOptions() + + @dataclasses.dataclass(frozen=True) class DynamicPlanningOptions: """ @@ -387,6 +430,9 @@ class DynamicPlanningOptions: sample_chunk_count The maximum number of chunks to sample before making dynamic-planning decisions. Default is 2. + partitioning_hints + Options controlling dynamic partitioning hints. ``None`` disables + partitioning-hint collection. """ _env_prefix = "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING" @@ -396,8 +442,24 @@ class DynamicPlanningOptions: f"{_env_prefix}__SAMPLE_CHUNK_COUNT", int, default=2 ) ) + partitioning_hints: PartitioningHintOptions | None = dataclasses.field( + default_factory=_default_partitioning_hints + ) def __post_init__(self) -> None: # noqa: D105 + if isinstance(self.partitioning_hints, dict): + object.__setattr__( + self, + "partitioning_hints", + PartitioningHintOptions(**self.partitioning_hints), + ) + if self.partitioning_hints is not None and not isinstance( + self.partitioning_hints, PartitioningHintOptions + ): + raise TypeError( + "partitioning_hints must be a PartitioningHintOptions " + "instance, dict, or None" + ) if not isinstance(self.sample_chunk_count, int): raise TypeError("sample_chunk_count must be an int") if self.sample_chunk_count < 1: @@ -878,6 +940,13 @@ def __post_init__(self) -> None: # noqa: D105 "dynamic_planning", DynamicPlanningOptions(**self.dynamic_planning), ) + if self.dynamic_planning is not None and not isinstance( + self.dynamic_planning, DynamicPlanningOptions + ): + raise TypeError( + "dynamic_planning must be a DynamicPlanningOptions " + "instance, dict, or None" + ) if isinstance(self.join_filter_pushdown, dict): object.__setattr__( diff --git a/python/cudf_polars/tests/streaming/test_groupby.py b/python/cudf_polars/tests/streaming/test_groupby.py index da6db146ca68..f8240678e000 100644 --- a/python/cudf_polars/tests/streaming/test_groupby.py +++ b/python/cudf_polars/tests/streaming/test_groupby.py @@ -19,6 +19,7 @@ from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.actor_graph import groupby as groupby_actor_graph from cudf_polars.streaming.actor_graph.collectives.shuffle import ShuffleManager +from cudf_polars.streaming.partitioning_hints import HashPartitioningHint from cudf_polars.testing.asserts import assert_gpu_result_equal @@ -73,8 +74,10 @@ async def fake_allgather_reduce(_context, _comm, _op_id, *local_values): True, # noqa: FBT003 [0], 1_000_000_000, + 0, False, # noqa: FBT003 False, # noqa: FBT003 + None, tracer, ) ) @@ -83,6 +86,76 @@ async def fake_allgather_reduce(_context, _comm, _op_id, *local_values): assert tracer.decision == "shuffle" +def test_dynamic_groupby_strategy_uses_peer_partition_count( + monkeypatch, strategy_chunk +): + """A partition-count hint may raise the shuffle output count.""" + + async def fake_allgather_reduce(_context, _comm, _op_id, *local_values): + estimated_size = strategy_chunk.data_alloc_size() * 4 + assert local_values == (estimated_size, 32, 4, 1) + return (estimated_size, 32, 4, 1) + + monkeypatch.setattr(groupby_actor_graph, "allgather_reduce", fake_allgather_reduce) + tracer = SimpleNamespace(decision=None) + + output_count = asyncio.run( + groupby_actor_graph._choose_strategy( + None, + None, + 4, + strategy_chunk, + 1, + False, # noqa: FBT003 + [0], + 1_000_000_000, + 0, + False, # noqa: FBT003 + False, # noqa: FBT003 + HashPartitioningHint(("key",), peer_partition_count=6), + tracer, + ) + ) + + assert output_count == 6 + assert tracer.decision == "shuffle" + + +def test_dynamic_groupby_strategy_skips_peer_partition_count_for_broadcastable_output( + monkeypatch, strategy_chunk +): + """Ignore peer partition count when grouped output is already broadcast-sized.""" + + async def fake_allgather_reduce(_context, _comm, _op_id, *local_values): + estimated_size = strategy_chunk.data_alloc_size() * 4 + assert local_values == (estimated_size, 32, 4, 1) + return (estimated_size, 32, 4, 1) + + monkeypatch.setattr(groupby_actor_graph, "allgather_reduce", fake_allgather_reduce) + tracer = SimpleNamespace(decision=None) + + output_count = asyncio.run( + groupby_actor_graph._choose_strategy( + None, + None, + 4, + strategy_chunk, + 1, + False, # noqa: FBT003 + [0], + 1_000_000_000, + 1_000_000_000, + False, # noqa: FBT003 + False, # noqa: FBT003 + HashPartitioningHint(("key",), peer_partition_count=6), + tracer, + ) + ) + + assert output_count == 2 + assert tracer.decision == "shuffle" + + @pytest.mark.parametrize("keys", [("key",), ("key", "key2")]) @pytest.mark.parametrize("agg", ["sum", "mean", "len", "min", "max"]) def test_dynamic_groupby_basic(df, streaming_engine, keys, agg): diff --git a/python/cudf_polars/tests/streaming/test_partitioning_hints.py b/python/cudf_polars/tests/streaming/test_partitioning_hints.py index 6ae51619b0fa..366b12aacc5f 100644 --- a/python/cudf_polars/tests/streaming/test_partitioning_hints.py +++ b/python/cudf_polars/tests/streaming/test_partitioning_hints.py @@ -9,7 +9,7 @@ from cudf_polars.containers import DataType from cudf_polars.dsl import expr -from cudf_polars.dsl.ir import DataFrameScan, GroupBy, Join, Select, Sort, Union +from cudf_polars.dsl.ir import DataFrameScan, Filter, GroupBy, Join, Select, Sort, Union from cudf_polars.streaming.base import PartitionInfo from cudf_polars.streaming.partitioning_hints import ( HashPartitioningHint, @@ -23,6 +23,7 @@ from cudf_polars.dsl.ir import IR I64 = DataType(pl.Int64()) +BOOL = DataType(pl.Boolean()) def make_scan(*names: str) -> DataFrameScan: @@ -106,6 +107,70 @@ def test_join_creates_hash_partitioning_hints() -> None: assert hints[right] == HashPartitioningHint(("k",), peer_partition_count=7) +def test_filter_clears_peer_partition_count() -> None: + scan = make_scan("k", "value") + mask = expr.NamedExpr("mask", expr.Literal(BOOL, True)) # noqa: FBT003 + filtered = Filter(scan.schema, mask, scan) + right = make_scan("k", "right_value") + join = Join( + {"k": I64, "value": I64, "right_value": I64}, + (named_col("k"),), + (named_col("k"),), + ("Inner", False, None, "_right", True, "none"), + filtered, + right, + ) + + hints = collect_partitioning_hints( + join, + { + join: PartitionInfo(7), + filtered: PartitionInfo(3), + scan: PartitionInfo(3), + right: PartitionInfo(7), + }, + ) + + assert hints[filtered] == HashPartitioningHint(("k",), peer_partition_count=7) + assert hints[scan] == HashPartitioningHint(("k",)) + assert hints[right] == HashPartitioningHint(("k",), peer_partition_count=7) + + +def test_groupby_preserves_peer_partition_count() -> None: + scan = make_scan("k", "value") + groupby = GroupBy( + {"k": I64}, + (named_col("k"),), + (), + maintain_order=False, + zlice=None, + df=scan, + ) + right = make_scan("k", "right_value") + join = Join( + {"k": I64, "right_value": I64}, + (named_col("k"),), + (named_col("k"),), + ("Inner", False, None, "_right", True, "none"), + groupby, + right, + ) + + hints = collect_partitioning_hints( + join, + { + join: PartitionInfo(7), + groupby: PartitionInfo(3), + scan: PartitionInfo(3), + right: PartitionInfo(7), + }, + ) + + assert hints[groupby] == HashPartitioningHint(("k",), peer_partition_count=7) + assert hints[scan] == HashPartitioningHint(("k",), peer_partition_count=7) + assert hints[right] == HashPartitioningHint(("k",), peer_partition_count=7) + + def test_groupby_creates_hash_partition_hint() -> None: scan = make_scan("a", "b") groupby = GroupBy( @@ -171,7 +236,8 @@ def test_groupby_remaps_order_partition_hint() -> None: ) assert hints[scan] == OrderPartitioningHint( - (NamedOrderKey("a", descending=False, nulls_last=False),) + (NamedOrderKey("a", descending=False, nulls_last=False),), + strict_key_count=1, ) @@ -232,6 +298,7 @@ def test_fanout_merges_compatible_hash_hint_into_order_hint() -> None: NamedOrderKey("a", descending=False, nulls_last=False), NamedOrderKey("b", descending=False, nulls_last=False), ), + strict_key_count=1, peer_partition_count=5, ) assert hints[right] == HashPartitioningHint(("a",), peer_partition_count=5) diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index fd184cf2f2c3..e23fb124e61e 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -35,6 +35,7 @@ JoinFilterPushdownOptions, MemoryResourceConfig, ParquetOptions, + PartitioningHintOptions, StreamingExecutor, Unspecified, ) @@ -693,9 +694,123 @@ def test_dynamic_planning_defaults() -> None: # Dynamic planning is enabled by default assert config.executor.dynamic_planning is not None assert config.executor.dynamic_planning.sample_chunk_count == 2 + assert config.executor.dynamic_planning.partitioning_hints is not None + assert config.executor.dynamic_planning.partitioning_hints.use_partition_counts + assert config.executor.dynamic_planning.partitioning_hints.use_ordering assert config.executor.join_filter_pushdown is None +def test_dynamic_planning_partitioning_hints_from_dict() -> None: + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "dynamic_planning": { + "partitioning_hints": { + "use_partition_counts": False, + "use_ordering": True, + } + } + }, + ) + ) + + assert config.executor.dynamic_planning is not None + assert config.executor.dynamic_planning.partitioning_hints is not None + assert not config.executor.dynamic_planning.partitioning_hints.use_partition_counts + assert config.executor.dynamic_planning.partitioning_hints.use_ordering + + +def test_dynamic_planning_partitioning_hints_from_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv( + "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__PARTITIONING_HINTS__USE_PARTITION_COUNTS", + "0", + ) + monkeypatch.setenv( + "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__PARTITIONING_HINTS__USE_ORDERING", + "1", + ) + + config = ConfigOptions.from_polars_engine(pl.GPUEngine()) + + assert config.executor.dynamic_planning is not None + assert config.executor.dynamic_planning.partitioning_hints is not None + assert not config.executor.dynamic_planning.partitioning_hints.use_partition_counts + assert config.executor.dynamic_planning.partitioning_hints.use_ordering + + +def test_dynamic_planning_partitioning_hints_disabled_from_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv( + "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__PARTITIONING_HINTS", "0" + ) + + config = ConfigOptions.from_polars_engine(pl.GPUEngine()) + + assert config.executor.dynamic_planning is not None + assert config.executor.dynamic_planning.partitioning_hints is None + + +def test_dynamic_planning_partitioning_hints_disabled() -> None: + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"dynamic_planning": {"partitioning_hints": None}}, + ) + ) + assert config.executor.dynamic_planning is not None + assert config.executor.dynamic_planning.partitioning_hints is None + + +def test_dynamic_planning_partitioning_hints_from_instance() -> None: + options = PartitioningHintOptions(use_partition_counts=False, use_ordering=False) + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "dynamic_planning": DynamicPlanningOptions(partitioning_hints=options), + }, + ) + ) + assert config.executor.dynamic_planning is not None + assert config.executor.dynamic_planning.partitioning_hints is options + + +def test_validate_dynamic_planning_partitioning_hints() -> None: + with pytest.raises(TypeError, match="use_partition_counts must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "dynamic_planning": { + "partitioning_hints": {"use_partition_counts": object()} + } + }, + ) + ) + with pytest.raises(TypeError, match="use_ordering must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "dynamic_planning": { + "partitioning_hints": {"use_ordering": object()} + } + }, + ) + ) + with pytest.raises(TypeError, match="partitioning_hints must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"dynamic_planning": {"partitioning_hints": object()}}, + ) + ) + + def test_dynamic_planning_disabled_from_env(monkeypatch: pytest.MonkeyPatch) -> None: # Test that env var can disable dynamic planning monkeypatch.setenv("CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING", "0") From 0a15c8825201f1a414cdb524a20424bb08158767 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 19 Aug 2026 12:09:47 -0700 Subject: [PATCH 05/12] roll-back rabbit-hole changes --- .../cudf_polars/streaming/actor_graph/core.py | 12 - .../streaming/actor_graph/dispatch.py | 6 +- .../streaming/actor_graph/groupby.py | 32 -- .../cudf_polars/streaming/actor_graph/join.py | 19 - .../streaming/partitioning_hints.py | 342 ++++++------------ .../cudf_polars/cudf_polars/utils/config.py | 62 ---- .../tests/streaming/test_groupby.py | 73 ---- .../streaming/test_partitioning_hints.py | 210 ++++------- python/cudf_polars/tests/test_config.py | 115 ------ 9 files changed, 175 insertions(+), 696 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py index 67ecb68bfd6d..fbec424d3104 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py @@ -23,7 +23,6 @@ metadata_drain_node, ) from cudf_polars.streaming.over import Over -from cudf_polars.streaming.partitioning_hints import collect_partitioning_hints from cudf_polars.utils.config import SPMDContext if TYPE_CHECKING: @@ -257,16 +256,6 @@ def generate_network( # Determine which nodes need fanout fanout_nodes = determine_fanout_nodes(ir, partition_info, ir_dep_count) - dynamic_planning = config_options.executor.dynamic_planning - hint_options = ( - None if dynamic_planning is None else dynamic_planning.partitioning_hints - ) - partitioning_hints = ( - collect_partitioning_hints(ir, partition_info) - if hint_options is not None - and (hint_options.use_partition_counts or hint_options.use_ordering) - else {} - ) # Generate the network state: GenState = { @@ -275,7 +264,6 @@ def generate_network( "config_options": config_options, "partition_info": partition_info, "fanout_nodes": fanout_nodes, - "partitioning_hints": partitioning_hints, "ir_context": ir_context, "max_concurrent_io_tasks": config_options.executor.max_concurrent_io_tasks, "stats": stats, diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py index 0dcdbe58673c..88d6784b0940 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py @@ -10,7 +10,7 @@ from cudf_polars.typing import GenericTransformer if TYPE_CHECKING: - from collections.abc import Mapping, MutableMapping + from collections.abc import MutableMapping from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.core.context import Context @@ -21,7 +21,6 @@ PartitionInfo, StatsCollector, ) - from cudf_polars.streaming.partitioning_hints import PartitioningHint from cudf_polars.utils.config import ConfigOptions, StreamingExecutor @@ -50,8 +49,6 @@ class GenState(TypedDict): Partition information. fanout_nodes Dictionary mapping IR nodes to fanout information. - partitioning_hints - Physical-layout hints for IR outputs. ir_context The execution context for the IR node. max_concurrent_io_tasks @@ -67,7 +64,6 @@ class GenState(TypedDict): config_options: ConfigOptions[StreamingExecutor] partition_info: MutableMapping[IR, PartitionInfo] fanout_nodes: dict[IR, FanoutInfo] - partitioning_hints: Mapping[IR, PartitioningHint] ir_context: IRExecutionContext max_concurrent_io_tasks: int stats: StatsCollector diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py index bfb1a2016f9b..e1c93620d1f0 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py @@ -45,7 +45,6 @@ shutdown_on_error, ) from cudf_polars.streaming.groupby import _has_stable_sorted_agg, combine, decompose -from cudf_polars.streaming.partitioning_hints import apply_peer_partition_count_hint from cudf_polars.streaming.repartition import Repartition if TYPE_CHECKING: @@ -56,7 +55,6 @@ from cudf_polars.dsl.ir import IRExecutionContext from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator from cudf_polars.streaming.actor_graph.tracing import ActorTracer - from cudf_polars.streaming.partitioning_hints import PartitioningHint from cudf_polars.typing import Schema @@ -517,10 +515,8 @@ async def _choose_strategy( input_drained: bool, # noqa: FBT001 collective_ids: list[int], target_partition_size: int, - broadcast_limit: int, skip_global_comm: bool, # noqa: FBT001 maintain_order: bool, # noqa: FBT001 - partitioning_hint: PartitioningHint | None, tracer: ActorTracer | None, ) -> int: """ @@ -544,14 +540,10 @@ async def _choose_strategy( The collective IDs. target_partition_size The target partition size. - broadcast_limit - The maximum size for broadcast join inputs. skip_global_comm Whether to skip the global communication. maintain_order Whether the operation should maintain input ordering semantics. - partitioning_hint - Optional downstream physical-layout hint. tracer Optional tracer for runtime metrics. @@ -610,14 +602,6 @@ async def _choose_strategy( output_count = min(ideal_count, output_count_limit) if not use_tree: output_count = max(2, output_count, min_row_limit_count) - broadcastable_output = ( - total_estimated_size < broadcast_limit - and total_estimated_rows < MAX_ROWS_PER_PARTITION - ) - if not broadcastable_output: - output_count = apply_peer_partition_count_hint( - output_count, partitioning_hint - ) if tracer is not None: tracer.decision = ( "tree_local" @@ -641,9 +625,7 @@ async def groupby_actor( ch_out: Channel[TableChunk], ch_in: Channel[TableChunk], target_partition_size: int, - broadcast_limit: int, collective_ids: list[int], - partitioning_hint: PartitioningHint | None, ) -> None: """ Dynamic GroupBy or Distinct actor that selects the best strategy at runtime. @@ -669,12 +651,8 @@ async def groupby_actor( The input channel. target_partition_size The target partition size. - broadcast_limit - The maximum size for broadcast join inputs. collective_ids The collective IDs. - partitioning_hint - Optional downstream physical-layout hint. """ async with shutdown_on_error( context, ch_in, ch_out, trace_ir=ir, ir_context=ir_context @@ -745,10 +723,8 @@ async def groupby_actor( input_drained, collective_ids, target_partition_size, - broadcast_limit, skip_global_comm, maintain_order, - partitioning_hint, tracer, ) @@ -803,12 +779,6 @@ def _( assert len(collective_ids) == 2, ( f"{type(ir).__name__} requires 2 collective IDs, got {len(collective_ids)}" ) - hint_options = config_options.executor.dynamic_planning.partitioning_hints - partitioning_hint = ( - rec.state["partitioning_hints"].get(ir) - if hint_options is not None and hint_options.use_partition_counts - else None - ) actors[ir] = [ groupby_actor( rec.state["context"], @@ -818,9 +788,7 @@ def _( channels[ir].reserve_input_slot(), channels[ir.children[0]].reserve_output_slot(), config_options.executor.target_partition_size, - config_options.executor.broadcast_limit, collective_ids, - partitioning_hint, ) ] diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index a21bef35515a..b2b8220eef70 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -56,7 +56,6 @@ send_metadata, shutdown_on_error, ) -from cudf_polars.streaming.partitioning_hints import apply_peer_partition_count_hint from cudf_polars.streaming.repartition import Repartition from cudf_polars.streaming.utils import _concat @@ -72,7 +71,6 @@ from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator from cudf_polars.streaming.actor_graph.tracing import ActorTracer from cudf_polars.streaming.base import PartitionInfo - from cudf_polars.streaming.partitioning_hints import PartitioningHint from cudf_polars.utils.config import StreamingExecutor @@ -716,7 +714,6 @@ def _choose_strategy_from_samples( left_sample: TableSizeStats, right_sample: TableSizeStats, chunkwise: bool, - partitioning_hint: PartitioningHint | None, tracer: ActorTracer | None, ) -> JoinStrategy: """Choose potential broadcast side and minimum shuffle modulus.""" @@ -798,9 +795,6 @@ def _choose_strategy_from_samples( estimated_rows_count + MAX_ROWS_PER_PARTITION - 1 ) // MAX_ROWS_PER_PARTITION min_shuffle_modulus = max(min_shuffle_modulus, min_partitions_for_row_limit) - min_shuffle_modulus = apply_peer_partition_count_hint( - min_shuffle_modulus, partitioning_hint - ) shuffle_modulus = _choose_shuffle_modulus( comm, @@ -865,7 +859,6 @@ async def _choose_strategy( right_metadata: ChannelMetadata, executor: StreamingExecutor, collective_ids: list[int], - partitioning_hint: PartitioningHint | None, *, tracer: ActorTracer | None, ) -> tuple[TableSizeStats, TableSizeStats, JoinStrategy]: @@ -940,7 +933,6 @@ async def _choose_strategy( left_sample=left_sample, right_sample=right_sample, chunkwise=chunkwise, - partitioning_hint=partitioning_hint, tracer=tracer, ) @@ -958,7 +950,6 @@ async def join_actor( ch_right: Channel[TableChunk], executor: StreamingExecutor, collective_ids: list[int], - partitioning_hint: PartitioningHint | None, ) -> None: """ Dynamic Join actor that selects the best strategy at runtime. @@ -987,8 +978,6 @@ async def join_actor( Streaming executor configuration. collective_ids List of collective IDs for shuffle/broadcast; consumed as needed. - partitioning_hint - Optional downstream physical-layout hint. """ async with shutdown_on_error( context, @@ -1013,7 +1002,6 @@ async def join_actor( right_metadata, executor, collective_ids, - partitioning_hint, tracer=tracer, ) ch_left_replay = context.create_channel() @@ -1163,12 +1151,6 @@ def _( f"{len(collective_ids)} for this Join. " "Ensure ReserveOpIDs is run with dynamic_planning enabled." ) - hint_options = executor.dynamic_planning.partitioning_hints - partitioning_hint = ( - rec.state["partitioning_hints"].get(ir) - if hint_options is not None and hint_options.use_partition_counts - else None - ) actors[ir] = [ join_actor( rec.state["context"], @@ -1180,7 +1162,6 @@ def _( channels[right].reserve_output_slot(), executor, collective_ids, - partitioning_hint, ) ] return actors, channels diff --git a/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py b/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py index e36ac9dbee1f..a648852c1409 100644 --- a/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py +++ b/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py @@ -4,7 +4,7 @@ from __future__ import annotations -import dataclasses +from dataclasses import dataclass from typing import TYPE_CHECKING, TypeAlias import pylibcudf as plc @@ -22,16 +22,14 @@ ) from cudf_polars.dsl.traversal import post_traversal from cudf_polars.dsl.utils.column_domain import column_domain_bindings -from cudf_polars.streaming.repartition import Repartition if TYPE_CHECKING: - from collections.abc import Mapping + from collections.abc import Iterator, Mapping from cudf_polars.dsl.ir import IR - from cudf_polars.streaming.base import PartitionInfo -@dataclasses.dataclass(frozen=True) +@dataclass(frozen=True) class NamedOrderKey: """Named sort key with logical Polars ordering options.""" @@ -40,167 +38,98 @@ class NamedOrderKey: nulls_last: bool -@dataclasses.dataclass(frozen=True) -class HashPartitioningHint: - """ - Hint that rows should be co-located by equality keys. - - ``peer_partition_count`` is a planning-time estimate of the partition - count that may align with a downstream peer input. - """ +@dataclass(frozen=True) +class StrictPartitioningHint: + """Hint that a downstream consumer wants strict partitioning by keys.""" keys: tuple[str, ...] - peer_partition_count: int | None = None -@dataclasses.dataclass(frozen=True) +@dataclass(frozen=True) class OrderPartitioningHint: - """ - Hint that rows should be ordered by a key sequence. - - ``strict_key_count`` means a downstream consumer would benefit from strict - partitioning on the first N ordering keys. - ``peer_partition_count`` is a planning-time estimate of the partition - count that may align with a downstream peer input. - """ + """Hint that upstream rows should be ordered, optionally with a strict prefix.""" keys: tuple[NamedOrderKey, ...] strict_key_count: int | None = None - peer_partition_count: int | None = None - - -PartitioningHint: TypeAlias = HashPartitioningHint | OrderPartitioningHint -_MaybeHint: TypeAlias = PartitioningHint | None - -__all__ = [ - "HashPartitioningHint", - "NamedOrderKey", - "OrderPartitioningHint", - "PartitioningHint", - "apply_peer_partition_count_hint", - "collect_partitioning_hints", -] - - -def collect_partitioning_hints( - ir: IR, - partition_info: Mapping[IR, PartitionInfo], -) -> dict[IR, PartitioningHint]: - """ - Collect physical-layout hints for each IR node. - - Parameters - ---------- - ir - The IR to collect partitioning hints for. - partition_info : Mapping[IR, PartitionInfo] - The partition information for each IR node. - - Returns - ------- - The physical-layout hints for each IR node. - - Notes - ----- - The returned hints describe layouts that downstream consumers would prefer, - if practical. These hints are not a strict runtime contract. - """ - hints: dict[IR, _MaybeHint] = {} - for node in reversed(list(post_traversal([ir]))): - for child, hint in _direct_child_hints(node, partition_info): - _record_hint(hints, child, hint) - if (current_hint := hints.get(node)) is not None: - for child, child_hint in _propagate_hint(node, current_hint): - _record_hint(hints, child, child_hint) - return {node: hint for node, hint in hints.items() if hint is not None} +PartitioningHint: TypeAlias = StrictPartitioningHint | OrderPartitioningHint -def apply_peer_partition_count_hint( - partition_count: int, - hint: PartitioningHint | None, -) -> int: - """Raise *partition_count* to *hint*'s peer count when present.""" - if hint is None or hint.peer_partition_count is None: - return partition_count - return max(partition_count, hint.peer_partition_count) +def collect_partitioning_hints(ir: IR) -> dict[IR, PartitioningHint]: + """Collect non-conflicting upstream partitioning hints for each IR node.""" + hints: dict[IR, PartitioningHint | None] = {} + for node in reversed(list(post_traversal([ir]))): + for child, child_hint in _construct_child_hints(node): + _record_hint(hints, child, child_hint) + + node_hint = hints.get(node) + if ( + node_hint is None + or len(node.children) != 1 + or not isinstance(node, (Projection, Select, Filter, Slice, GroupBy)) + ): + continue + + remapped = _remap_hint( + node_hint, + { + output_name: binding.name + for output_name, binding in column_domain_bindings(node).items() + if binding.child_index == 0 + }, + ) + if remapped is not None: + _record_hint(hints, node.children[0], remapped) + return {node: hint for node, hint in hints.items() if hint is not None} -def _direct_child_hints( - ir: IR, - partition_info: Mapping[IR, PartitionInfo], -) -> tuple[tuple[IR, PartitioningHint], ...]: - """Return hints created directly by *ir* for its children.""" + +def _construct_child_hints(ir: IR) -> Iterator[tuple[IR, PartitioningHint]]: + """Construct hints for the upstream children of *ir*.""" if isinstance(ir, Sort): - hint = _sort_hint(ir) - return () if hint is None else ((ir.children[0], hint),) + names = _column_names(ir.by) + if names is not None: + yield ( + ir.children[0], + _order_hint( + names, + tuple(order == plc.types.Order.DESCENDING for order in ir.order), + tuple( + (order == plc.types.Order.ASCENDING) + == (null_order == plc.types.NullOrder.AFTER) + for order, null_order in zip( + ir.order, ir.null_order, strict=True + ) + ), + ), + ) - if isinstance(ir, Join) and ir.options[0] != "Cross": + elif isinstance(ir, MapFunction) and ir.name == "hint_sorted": + yield ir.children[0], _order_hint(*ir.options) + + elif isinstance(ir, Join) and ir.options[0] != "Cross": left_keys = _column_names(ir.left_on) right_keys = _column_names(ir.right_on) - if left_keys is None or right_keys is None: - return () - peer_partition_count = max( - partition_info[ir.children[0]].count, - partition_info[ir.children[1]].count, - ) - return ( - (ir.children[0], HashPartitioningHint(left_keys, peer_partition_count)), - (ir.children[1], HashPartitioningHint(right_keys, peer_partition_count)), - ) + if left_keys is not None and right_keys is not None: + yield ir.children[0], StrictPartitioningHint(left_keys) + yield ir.children[1], StrictPartitioningHint(right_keys) - if isinstance(ir, GroupBy) and not ir.maintain_order: + elif isinstance(ir, GroupBy) and not ir.maintain_order: keys = _column_names(ir.keys) - return () if keys is None else ((ir.children[0], HashPartitioningHint(keys)),) + if keys is not None: + yield ir.children[0], StrictPartitioningHint(keys) - return () - -def _propagate_hint( - ir: IR, hint: PartitioningHint -) -> tuple[tuple[IR, PartitioningHint], ...]: - """Propagate *hint* through nodes whose children may benefit from it.""" - if len(ir.children) != 1: - return () - - (child,) = ir.children - if isinstance(ir, (Projection, Select)): - remapped = _remap_hint(hint, _child_zero_remapping(ir)) - elif isinstance(ir, (Filter, Slice)): - remapped = _clear_peer_partition_count( - _remap_hint(hint, _child_zero_remapping(ir)) - ) - elif isinstance(ir, GroupBy): - # Only group-key outputs remap to the child; aggregate outputs stop here. - remapped = _remap_hint(hint, _child_zero_remapping(ir)) - elif isinstance(ir, Repartition) or ( - isinstance(ir, MapFunction) and ir.name in {"hint_sorted", "rechunk"} - ): - remapped = _remap_hint(hint, {name: name for name in ir.schema}) - elif isinstance(ir, MapFunction) and ir.name == "row_index": - remapped = _remap_hint(hint, {name: name for name in child.schema}) - else: - remapped = None - - return () if remapped is None else ((child, remapped),) - - -def _sort_hint(ir: Sort) -> OrderPartitioningHint | None: - names = _column_names(ir.by) - if names is None: - return None +def _order_hint( + names: tuple[str, ...], + descending: tuple[bool, ...], + nulls_last: tuple[bool, ...], +) -> OrderPartitioningHint: return OrderPartitioningHint( tuple( - NamedOrderKey( - name, - order == plc.types.Order.DESCENDING, - (order == plc.types.Order.ASCENDING) - == (null_order == plc.types.NullOrder.AFTER), - ) - for name, order, null_order in zip( - names, ir.order, ir.null_order, strict=True - ) + NamedOrderKey(name, desc, null_last) + for name, desc, null_last in zip(names, descending, nulls_last, strict=True) ) ) @@ -217,41 +146,25 @@ def _column_names(named_exprs: tuple[expr.NamedExpr, ...]) -> tuple[str, ...] | def _remap_hint( hint: PartitioningHint, remapping: Mapping[str, str] ) -> PartitioningHint | None: - if isinstance(hint, HashPartitioningHint): + if isinstance(hint, StrictPartitioningHint): remapped_names = [] for name in hint.keys: if (new_name := remapping.get(name)) is None: return None remapped_names.append(new_name) - return dataclasses.replace(hint, keys=tuple(remapped_names)) + return StrictPartitioningHint(tuple(remapped_names)) remapped_keys = [] for key in hint.keys: new_name = remapping.get(key.name) if new_name is None: return None - remapped_keys.append(dataclasses.replace(key, name=new_name)) - return dataclasses.replace(hint, keys=tuple(remapped_keys)) - - -def _clear_peer_partition_count( - hint: PartitioningHint | None, -) -> PartitioningHint | None: - if hint is None or hint.peer_partition_count is None: - return hint - return dataclasses.replace(hint, peer_partition_count=None) - - -def _child_zero_remapping(ir: IR) -> dict[str, str]: - return { - output_name: binding.name - for output_name, binding in column_domain_bindings(ir).items() - if binding.child_index == 0 - } + remapped_keys.append(NamedOrderKey(new_name, key.descending, key.nulls_last)) + return OrderPartitioningHint(tuple(remapped_keys), hint.strict_key_count) def _record_hint( - hints: dict[IR, _MaybeHint], + hints: dict[IR, PartitioningHint | None], node: IR, hint: PartitioningHint, ) -> None: @@ -264,93 +177,48 @@ def _record_hint( def _merge_hints( left: PartitioningHint, right: PartitioningHint ) -> PartitioningHint | None: - peer_partition_count = _merge_peer_partition_count(left, right) - if isinstance(left, HashPartitioningHint) and isinstance( - right, HashPartitioningHint - ): - hash_keys = _shortest_common_prefix(left.keys, right.keys) - return ( - None - if hash_keys is None - else HashPartitioningHint( - hash_keys, peer_partition_count=peer_partition_count - ) - ) + if isinstance(left, StrictPartitioningHint): + if isinstance(right, StrictPartitioningHint): + if _is_prefix(left.keys, right.keys): + return left + if _is_prefix(right.keys, left.keys): + return right + return None + return _merge_order_with_strict(right, left) - if isinstance(left, OrderPartitioningHint) and isinstance( - right, OrderPartitioningHint - ): - order_keys = _longest_common_extension(left.keys, right.keys) - return ( - None - if order_keys is None - else OrderPartitioningHint( - order_keys, - strict_key_count=_merge_strict_key_count( - left.strict_key_count, right.strict_key_count - ), - peer_partition_count=peer_partition_count, - ) - ) + if isinstance(right, StrictPartitioningHint): + return _merge_order_with_strict(left, right) - if isinstance(left, OrderPartitioningHint): - order_hint = left - assert isinstance(right, HashPartitioningHint) - hash_hint = right + if _is_prefix(left.keys, right.keys): + keys = right.keys + elif _is_prefix(right.keys, left.keys): + keys = left.keys else: - assert isinstance(right, OrderPartitioningHint) - order_hint = right - hash_hint = left + return None + return OrderPartitioningHint( + keys, _merge_strict_key_count(left.strict_key_count, right.strict_key_count) + ) + + +def _merge_order_with_strict( + order_hint: OrderPartitioningHint, strict_hint: StrictPartitioningHint +) -> OrderPartitioningHint | None: order_names = tuple(key.name for key in order_hint.keys) - if _is_prefix(hash_hint.keys, order_names) or _is_prefix( - order_names, hash_hint.keys + if _is_prefix(strict_hint.keys, order_names) or _is_prefix( + order_names, strict_hint.keys ): - strict_key_count = min(len(hash_hint.keys), len(order_names)) - return dataclasses.replace( - order_hint, - strict_key_count=_merge_strict_key_count( - order_hint.strict_key_count, strict_key_count - ), - peer_partition_count=peer_partition_count, + strict_key_count = min(len(strict_hint.keys), len(order_names)) + return OrderPartitioningHint( + order_hint.keys, + _merge_strict_key_count(order_hint.strict_key_count, strict_key_count), ) return None -def _merge_peer_partition_count( - left: PartitioningHint, right: PartitioningHint -) -> int | None: - counts = [ - count - for count in (left.peer_partition_count, right.peer_partition_count) - if count is not None - ] - return max(counts) if counts else None - - def _merge_strict_key_count(*counts: int | None) -> int | None: count = max((count for count in counts if count is not None), default=0) return count or None -def _shortest_common_prefix( - left: tuple[str, ...], right: tuple[str, ...] -) -> tuple[str, ...] | None: - if _is_prefix(left, right): - return left - if _is_prefix(right, left): - return right - return None - - -def _longest_common_extension( - left: tuple[NamedOrderKey, ...], right: tuple[NamedOrderKey, ...] -) -> tuple[NamedOrderKey, ...] | None: - if _is_prefix(left, right): - return right - if _is_prefix(right, left): - return left - return None - - def _is_prefix(left: tuple[object, ...], right: tuple[object, ...]) -> bool: return len(left) <= len(right) and left == right[: len(left)] diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 225785a5eefa..58de100394c1 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -57,7 +57,6 @@ "InMemoryExecutor", "JoinFilterPushdownOptions", "ParquetOptions", - "PartitioningHintOptions", "RayContext", "SPMDContext", "StreamingExecutor", @@ -368,48 +367,6 @@ def default_broadcast_limit(min_device_size: int | None) -> int: return min(max(int(min_device_size * 0.15), 1), _DEFAULT_BROADCAST_LIMIT) -@dataclasses.dataclass(frozen=True) -class PartitioningHintOptions: - """ - Configuration for dynamic partitioning-hint usage. - - Parameters - ---------- - use_partition_counts - Whether shuffle actors may use compatible downstream partition-count - hints. Default is True. - use_ordering - Whether ordered actors may use compatible downstream ordering hints. - Default is True. - """ - - _env_prefix = "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__PARTITIONING_HINTS" - - use_partition_counts: bool = dataclasses.field( - default_factory=_make_default_factory( - f"{_env_prefix}__USE_PARTITION_COUNTS", _bool_converter, default=True - ) - ) - use_ordering: bool = dataclasses.field( - default_factory=_make_default_factory( - f"{_env_prefix}__USE_ORDERING", _bool_converter, default=True - ) - ) - - def __post_init__(self) -> None: # noqa: D105 - if not isinstance(self.use_partition_counts, bool): - raise TypeError("use_partition_counts must be a bool") - if not isinstance(self.use_ordering, bool): - raise TypeError("use_ordering must be a bool") - - -def _default_partitioning_hints() -> PartitioningHintOptions | None: - enabled = os.environ.get(PartitioningHintOptions._env_prefix) - if enabled is not None and not _bool_converter(enabled): - return None - return PartitioningHintOptions() - - @dataclasses.dataclass(frozen=True) class DynamicPlanningOptions: """ @@ -430,9 +387,6 @@ class DynamicPlanningOptions: sample_chunk_count The maximum number of chunks to sample before making dynamic-planning decisions. Default is 2. - partitioning_hints - Options controlling dynamic partitioning hints. ``None`` disables - partitioning-hint collection. """ _env_prefix = "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING" @@ -442,24 +396,8 @@ class DynamicPlanningOptions: f"{_env_prefix}__SAMPLE_CHUNK_COUNT", int, default=2 ) ) - partitioning_hints: PartitioningHintOptions | None = dataclasses.field( - default_factory=_default_partitioning_hints - ) def __post_init__(self) -> None: # noqa: D105 - if isinstance(self.partitioning_hints, dict): - object.__setattr__( - self, - "partitioning_hints", - PartitioningHintOptions(**self.partitioning_hints), - ) - if self.partitioning_hints is not None and not isinstance( - self.partitioning_hints, PartitioningHintOptions - ): - raise TypeError( - "partitioning_hints must be a PartitioningHintOptions " - "instance, dict, or None" - ) if not isinstance(self.sample_chunk_count, int): raise TypeError("sample_chunk_count must be an int") if self.sample_chunk_count < 1: diff --git a/python/cudf_polars/tests/streaming/test_groupby.py b/python/cudf_polars/tests/streaming/test_groupby.py index f8240678e000..da6db146ca68 100644 --- a/python/cudf_polars/tests/streaming/test_groupby.py +++ b/python/cudf_polars/tests/streaming/test_groupby.py @@ -19,7 +19,6 @@ from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.actor_graph import groupby as groupby_actor_graph from cudf_polars.streaming.actor_graph.collectives.shuffle import ShuffleManager -from cudf_polars.streaming.partitioning_hints import HashPartitioningHint from cudf_polars.testing.asserts import assert_gpu_result_equal @@ -74,10 +73,8 @@ async def fake_allgather_reduce(_context, _comm, _op_id, *local_values): True, # noqa: FBT003 [0], 1_000_000_000, - 0, False, # noqa: FBT003 False, # noqa: FBT003 - None, tracer, ) ) @@ -86,76 +83,6 @@ async def fake_allgather_reduce(_context, _comm, _op_id, *local_values): assert tracer.decision == "shuffle" -def test_dynamic_groupby_strategy_uses_peer_partition_count( - monkeypatch, strategy_chunk -): - """A partition-count hint may raise the shuffle output count.""" - - async def fake_allgather_reduce(_context, _comm, _op_id, *local_values): - estimated_size = strategy_chunk.data_alloc_size() * 4 - assert local_values == (estimated_size, 32, 4, 1) - return (estimated_size, 32, 4, 1) - - monkeypatch.setattr(groupby_actor_graph, "allgather_reduce", fake_allgather_reduce) - tracer = SimpleNamespace(decision=None) - - output_count = asyncio.run( - groupby_actor_graph._choose_strategy( - None, - None, - 4, - strategy_chunk, - 1, - False, # noqa: FBT003 - [0], - 1_000_000_000, - 0, - False, # noqa: FBT003 - False, # noqa: FBT003 - HashPartitioningHint(("key",), peer_partition_count=6), - tracer, - ) - ) - - assert output_count == 6 - assert tracer.decision == "shuffle" - - -def test_dynamic_groupby_strategy_skips_peer_partition_count_for_broadcastable_output( - monkeypatch, strategy_chunk -): - """Ignore peer partition count when grouped output is already broadcast-sized.""" - - async def fake_allgather_reduce(_context, _comm, _op_id, *local_values): - estimated_size = strategy_chunk.data_alloc_size() * 4 - assert local_values == (estimated_size, 32, 4, 1) - return (estimated_size, 32, 4, 1) - - monkeypatch.setattr(groupby_actor_graph, "allgather_reduce", fake_allgather_reduce) - tracer = SimpleNamespace(decision=None) - - output_count = asyncio.run( - groupby_actor_graph._choose_strategy( - None, - None, - 4, - strategy_chunk, - 1, - False, # noqa: FBT003 - [0], - 1_000_000_000, - 1_000_000_000, - False, # noqa: FBT003 - False, # noqa: FBT003 - HashPartitioningHint(("key",), peer_partition_count=6), - tracer, - ) - ) - - assert output_count == 2 - assert tracer.decision == "shuffle" - - @pytest.mark.parametrize("keys", [("key",), ("key", "key2")]) @pytest.mark.parametrize("agg", ["sum", "mean", "len", "min", "max"]) def test_dynamic_groupby_basic(df, streaming_engine, keys, agg): diff --git a/python/cudf_polars/tests/streaming/test_partitioning_hints.py b/python/cudf_polars/tests/streaming/test_partitioning_hints.py index 366b12aacc5f..878b4f479b7c 100644 --- a/python/cudf_polars/tests/streaming/test_partitioning_hints.py +++ b/python/cudf_polars/tests/streaming/test_partitioning_hints.py @@ -9,12 +9,19 @@ from cudf_polars.containers import DataType from cudf_polars.dsl import expr -from cudf_polars.dsl.ir import DataFrameScan, Filter, GroupBy, Join, Select, Sort, Union -from cudf_polars.streaming.base import PartitionInfo +from cudf_polars.dsl.ir import ( + DataFrameScan, + GroupBy, + Join, + MapFunction, + Select, + Sort, + Union, +) from cudf_polars.streaming.partitioning_hints import ( - HashPartitioningHint, NamedOrderKey, OrderPartitioningHint, + StrictPartitioningHint, collect_partitioning_hints, ) from cudf_polars.utils.sorting import sort_order @@ -23,7 +30,6 @@ from cudf_polars.dsl.ir import IR I64 = DataType(pl.Int64()) -BOOL = DataType(pl.Boolean()) def make_scan(*names: str) -> DataFrameScan: @@ -59,6 +65,21 @@ def make_sort( ) +def make_hint_sorted( + child: IR, + *names: str, + descending: tuple[bool, ...] | None = None, + nulls_last: tuple[bool, ...] | None = None, +) -> MapFunction: + if descending is None: + descending = (False,) * len(names) + if nulls_last is None: + nulls_last = (False,) * len(names) + return MapFunction( + child.schema, "hint_sorted", (names, descending, nulls_last), child + ) + + def test_sort_creates_order_partition_hint() -> None: scan = make_scan("a", "b") sort = make_sort( @@ -69,10 +90,7 @@ def test_sort_creates_order_partition_hint() -> None: nulls_last=(True, False), ) - hints = collect_partitioning_hints( - sort, - {sort: PartitionInfo(2), scan: PartitionInfo(2)}, - ) + hints = collect_partitioning_hints(sort) assert hints[scan] == OrderPartitioningHint( ( @@ -82,135 +100,67 @@ def test_sort_creates_order_partition_hint() -> None: ) -def test_join_creates_hash_partitioning_hints() -> None: - left = make_scan("k", "left_value") - right = make_scan("k", "right_value") - join = Join( - {"k": I64, "left_value": I64, "right_value": I64}, - (named_col("k"),), - (named_col("k"),), - ("Inner", False, None, "_right", True, "none"), - left, - right, +def test_select_remaps_order_partition_hint() -> None: + scan = make_scan("a", "b") + select = Select( + {"x": I64, "b": I64}, + (expr.NamedExpr("x", expr.Col(I64, "a")), named_col("b")), + should_broadcast=False, + df=scan, ) + sort = make_sort(select, "x") - hints = collect_partitioning_hints( - join, - { - join: PartitionInfo(7), - left: PartitionInfo(3), - right: PartitionInfo(7), - }, + hints = collect_partitioning_hints(sort) + + assert hints[scan] == OrderPartitioningHint( + (NamedOrderKey("a", descending=False, nulls_last=False),) ) - assert hints[left] == HashPartitioningHint(("k",), peer_partition_count=7) - assert hints[right] == HashPartitioningHint(("k",), peer_partition_count=7) +def test_hint_sorted_creates_order_partition_hint() -> None: + scan = make_scan("a", "b") + hint_sorted = make_hint_sorted(scan, "a", descending=(True,)) -def test_filter_clears_peer_partition_count() -> None: - scan = make_scan("k", "value") - mask = expr.NamedExpr("mask", expr.Literal(BOOL, True)) # noqa: FBT003 - filtered = Filter(scan.schema, mask, scan) - right = make_scan("k", "right_value") - join = Join( - {"k": I64, "value": I64, "right_value": I64}, - (named_col("k"),), - (named_col("k"),), - ("Inner", False, None, "_right", True, "none"), - filtered, - right, - ) + hints = collect_partitioning_hints(hint_sorted) - hints = collect_partitioning_hints( - join, - { - join: PartitionInfo(7), - filtered: PartitionInfo(3), - scan: PartitionInfo(3), - right: PartitionInfo(7), - }, + assert hints[scan] == OrderPartitioningHint( + (NamedOrderKey("a", descending=True, nulls_last=False),) ) - assert hints[filtered] == HashPartitioningHint(("k",), peer_partition_count=7) - assert hints[scan] == HashPartitioningHint(("k",)) - assert hints[right] == HashPartitioningHint(("k",), peer_partition_count=7) +def test_hint_sorted_keeps_declared_order_with_compatible_downstream_sort() -> None: + scan = make_scan("a", "b") + hint_sorted = make_hint_sorted(scan, "a") + sort = make_sort(hint_sorted, "a") -def test_groupby_preserves_peer_partition_count() -> None: - scan = make_scan("k", "value") - groupby = GroupBy( - {"k": I64}, - (named_col("k"),), - (), - maintain_order=False, - zlice=None, - df=scan, - ) - right = make_scan("k", "right_value") - join = Join( - {"k": I64, "right_value": I64}, - (named_col("k"),), - (named_col("k"),), - ("Inner", False, None, "_right", True, "none"), - groupby, - right, - ) + hints = collect_partitioning_hints(sort) - hints = collect_partitioning_hints( - join, - { - join: PartitionInfo(7), - groupby: PartitionInfo(3), - scan: PartitionInfo(3), - right: PartitionInfo(7), - }, + assert hints[scan] == OrderPartitioningHint( + (NamedOrderKey("a", descending=False, nulls_last=False),) ) - assert hints[groupby] == HashPartitioningHint(("k",), peer_partition_count=7) - assert hints[scan] == HashPartitioningHint(("k",), peer_partition_count=7) - assert hints[right] == HashPartitioningHint(("k",), peer_partition_count=7) - -def test_groupby_creates_hash_partition_hint() -> None: +def test_hint_sorted_keeps_declared_order_with_extended_downstream_sort() -> None: scan = make_scan("a", "b") - groupby = GroupBy( - {"a": I64}, - (named_col("a"),), - (), - maintain_order=False, - zlice=None, - df=scan, - ) + hint_sorted = make_hint_sorted(scan, "a") + sort = make_sort(hint_sorted, "a", "b") - hints = collect_partitioning_hints( - groupby, - {groupby: PartitionInfo(2), scan: PartitionInfo(2)}, - ) + hints = collect_partitioning_hints(sort) - assert hints[scan] == HashPartitioningHint(("a",)) + assert hints[scan] == OrderPartitioningHint( + (NamedOrderKey("a", descending=False, nulls_last=False),) + ) -def test_select_remaps_order_partition_hint() -> None: +def test_hint_sorted_keeps_declared_order_with_incompatible_downstream_sort() -> None: scan = make_scan("a", "b") - select = Select( - {"x": I64, "b": I64}, - (expr.NamedExpr("x", expr.Col(I64, "a")), named_col("b")), - should_broadcast=False, - df=scan, - ) - sort = make_sort(select, "x") + hint_sorted = make_hint_sorted(scan, "a", descending=(True,)) + sort = make_sort(hint_sorted, "a") - hints = collect_partitioning_hints( - sort, - { - sort: PartitionInfo(2), - select: PartitionInfo(2), - scan: PartitionInfo(2), - }, - ) + hints = collect_partitioning_hints(sort) assert hints[scan] == OrderPartitioningHint( - (NamedOrderKey("a", descending=False, nulls_last=False),) + (NamedOrderKey("a", descending=True, nulls_last=False),) ) @@ -226,14 +176,7 @@ def test_groupby_remaps_order_partition_hint() -> None: ) sort = make_sort(groupby, "key") - hints = collect_partitioning_hints( - sort, - { - sort: PartitionInfo(2), - groupby: PartitionInfo(2), - scan: PartitionInfo(2), - }, - ) + hints = collect_partitioning_hints(sort) assert hints[scan] == OrderPartitioningHint( (NamedOrderKey("a", descending=False, nulls_last=False),), @@ -251,10 +194,7 @@ def test_fanout_keeps_more_specific_compatible_order_hint() -> None: make_sort(scan, "a", "b"), ) - hints = collect_partitioning_hints( - root, - {root: PartitionInfo(2), scan: PartitionInfo(2)}, - ) + hints = collect_partitioning_hints(root) assert hints[scan] == OrderPartitioningHint( ( @@ -264,7 +204,7 @@ def test_fanout_keeps_more_specific_compatible_order_hint() -> None: ) -def test_fanout_merges_compatible_hash_hint_into_order_hint() -> None: +def test_fanout_marks_compatible_order_hint_as_strict() -> None: scan = make_scan("a", "b") right = make_scan("a", "right_value") join = Join( @@ -283,15 +223,7 @@ def test_fanout_merges_compatible_hash_hint_into_order_hint() -> None: join, ) - hints = collect_partitioning_hints( - root, - { - root: PartitionInfo(5), - join: PartitionInfo(5), - scan: PartitionInfo(3), - right: PartitionInfo(5), - }, - ) + hints = collect_partitioning_hints(root) assert hints[scan] == OrderPartitioningHint( ( @@ -299,9 +231,8 @@ def test_fanout_merges_compatible_hash_hint_into_order_hint() -> None: NamedOrderKey("b", descending=False, nulls_last=False), ), strict_key_count=1, - peer_partition_count=5, ) - assert hints[right] == HashPartitioningHint(("a",), peer_partition_count=5) + assert hints[right] == StrictPartitioningHint(("a",)) def test_conflicting_fanout_drops_hint() -> None: @@ -314,9 +245,6 @@ def test_conflicting_fanout_drops_hint() -> None: make_sort(scan, "b"), ) - hints = collect_partitioning_hints( - root, - {root: PartitionInfo(2), scan: PartitionInfo(2)}, - ) + hints = collect_partitioning_hints(root) assert scan not in hints diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index e23fb124e61e..fd184cf2f2c3 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -35,7 +35,6 @@ JoinFilterPushdownOptions, MemoryResourceConfig, ParquetOptions, - PartitioningHintOptions, StreamingExecutor, Unspecified, ) @@ -694,123 +693,9 @@ def test_dynamic_planning_defaults() -> None: # Dynamic planning is enabled by default assert config.executor.dynamic_planning is not None assert config.executor.dynamic_planning.sample_chunk_count == 2 - assert config.executor.dynamic_planning.partitioning_hints is not None - assert config.executor.dynamic_planning.partitioning_hints.use_partition_counts - assert config.executor.dynamic_planning.partitioning_hints.use_ordering assert config.executor.join_filter_pushdown is None -def test_dynamic_planning_partitioning_hints_from_dict() -> None: - config = ConfigOptions.from_polars_engine( - pl.GPUEngine( - executor="streaming", - executor_options={ - "dynamic_planning": { - "partitioning_hints": { - "use_partition_counts": False, - "use_ordering": True, - } - } - }, - ) - ) - - assert config.executor.dynamic_planning is not None - assert config.executor.dynamic_planning.partitioning_hints is not None - assert not config.executor.dynamic_planning.partitioning_hints.use_partition_counts - assert config.executor.dynamic_planning.partitioning_hints.use_ordering - - -def test_dynamic_planning_partitioning_hints_from_env( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv( - "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__PARTITIONING_HINTS__USE_PARTITION_COUNTS", - "0", - ) - monkeypatch.setenv( - "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__PARTITIONING_HINTS__USE_ORDERING", - "1", - ) - - config = ConfigOptions.from_polars_engine(pl.GPUEngine()) - - assert config.executor.dynamic_planning is not None - assert config.executor.dynamic_planning.partitioning_hints is not None - assert not config.executor.dynamic_planning.partitioning_hints.use_partition_counts - assert config.executor.dynamic_planning.partitioning_hints.use_ordering - - -def test_dynamic_planning_partitioning_hints_disabled_from_env( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv( - "CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING__PARTITIONING_HINTS", "0" - ) - - config = ConfigOptions.from_polars_engine(pl.GPUEngine()) - - assert config.executor.dynamic_planning is not None - assert config.executor.dynamic_planning.partitioning_hints is None - - -def test_dynamic_planning_partitioning_hints_disabled() -> None: - config = ConfigOptions.from_polars_engine( - pl.GPUEngine( - executor="streaming", - executor_options={"dynamic_planning": {"partitioning_hints": None}}, - ) - ) - assert config.executor.dynamic_planning is not None - assert config.executor.dynamic_planning.partitioning_hints is None - - -def test_dynamic_planning_partitioning_hints_from_instance() -> None: - options = PartitioningHintOptions(use_partition_counts=False, use_ordering=False) - config = ConfigOptions.from_polars_engine( - pl.GPUEngine( - executor="streaming", - executor_options={ - "dynamic_planning": DynamicPlanningOptions(partitioning_hints=options), - }, - ) - ) - assert config.executor.dynamic_planning is not None - assert config.executor.dynamic_planning.partitioning_hints is options - - -def test_validate_dynamic_planning_partitioning_hints() -> None: - with pytest.raises(TypeError, match="use_partition_counts must be"): - ConfigOptions.from_polars_engine( - pl.GPUEngine( - executor="streaming", - executor_options={ - "dynamic_planning": { - "partitioning_hints": {"use_partition_counts": object()} - } - }, - ) - ) - with pytest.raises(TypeError, match="use_ordering must be"): - ConfigOptions.from_polars_engine( - pl.GPUEngine( - executor="streaming", - executor_options={ - "dynamic_planning": { - "partitioning_hints": {"use_ordering": object()} - } - }, - ) - ) - with pytest.raises(TypeError, match="partitioning_hints must be"): - ConfigOptions.from_polars_engine( - pl.GPUEngine( - executor="streaming", - executor_options={"dynamic_planning": {"partitioning_hints": object()}}, - ) - ) - - def test_dynamic_planning_disabled_from_env(monkeypatch: pytest.MonkeyPatch) -> None: # Test that env var can disable dynamic planning monkeypatch.setenv("CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING", "0") From 4e1ca4d4afa4370fa003fa4ee8ba3596f8b67ea9 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 19 Aug 2026 13:01:32 -0700 Subject: [PATCH 06/12] track multiple hints --- .../streaming/partitioning_hints.py | 63 ++-- .../streaming/test_partitioning_hints.py | 317 +++++++++++++++--- 2 files changed, 312 insertions(+), 68 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py b/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py index a648852c1409..6e47027dc0d7 100644 --- a/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py +++ b/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py @@ -56,33 +56,47 @@ class OrderPartitioningHint: PartitioningHint: TypeAlias = StrictPartitioningHint | OrderPartitioningHint -def collect_partitioning_hints(ir: IR) -> dict[IR, PartitioningHint]: - """Collect non-conflicting upstream partitioning hints for each IR node.""" - hints: dict[IR, PartitioningHint | None] = {} +def collect_partitioning_hints(ir: IR) -> dict[IR, tuple[PartitioningHint, ...]]: + """Collect candidate upstream partitioning hints for each IR node.""" + hints: dict[IR, tuple[PartitioningHint, ...]] = {} for node in reversed(list(post_traversal([ir]))): - for child, child_hint in _construct_child_hints(node): - _record_hint(hints, child, child_hint) + child_hints = list(_construct_child_hints(node)) - node_hint = hints.get(node) + node_hints = hints.get(node) if ( - node_hint is None - or len(node.children) != 1 - or not isinstance(node, (Projection, Select, Filter, Slice, GroupBy)) + node_hints is not None + and len(node.children) == 1 + and isinstance(node, (Projection, Select, Filter, Slice, GroupBy)) ): - continue - - remapped = _remap_hint( - node_hint, - { + remapping = { output_name: binding.name for output_name, binding in column_domain_bindings(node).items() if binding.child_index == 0 - }, - ) - if remapped is not None: - _record_hint(hints, node.children[0], remapped) + } + child_hints.extend( + (node.children[0], remapped) + for node_hint in node_hints + if (remapped := _remap_hint(node_hint, remapping)) is not None + ) - return {node: hint for node, hint in hints.items() if hint is not None} + for child, child_hint in child_hints: + candidates: list[PartitioningHint] = [] + new_hint = child_hint + insertion_index: int | None = None + for existing_hint in hints.get(child, ()): + if (merged := _merge_hints(existing_hint, new_hint)) is None: + candidates.append(existing_hint) + else: + new_hint = merged + if insertion_index is None: + insertion_index = len(candidates) + if insertion_index is None: + candidates.append(new_hint) + else: + candidates.insert(insertion_index, new_hint) + hints[child] = tuple(candidates) + + return hints def _construct_child_hints(ir: IR) -> Iterator[tuple[IR, PartitioningHint]]: @@ -163,17 +177,6 @@ def _remap_hint( return OrderPartitioningHint(tuple(remapped_keys), hint.strict_key_count) -def _record_hint( - hints: dict[IR, PartitioningHint | None], - node: IR, - hint: PartitioningHint, -) -> None: - if node not in hints: - hints[node] = hint - elif (current := hints[node]) is not None: - hints[node] = _merge_hints(current, hint) - - def _merge_hints( left: PartitioningHint, right: PartitioningHint ) -> PartitioningHint | None: diff --git a/python/cudf_polars/tests/streaming/test_partitioning_hints.py b/python/cudf_polars/tests/streaming/test_partitioning_hints.py index 878b4f479b7c..79c363c67e99 100644 --- a/python/cudf_polars/tests/streaming/test_partitioning_hints.py +++ b/python/cudf_polars/tests/streaming/test_partitioning_hints.py @@ -80,6 +80,17 @@ def make_hint_sorted( ) +def make_groupby(child: IR, *names: str) -> GroupBy: + return GroupBy( + dict.fromkeys(names, I64), + tuple(named_col(name) for name in names), + (), + maintain_order=False, + zlice=None, + df=child, + ) + + def test_sort_creates_order_partition_hint() -> None: scan = make_scan("a", "b") sort = make_sort( @@ -92,11 +103,13 @@ def test_sort_creates_order_partition_hint() -> None: hints = collect_partitioning_hints(sort) - assert hints[scan] == OrderPartitioningHint( - ( - NamedOrderKey("a", descending=False, nulls_last=True), - NamedOrderKey("b", descending=True, nulls_last=False), - ) + assert hints[scan] == ( + OrderPartitioningHint( + ( + NamedOrderKey("a", descending=False, nulls_last=True), + NamedOrderKey("b", descending=True, nulls_last=False), + ) + ), ) @@ -112,10 +125,30 @@ def test_select_remaps_order_partition_hint() -> None: hints = collect_partitioning_hints(sort) - assert hints[scan] == OrderPartitioningHint( - (NamedOrderKey("a", descending=False, nulls_last=False),) + assert hints[scan] == ( + OrderPartitioningHint( + (NamedOrderKey("a", descending=False, nulls_last=False),) + ), + ) + + +def test_non_column_sort_does_not_create_hint() -> None: + scan = make_scan("a") + order, null_order = sort_order((False,), nulls_last=(False,), num_keys=1) + sort = Sort( + scan.schema, + (expr.NamedExpr("literal", expr.Literal(I64, 1)),), + order, + null_order, + stable=False, + zlice=None, + df=scan, ) + hints = collect_partitioning_hints(sort) + + assert hints == {} + def test_hint_sorted_creates_order_partition_hint() -> None: scan = make_scan("a", "b") @@ -123,10 +156,79 @@ def test_hint_sorted_creates_order_partition_hint() -> None: hints = collect_partitioning_hints(hint_sorted) - assert hints[scan] == OrderPartitioningHint( - (NamedOrderKey("a", descending=True, nulls_last=False),) + assert hints[scan] == ( + OrderPartitioningHint((NamedOrderKey("a", descending=True, nulls_last=False),)), + ) + + +def test_select_remaps_strict_partition_hint() -> None: + scan = make_scan("a") + select = Select( + {"x": I64}, + (expr.NamedExpr("x", expr.Col(I64, "a")),), + should_broadcast=False, + df=scan, + ) + right = make_scan("x") + join = Join( + {"x": I64}, + (named_col("x"),), + (named_col("x"),), + ("Inner", False, None, "_right", True, "none"), + select, + right, + ) + + hints = collect_partitioning_hints(join) + + assert hints[scan] == (StrictPartitioningHint(("a",)),) + assert hints[right] == (StrictPartitioningHint(("x",)),) + + +def test_select_drops_order_hint_on_non_column_output() -> None: + scan = make_scan("a") + select = Select( + {"x": I64}, + (expr.NamedExpr("x", expr.Literal(I64, 1)),), + should_broadcast=False, + df=scan, + ) + sort = make_sort(select, "x") + + hints = collect_partitioning_hints(sort) + + assert hints[select] == ( + OrderPartitioningHint( + (NamedOrderKey("x", descending=False, nulls_last=False),) + ), + ) + assert scan not in hints + + +def test_select_drops_strict_hint_on_non_column_output() -> None: + scan = make_scan("a") + select = Select( + {"x": I64}, + (expr.NamedExpr("x", expr.Literal(I64, 1)),), + should_broadcast=False, + df=scan, + ) + right = make_scan("x") + join = Join( + {"x": I64}, + (named_col("x"),), + (named_col("x"),), + ("Inner", False, None, "_right", True, "none"), + select, + right, ) + hints = collect_partitioning_hints(join) + + assert hints[select] == (StrictPartitioningHint(("x",)),) + assert hints[right] == (StrictPartitioningHint(("x",)),) + assert scan not in hints + def test_hint_sorted_keeps_declared_order_with_compatible_downstream_sort() -> None: scan = make_scan("a", "b") @@ -135,8 +237,10 @@ def test_hint_sorted_keeps_declared_order_with_compatible_downstream_sort() -> N hints = collect_partitioning_hints(sort) - assert hints[scan] == OrderPartitioningHint( - (NamedOrderKey("a", descending=False, nulls_last=False),) + assert hints[scan] == ( + OrderPartitioningHint( + (NamedOrderKey("a", descending=False, nulls_last=False),) + ), ) @@ -147,8 +251,10 @@ def test_hint_sorted_keeps_declared_order_with_extended_downstream_sort() -> Non hints = collect_partitioning_hints(sort) - assert hints[scan] == OrderPartitioningHint( - (NamedOrderKey("a", descending=False, nulls_last=False),) + assert hints[scan] == ( + OrderPartitioningHint( + (NamedOrderKey("a", descending=False, nulls_last=False),) + ), ) @@ -159,28 +265,31 @@ def test_hint_sorted_keeps_declared_order_with_incompatible_downstream_sort() -> hints = collect_partitioning_hints(sort) - assert hints[scan] == OrderPartitioningHint( - (NamedOrderKey("a", descending=True, nulls_last=False),) + assert hints[scan] == ( + OrderPartitioningHint((NamedOrderKey("a", descending=True, nulls_last=False),)), ) def test_groupby_remaps_order_partition_hint() -> None: scan = make_scan("a", "b") - groupby = GroupBy( - {"key": I64}, - (expr.NamedExpr("key", expr.Col(I64, "a")),), - (), - maintain_order=False, - zlice=None, - df=scan, + groupby = make_groupby( + Select( + {"key": I64}, + (expr.NamedExpr("key", expr.Col(I64, "a")),), + should_broadcast=False, + df=scan, + ), + "key", ) sort = make_sort(groupby, "key") hints = collect_partitioning_hints(sort) - assert hints[scan] == OrderPartitioningHint( - (NamedOrderKey("a", descending=False, nulls_last=False),), - strict_key_count=1, + assert hints[scan] == ( + OrderPartitioningHint( + (NamedOrderKey("a", descending=False, nulls_last=False),), + strict_key_count=1, + ), ) @@ -196,11 +305,13 @@ def test_fanout_keeps_more_specific_compatible_order_hint() -> None: hints = collect_partitioning_hints(root) - assert hints[scan] == OrderPartitioningHint( - ( - NamedOrderKey("a", descending=False, nulls_last=False), - NamedOrderKey("b", descending=False, nulls_last=False), - ) + assert hints[scan] == ( + OrderPartitioningHint( + ( + NamedOrderKey("a", descending=False, nulls_last=False), + NamedOrderKey("b", descending=False, nulls_last=False), + ) + ), ) @@ -225,17 +336,116 @@ def test_fanout_marks_compatible_order_hint_as_strict() -> None: hints = collect_partitioning_hints(root) - assert hints[scan] == OrderPartitioningHint( - ( - NamedOrderKey("a", descending=False, nulls_last=False), - NamedOrderKey("b", descending=False, nulls_last=False), + assert hints[scan] == ( + OrderPartitioningHint( + ( + NamedOrderKey("a", descending=False, nulls_last=False), + NamedOrderKey("b", descending=False, nulls_last=False), + ), + strict_key_count=1, ), - strict_key_count=1, ) - assert hints[right] == StrictPartitioningHint(("a",)) + assert hints[right] == (StrictPartitioningHint(("a",)),) -def test_conflicting_fanout_drops_hint() -> None: +def test_fanout_merges_compatible_strict_hints() -> None: + scan = make_scan("a", "b") + root = Union( + scan.schema, + None, + False, # noqa: FBT003 + make_groupby(scan, "a"), + make_groupby(scan, "a", "b"), + ) + + hints = collect_partitioning_hints(root) + + assert hints[scan] == (StrictPartitioningHint(("a",)),) + + +def test_fanout_keeps_incompatible_strict_candidates() -> None: + scan = make_scan("a", "b") + root = Union( + scan.schema, + None, + False, # noqa: FBT003 + make_groupby(scan, "a"), + make_groupby(scan, "b"), + ) + + hints = collect_partitioning_hints(root) + + assert set(hints[scan]) == { + StrictPartitioningHint(("a",)), + StrictPartitioningHint(("b",)), + } + + +def test_fanout_keeps_incompatible_order_and_strict_candidates() -> None: + scan = make_scan("a", "b") + root = Union( + scan.schema, + None, + False, # noqa: FBT003 + make_sort(scan, "a"), + make_groupby(scan, "b"), + ) + + hints = collect_partitioning_hints(root) + + assert set(hints[scan]) == { + OrderPartitioningHint( + (NamedOrderKey("a", descending=False, nulls_last=False),) + ), + StrictPartitioningHint(("b",)), + } + + +def test_compatible_hint_merging_is_not_directional() -> None: + scan = make_scan("a", "b") + root = Union( + scan.schema, + None, + False, # noqa: FBT003 + make_groupby(scan, "a", "b"), + make_groupby(scan, "a"), + ) + assert collect_partitioning_hints(root)[scan] == (StrictPartitioningHint(("a",)),) + + scan = make_scan("a", "b") + root = Union( + scan.schema, + None, + False, # noqa: FBT003 + make_groupby(scan, "b"), + make_sort(scan, "a"), + ) + assert set(collect_partitioning_hints(root)[scan]) == { + StrictPartitioningHint(("b",)), + OrderPartitioningHint( + (NamedOrderKey("a", descending=False, nulls_last=False),) + ), + } + + scan = make_scan("a", "b") + root = Union( + scan.schema, + None, + False, # noqa: FBT003 + make_sort(scan, "a", "b"), + make_sort(scan, "a"), + ) + assert collect_partitioning_hints(root)[scan] == ( + OrderPartitioningHint( + ( + NamedOrderKey("a", descending=False, nulls_last=False), + NamedOrderKey("b", descending=False, nulls_last=False), + ) + ), + ) + + +def test_conflicting_fanout_keeps_candidate_hints() -> None: scan = make_scan("a", "b") root = Union( scan.schema, @@ -247,4 +457,35 @@ def test_conflicting_fanout_drops_hint() -> None: hints = collect_partitioning_hints(root) - assert scan not in hints + assert set(hints[scan]) == { + OrderPartitioningHint( + (NamedOrderKey("a", descending=False, nulls_last=False),) + ), + OrderPartitioningHint( + (NamedOrderKey("b", descending=False, nulls_last=False),) + ), + } + + +def test_repeated_fanout_candidate_is_merged() -> None: + scan = make_scan("a", "b") + root = Union( + scan.schema, + None, + False, # noqa: FBT003 + make_sort(scan, "a"), + make_sort(scan, "b"), + make_sort(scan, "a"), + ) + + hints = collect_partitioning_hints(root) + + assert len(hints[scan]) == 2 + assert set(hints[scan]) == { + OrderPartitioningHint( + (NamedOrderKey("a", descending=False, nulls_last=False),) + ), + OrderPartitioningHint( + (NamedOrderKey("b", descending=False, nulls_last=False),) + ), + } From c670199bf8c22a3d3bef0e489323d3a35afa800a Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 19 Aug 2026 13:26:58 -0700 Subject: [PATCH 07/12] cleanup --- .../streaming/partitioning_hints.py | 142 ++++++++++-------- 1 file changed, 83 insertions(+), 59 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py b/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py index 6e47027dc0d7..e8bce4031f1d 100644 --- a/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py +++ b/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py @@ -24,7 +24,7 @@ from cudf_polars.dsl.utils.column_domain import column_domain_bindings if TYPE_CHECKING: - from collections.abc import Iterator, Mapping + from collections.abc import Mapping from cudf_polars.dsl.ir import IR @@ -57,82 +57,103 @@ class OrderPartitioningHint: def collect_partitioning_hints(ir: IR) -> dict[IR, tuple[PartitioningHint, ...]]: - """Collect candidate upstream partitioning hints for each IR node.""" + """Collect upstream partitioning hints for an IR graph.""" hints: dict[IR, tuple[PartitioningHint, ...]] = {} for node in reversed(list(post_traversal([ir]))): - child_hints = list(_construct_child_hints(node)) - - node_hints = hints.get(node) - if ( - node_hints is not None - and len(node.children) == 1 - and isinstance(node, (Projection, Select, Filter, Slice, GroupBy)) - ): - remapping = { - output_name: binding.name - for output_name, binding in column_domain_bindings(node).items() - if binding.child_index == 0 - } - child_hints.extend( - (node.children[0], remapped) - for node_hint in node_hints - if (remapped := _remap_hint(node_hint, remapping)) is not None - ) - + child_hints = _direct_child_hints(node) + child_hints.extend(_propagated_child_hints(node, hints)) for child, child_hint in child_hints: - candidates: list[PartitioningHint] = [] - new_hint = child_hint - insertion_index: int | None = None - for existing_hint in hints.get(child, ()): - if (merged := _merge_hints(existing_hint, new_hint)) is None: - candidates.append(existing_hint) - else: - new_hint = merged - if insertion_index is None: - insertion_index = len(candidates) - if insertion_index is None: - candidates.append(new_hint) - else: - candidates.insert(insertion_index, new_hint) - hints[child] = tuple(candidates) - + hints[child] = _merge_candidate_hint(hints.get(child, ()), child_hint) return hints -def _construct_child_hints(ir: IR) -> Iterator[tuple[IR, PartitioningHint]]: - """Construct hints for the upstream children of *ir*.""" +def _direct_child_hints(ir: IR) -> list[tuple[IR, PartitioningHint]]: + """Create hints implied directly by operators that consume a partitioning.""" if isinstance(ir, Sort): names = _column_names(ir.by) if names is not None: - yield ( - ir.children[0], - _order_hint( - names, - tuple(order == plc.types.Order.DESCENDING for order in ir.order), - tuple( - (order == plc.types.Order.ASCENDING) - == (null_order == plc.types.NullOrder.AFTER) - for order, null_order in zip( - ir.order, ir.null_order, strict=True - ) + return [ + ( + ir.children[0], + _order_hint( + names, + tuple( + order == plc.types.Order.DESCENDING for order in ir.order + ), + tuple( + (order == plc.types.Order.ASCENDING) + == (null_order == plc.types.NullOrder.AFTER) + for order, null_order in zip( + ir.order, ir.null_order, strict=True + ) + ), ), - ), - ) + ) + ] - elif isinstance(ir, MapFunction) and ir.name == "hint_sorted": - yield ir.children[0], _order_hint(*ir.options) + if isinstance(ir, MapFunction) and ir.name == "hint_sorted": + return [(ir.children[0], _order_hint(*ir.options))] - elif isinstance(ir, Join) and ir.options[0] != "Cross": + if isinstance(ir, Join) and ir.options[0] != "Cross": left_keys = _column_names(ir.left_on) right_keys = _column_names(ir.right_on) if left_keys is not None and right_keys is not None: - yield ir.children[0], StrictPartitioningHint(left_keys) - yield ir.children[1], StrictPartitioningHint(right_keys) + return [ + (ir.children[0], StrictPartitioningHint(left_keys)), + (ir.children[1], StrictPartitioningHint(right_keys)), + ] - elif isinstance(ir, GroupBy) and not ir.maintain_order: + if isinstance(ir, GroupBy) and not ir.maintain_order: keys = _column_names(ir.keys) if keys is not None: - yield ir.children[0], StrictPartitioningHint(keys) + return [(ir.children[0], StrictPartitioningHint(keys))] + + return [] + + +def _propagated_child_hints( + node: IR, hints: dict[IR, tuple[PartitioningHint, ...]] +) -> list[tuple[IR, PartitioningHint]]: + """Push compatible downstream hints through single-child operators.""" + child_hints: list[tuple[IR, PartitioningHint]] = [] + node_hints = hints.get(node) + if ( + node_hints is not None + and len(node.children) == 1 + and isinstance(node, (Projection, Select, Filter, Slice, GroupBy)) + ): + remapping = { + output_name: binding.name + for output_name, binding in column_domain_bindings(node).items() + if binding.child_index == 0 + } + child_hints.extend( + (node.children[0], remapped) + for node_hint in node_hints + if (remapped := _remap_hint(node_hint, remapping)) is not None + ) + return child_hints + + +def _merge_candidate_hint( + existing_hints: tuple[PartitioningHint, ...], child_hint: PartitioningHint +) -> tuple[PartitioningHint, ...]: + """Merge one hint into compatible candidates and preserve incompatible ones.""" + candidates: list[PartitioningHint] = [] + new_hint = child_hint + insertion_index: int | None = None + for existing_hint in existing_hints: + if (merged := _merge_hints(existing_hint, new_hint)) is None: + candidates.append(existing_hint) + else: + new_hint = merged + if insertion_index is None: + insertion_index = len(candidates) + if insertion_index is None: + candidates.append(new_hint) + else: + candidates.insert(insertion_index, new_hint) + return tuple(candidates) def _order_hint( @@ -160,6 +181,7 @@ def _column_names(named_exprs: tuple[expr.NamedExpr, ...]) -> tuple[str, ...] | def _remap_hint( hint: PartitioningHint, remapping: Mapping[str, str] ) -> PartitioningHint | None: + """Rewrite hint column names through a child-to-parent name mapping.""" if isinstance(hint, StrictPartitioningHint): remapped_names = [] for name in hint.keys: @@ -180,6 +202,7 @@ def _remap_hint( def _merge_hints( left: PartitioningHint, right: PartitioningHint ) -> PartitioningHint | None: + """Merge compatible hints, or return None if both should remain candidates.""" if isinstance(left, StrictPartitioningHint): if isinstance(right, StrictPartitioningHint): if _is_prefix(left.keys, right.keys): @@ -206,6 +229,7 @@ def _merge_hints( def _merge_order_with_strict( order_hint: OrderPartitioningHint, strict_hint: StrictPartitioningHint ) -> OrderPartitioningHint | None: + """Fold strict-key requirements into compatible ordering hints.""" order_names = tuple(key.name for key in order_hint.keys) if _is_prefix(strict_hint.keys, order_names) or _is_prefix( order_names, strict_hint.keys From 39c2150f7f626256f920b1e6e3832afc90b70af3 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 20 Aug 2026 10:43:05 -0700 Subject: [PATCH 08/12] remove stale code --- python/cudf_polars/cudf_polars/utils/config.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 58de100394c1..315aa0c8f8cd 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -878,13 +878,6 @@ def __post_init__(self) -> None: # noqa: D105 "dynamic_planning", DynamicPlanningOptions(**self.dynamic_planning), ) - if self.dynamic_planning is not None and not isinstance( - self.dynamic_planning, DynamicPlanningOptions - ): - raise TypeError( - "dynamic_planning must be a DynamicPlanningOptions " - "instance, dict, or None" - ) if isinstance(self.join_filter_pushdown, dict): object.__setattr__( From d235c52dfd0675c8d708cba700d6a06f3ada820c Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 20 Aug 2026 10:45:28 -0700 Subject: [PATCH 09/12] roll back accidental change --- python/cudf_polars/cudf_polars/streaming/actor_graph/join.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index b2b8220eef70..d944847e621b 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -878,10 +878,9 @@ async def _choose_strategy( hash_chunkwise = isinstance( left_partitioning.inter_rank_scheme, HashScheme ) and isinstance(right_partitioning.inter_rank_scheme, HashScheme) - aligned = hash_chunkwise and left_partitioning.is_aligned_with( + if hash_chunkwise and left_partitioning.is_aligned_with( right_partitioning, context.br() - ) - if aligned: + ): # We can use a chunkwise join chunkwise = True left_sample = TableSizeStats( From 5843c994f20b5831e7beb57b4139d6f6dc7d76f3 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 21 Aug 2026 09:36:44 -0700 Subject: [PATCH 10/12] adopt 'request' language instead of 'hint' --- .../streaming/partitioning_hints.py | 251 ---------------- .../streaming/partitioning_requests.py | 271 ++++++++++++++++++ ...hints.py => test_partitioning_requests.py} | 172 +++++------ 3 files changed, 360 insertions(+), 334 deletions(-) delete mode 100644 python/cudf_polars/cudf_polars/streaming/partitioning_hints.py create mode 100644 python/cudf_polars/cudf_polars/streaming/partitioning_requests.py rename python/cudf_polars/tests/streaming/{test_partitioning_hints.py => test_partitioning_requests.py} (69%) diff --git a/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py b/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py deleted file mode 100644 index e8bce4031f1d..000000000000 --- a/python/cudf_polars/cudf_polars/streaming/partitioning_hints.py +++ /dev/null @@ -1,251 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Partitioning hints for streaming actor-graph construction.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import TYPE_CHECKING, TypeAlias - -import pylibcudf as plc - -from cudf_polars.dsl import expr -from cudf_polars.dsl.ir import ( - Filter, - GroupBy, - Join, - MapFunction, - Projection, - Select, - Slice, - Sort, -) -from cudf_polars.dsl.traversal import post_traversal -from cudf_polars.dsl.utils.column_domain import column_domain_bindings - -if TYPE_CHECKING: - from collections.abc import Mapping - - from cudf_polars.dsl.ir import IR - - -@dataclass(frozen=True) -class NamedOrderKey: - """Named sort key with logical Polars ordering options.""" - - name: str - descending: bool - nulls_last: bool - - -@dataclass(frozen=True) -class StrictPartitioningHint: - """Hint that a downstream consumer wants strict partitioning by keys.""" - - keys: tuple[str, ...] - - -@dataclass(frozen=True) -class OrderPartitioningHint: - """Hint that upstream rows should be ordered, optionally with a strict prefix.""" - - keys: tuple[NamedOrderKey, ...] - strict_key_count: int | None = None - - -PartitioningHint: TypeAlias = StrictPartitioningHint | OrderPartitioningHint - - -def collect_partitioning_hints(ir: IR) -> dict[IR, tuple[PartitioningHint, ...]]: - """Collect upstream partitioning hints for an IR graph.""" - hints: dict[IR, tuple[PartitioningHint, ...]] = {} - for node in reversed(list(post_traversal([ir]))): - child_hints = _direct_child_hints(node) - child_hints.extend(_propagated_child_hints(node, hints)) - for child, child_hint in child_hints: - hints[child] = _merge_candidate_hint(hints.get(child, ()), child_hint) - return hints - - -def _direct_child_hints(ir: IR) -> list[tuple[IR, PartitioningHint]]: - """Create hints implied directly by operators that consume a partitioning.""" - if isinstance(ir, Sort): - names = _column_names(ir.by) - if names is not None: - return [ - ( - ir.children[0], - _order_hint( - names, - tuple( - order == plc.types.Order.DESCENDING for order in ir.order - ), - tuple( - (order == plc.types.Order.ASCENDING) - == (null_order == plc.types.NullOrder.AFTER) - for order, null_order in zip( - ir.order, ir.null_order, strict=True - ) - ), - ), - ) - ] - - if isinstance(ir, MapFunction) and ir.name == "hint_sorted": - return [(ir.children[0], _order_hint(*ir.options))] - - if isinstance(ir, Join) and ir.options[0] != "Cross": - left_keys = _column_names(ir.left_on) - right_keys = _column_names(ir.right_on) - if left_keys is not None and right_keys is not None: - return [ - (ir.children[0], StrictPartitioningHint(left_keys)), - (ir.children[1], StrictPartitioningHint(right_keys)), - ] - - if isinstance(ir, GroupBy) and not ir.maintain_order: - keys = _column_names(ir.keys) - if keys is not None: - return [(ir.children[0], StrictPartitioningHint(keys))] - - return [] - - -def _propagated_child_hints( - node: IR, hints: dict[IR, tuple[PartitioningHint, ...]] -) -> list[tuple[IR, PartitioningHint]]: - """Push compatible downstream hints through single-child operators.""" - child_hints: list[tuple[IR, PartitioningHint]] = [] - node_hints = hints.get(node) - if ( - node_hints is not None - and len(node.children) == 1 - and isinstance(node, (Projection, Select, Filter, Slice, GroupBy)) - ): - remapping = { - output_name: binding.name - for output_name, binding in column_domain_bindings(node).items() - if binding.child_index == 0 - } - child_hints.extend( - (node.children[0], remapped) - for node_hint in node_hints - if (remapped := _remap_hint(node_hint, remapping)) is not None - ) - return child_hints - - -def _merge_candidate_hint( - existing_hints: tuple[PartitioningHint, ...], child_hint: PartitioningHint -) -> tuple[PartitioningHint, ...]: - """Merge one hint into compatible candidates and preserve incompatible ones.""" - candidates: list[PartitioningHint] = [] - new_hint = child_hint - insertion_index: int | None = None - for existing_hint in existing_hints: - if (merged := _merge_hints(existing_hint, new_hint)) is None: - candidates.append(existing_hint) - else: - new_hint = merged - if insertion_index is None: - insertion_index = len(candidates) - if insertion_index is None: - candidates.append(new_hint) - else: - candidates.insert(insertion_index, new_hint) - return tuple(candidates) - - -def _order_hint( - names: tuple[str, ...], - descending: tuple[bool, ...], - nulls_last: tuple[bool, ...], -) -> OrderPartitioningHint: - return OrderPartitioningHint( - tuple( - NamedOrderKey(name, desc, null_last) - for name, desc, null_last in zip(names, descending, nulls_last, strict=True) - ) - ) - - -def _column_names(named_exprs: tuple[expr.NamedExpr, ...]) -> tuple[str, ...] | None: - names = [] - for named_expr in named_exprs: - if not isinstance(named_expr.value, expr.Col): - return None - names.append(named_expr.value.name) - return tuple(names) - - -def _remap_hint( - hint: PartitioningHint, remapping: Mapping[str, str] -) -> PartitioningHint | None: - """Rewrite hint column names through a child-to-parent name mapping.""" - if isinstance(hint, StrictPartitioningHint): - remapped_names = [] - for name in hint.keys: - if (new_name := remapping.get(name)) is None: - return None - remapped_names.append(new_name) - return StrictPartitioningHint(tuple(remapped_names)) - - remapped_keys = [] - for key in hint.keys: - new_name = remapping.get(key.name) - if new_name is None: - return None - remapped_keys.append(NamedOrderKey(new_name, key.descending, key.nulls_last)) - return OrderPartitioningHint(tuple(remapped_keys), hint.strict_key_count) - - -def _merge_hints( - left: PartitioningHint, right: PartitioningHint -) -> PartitioningHint | None: - """Merge compatible hints, or return None if both should remain candidates.""" - if isinstance(left, StrictPartitioningHint): - if isinstance(right, StrictPartitioningHint): - if _is_prefix(left.keys, right.keys): - return left - if _is_prefix(right.keys, left.keys): - return right - return None - return _merge_order_with_strict(right, left) - - if isinstance(right, StrictPartitioningHint): - return _merge_order_with_strict(left, right) - - if _is_prefix(left.keys, right.keys): - keys = right.keys - elif _is_prefix(right.keys, left.keys): - keys = left.keys - else: - return None - return OrderPartitioningHint( - keys, _merge_strict_key_count(left.strict_key_count, right.strict_key_count) - ) - - -def _merge_order_with_strict( - order_hint: OrderPartitioningHint, strict_hint: StrictPartitioningHint -) -> OrderPartitioningHint | None: - """Fold strict-key requirements into compatible ordering hints.""" - order_names = tuple(key.name for key in order_hint.keys) - if _is_prefix(strict_hint.keys, order_names) or _is_prefix( - order_names, strict_hint.keys - ): - strict_key_count = min(len(strict_hint.keys), len(order_names)) - return OrderPartitioningHint( - order_hint.keys, - _merge_strict_key_count(order_hint.strict_key_count, strict_key_count), - ) - return None - - -def _merge_strict_key_count(*counts: int | None) -> int | None: - count = max((count for count in counts if count is not None), default=0) - return count or None - - -def _is_prefix(left: tuple[object, ...], right: tuple[object, ...]) -> bool: - return len(left) <= len(right) and left == right[: len(left)] diff --git a/python/cudf_polars/cudf_polars/streaming/partitioning_requests.py b/python/cudf_polars/cudf_polars/streaming/partitioning_requests.py new file mode 100644 index 000000000000..db948e67a494 --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/partitioning_requests.py @@ -0,0 +1,271 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +""" +Downstream partitioning requests for streaming actor-graph construction. + +A partitioning request is attached to an IR node when a downstream consumer +may benefit if that node produces data with a specific layout. These requests +do not describe or guarantee the current layout of the node's output; actual +runtime layout metadata is tracked separately in ``ChannelMetadata``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, TypeAlias + +import pylibcudf as plc + +from cudf_polars.dsl import expr +from cudf_polars.dsl.ir import ( + Filter, + GroupBy, + Join, + MapFunction, + Projection, + Select, + Slice, + Sort, +) +from cudf_polars.dsl.traversal import post_traversal +from cudf_polars.dsl.utils.column_domain import column_domain_bindings + +if TYPE_CHECKING: + from collections.abc import Mapping + + from cudf_polars.dsl.ir import IR + + +@dataclass(frozen=True) +class NamedOrderKey: + """Named sort key with logical Polars ordering options.""" + + name: str + descending: bool + nulls_last: bool + + +@dataclass(frozen=True) +class StrictPartitioningRequest: + """Request for upstream output to strictly partition equal keys.""" + + keys: tuple[str, ...] + + +@dataclass(frozen=True) +class OrderPartitioningRequest: + """Request for upstream output to be ordered by a key sequence.""" + + keys: tuple[NamedOrderKey, ...] + strict_key_count: int | None = None + + +PartitioningRequest: TypeAlias = StrictPartitioningRequest | OrderPartitioningRequest + + +def collect_partitioning_requests( + ir: IR, +) -> dict[IR, tuple[PartitioningRequest, ...]]: + """ + Collect downstream layout requests for each IR node. + + The returned mapping answers "which layouts could make downstream consumers + cheaper if this node produced them?" A request is therefore aspirational: + it is not evidence that the data is currently sorted, hash partitioned, or + otherwise laid out that way. + """ + requests: dict[IR, tuple[PartitioningRequest, ...]] = {} + # Reverse post-order ensures every downstream consumer is processed before + # any shared upstream producer whose requests must be propagated further. + for node in reversed(list(post_traversal([ir]))): + child_requests = _direct_child_requests(node) + child_requests.extend(_propagated_child_requests(node, requests)) + for child, request in child_requests: + requests[child] = _merge_candidate_request(requests.get(child, ()), request) + return requests + + +def _direct_child_requests(ir: IR) -> list[tuple[IR, PartitioningRequest]]: + """Create child requests implied by operators that consume a layout.""" + if isinstance(ir, Sort): + names = _column_names(ir.by) + if names is not None: + return [ + ( + ir.children[0], + _order_request( + names, + tuple( + order == plc.types.Order.DESCENDING for order in ir.order + ), + tuple( + (order == plc.types.Order.ASCENDING) + == (null_order == plc.types.NullOrder.AFTER) + for order, null_order in zip( + ir.order, ir.null_order, strict=True + ) + ), + ), + ) + ] + + if isinstance(ir, MapFunction) and ir.name == "hint_sorted": + return [(ir.children[0], _order_request(*ir.options))] + + if isinstance(ir, Join) and ir.options[0] != "Cross": + left_keys = _column_names(ir.left_on) + right_keys = _column_names(ir.right_on) + if left_keys is not None and right_keys is not None: + return [ + (ir.children[0], StrictPartitioningRequest(left_keys)), + (ir.children[1], StrictPartitioningRequest(right_keys)), + ] + + if isinstance(ir, GroupBy) and not ir.maintain_order: + keys = _column_names(ir.keys) + if keys is not None: + return [(ir.children[0], StrictPartitioningRequest(keys))] + + return [] + + +def _propagated_child_requests( + node: IR, requests: dict[IR, tuple[PartitioningRequest, ...]] +) -> list[tuple[IR, PartitioningRequest]]: + """Push compatible downstream requests through single-child operators.""" + child_requests: list[tuple[IR, PartitioningRequest]] = [] + node_requests = requests.get(node) + if ( + node_requests is not None + and len(node.children) == 1 + and isinstance(node, (Projection, Select, Filter, Slice, GroupBy)) + ): + remapping = { + output_name: binding.name + for output_name, binding in column_domain_bindings(node).items() + if binding.child_index == 0 + } + child_requests.extend( + (node.children[0], remapped) + for node_request in node_requests + if (remapped := _remap_request(node_request, remapping)) is not None + ) + return child_requests + + +def _merge_candidate_request( + existing_requests: tuple[PartitioningRequest, ...], + request: PartitioningRequest, +) -> tuple[PartitioningRequest, ...]: + """Merge compatible requests while preserving incompatible candidates.""" + candidates: list[PartitioningRequest] = [] + new_request = request + insertion_index: int | None = None + for existing_request in existing_requests: + if (merged := _merge_requests(existing_request, new_request)) is None: + candidates.append(existing_request) + else: + new_request = merged + if insertion_index is None: + insertion_index = len(candidates) + if insertion_index is None: + candidates.append(new_request) + else: + candidates.insert(insertion_index, new_request) + return tuple(candidates) + + +def _order_request( + names: tuple[str, ...], + descending: tuple[bool, ...], + nulls_last: tuple[bool, ...], +) -> OrderPartitioningRequest: + return OrderPartitioningRequest( + tuple( + NamedOrderKey(name, desc, null_last) + for name, desc, null_last in zip(names, descending, nulls_last, strict=True) + ) + ) + + +def _column_names(named_exprs: tuple[expr.NamedExpr, ...]) -> tuple[str, ...] | None: + names = [] + for named_expr in named_exprs: + if not isinstance(named_expr.value, expr.Col): + return None + names.append(named_expr.value.name) + return tuple(names) + + +def _remap_request( + request: PartitioningRequest, remapping: Mapping[str, str] +) -> PartitioningRequest | None: + """Rewrite request column names through a child-to-parent name mapping.""" + if isinstance(request, StrictPartitioningRequest): + remapped_names = [] + for name in request.keys: + if (new_name := remapping.get(name)) is None: + return None + remapped_names.append(new_name) + return StrictPartitioningRequest(tuple(remapped_names)) + + remapped_keys = [] + for key in request.keys: + new_name = remapping.get(key.name) + if new_name is None: + return None + remapped_keys.append(NamedOrderKey(new_name, key.descending, key.nulls_last)) + return OrderPartitioningRequest(tuple(remapped_keys), request.strict_key_count) + + +def _merge_requests( + left: PartitioningRequest, right: PartitioningRequest +) -> PartitioningRequest | None: + """Merge compatible requests, or keep both candidates if incompatible.""" + if isinstance(left, StrictPartitioningRequest): + if isinstance(right, StrictPartitioningRequest): + if _is_prefix(left.keys, right.keys): + return left + if _is_prefix(right.keys, left.keys): + return right + return None + return _merge_order_with_strict(right, left) + + if isinstance(right, StrictPartitioningRequest): + return _merge_order_with_strict(left, right) + + if _is_prefix(left.keys, right.keys): + keys = right.keys + elif _is_prefix(right.keys, left.keys): + keys = left.keys + else: + return None + return OrderPartitioningRequest( + keys, _merge_strict_key_count(left.strict_key_count, right.strict_key_count) + ) + + +def _merge_order_with_strict( + order_request: OrderPartitioningRequest, + strict_request: StrictPartitioningRequest, +) -> OrderPartitioningRequest | None: + """Fold strict-key requirements into compatible ordering requests.""" + order_names = tuple(key.name for key in order_request.keys) + if _is_prefix(strict_request.keys, order_names) or _is_prefix( + order_names, strict_request.keys + ): + strict_key_count = min(len(strict_request.keys), len(order_names)) + return OrderPartitioningRequest( + order_request.keys, + _merge_strict_key_count(order_request.strict_key_count, strict_key_count), + ) + return None + + +def _merge_strict_key_count(*counts: int | None) -> int | None: + count = max((count for count in counts if count is not None), default=0) + return count or None + + +def _is_prefix(left: tuple[object, ...], right: tuple[object, ...]) -> bool: + return len(left) <= len(right) and left == right[: len(left)] diff --git a/python/cudf_polars/tests/streaming/test_partitioning_hints.py b/python/cudf_polars/tests/streaming/test_partitioning_requests.py similarity index 69% rename from python/cudf_polars/tests/streaming/test_partitioning_hints.py rename to python/cudf_polars/tests/streaming/test_partitioning_requests.py index 79c363c67e99..7088e67ddea8 100644 --- a/python/cudf_polars/tests/streaming/test_partitioning_hints.py +++ b/python/cudf_polars/tests/streaming/test_partitioning_requests.py @@ -18,11 +18,11 @@ Sort, Union, ) -from cudf_polars.streaming.partitioning_hints import ( +from cudf_polars.streaming.partitioning_requests import ( NamedOrderKey, - OrderPartitioningHint, - StrictPartitioningHint, - collect_partitioning_hints, + OrderPartitioningRequest, + StrictPartitioningRequest, + collect_partitioning_requests, ) from cudf_polars.utils.sorting import sort_order @@ -91,7 +91,7 @@ def make_groupby(child: IR, *names: str) -> GroupBy: ) -def test_sort_creates_order_partition_hint() -> None: +def test_sort_creates_order_partition_request() -> None: scan = make_scan("a", "b") sort = make_sort( scan, @@ -101,10 +101,10 @@ def test_sort_creates_order_partition_hint() -> None: nulls_last=(True, False), ) - hints = collect_partitioning_hints(sort) + requests = collect_partitioning_requests(sort) - assert hints[scan] == ( - OrderPartitioningHint( + assert requests[scan] == ( + OrderPartitioningRequest( ( NamedOrderKey("a", descending=False, nulls_last=True), NamedOrderKey("b", descending=True, nulls_last=False), @@ -113,7 +113,7 @@ def test_sort_creates_order_partition_hint() -> None: ) -def test_select_remaps_order_partition_hint() -> None: +def test_select_remaps_order_partition_request() -> None: scan = make_scan("a", "b") select = Select( {"x": I64, "b": I64}, @@ -123,16 +123,16 @@ def test_select_remaps_order_partition_hint() -> None: ) sort = make_sort(select, "x") - hints = collect_partitioning_hints(sort) + requests = collect_partitioning_requests(sort) - assert hints[scan] == ( - OrderPartitioningHint( + assert requests[scan] == ( + OrderPartitioningRequest( (NamedOrderKey("a", descending=False, nulls_last=False),) ), ) -def test_non_column_sort_does_not_create_hint() -> None: +def test_non_column_sort_does_not_create_request() -> None: scan = make_scan("a") order, null_order = sort_order((False,), nulls_last=(False,), num_keys=1) sort = Sort( @@ -145,23 +145,25 @@ def test_non_column_sort_does_not_create_hint() -> None: df=scan, ) - hints = collect_partitioning_hints(sort) + requests = collect_partitioning_requests(sort) - assert hints == {} + assert requests == {} -def test_hint_sorted_creates_order_partition_hint() -> None: +def test_hint_sorted_creates_order_partition_request() -> None: scan = make_scan("a", "b") hint_sorted = make_hint_sorted(scan, "a", descending=(True,)) - hints = collect_partitioning_hints(hint_sorted) + requests = collect_partitioning_requests(hint_sorted) - assert hints[scan] == ( - OrderPartitioningHint((NamedOrderKey("a", descending=True, nulls_last=False),)), + assert requests[scan] == ( + OrderPartitioningRequest( + (NamedOrderKey("a", descending=True, nulls_last=False),) + ), ) -def test_select_remaps_strict_partition_hint() -> None: +def test_select_remaps_strict_partition_request() -> None: scan = make_scan("a") select = Select( {"x": I64}, @@ -179,13 +181,13 @@ def test_select_remaps_strict_partition_hint() -> None: right, ) - hints = collect_partitioning_hints(join) + requests = collect_partitioning_requests(join) - assert hints[scan] == (StrictPartitioningHint(("a",)),) - assert hints[right] == (StrictPartitioningHint(("x",)),) + assert requests[scan] == (StrictPartitioningRequest(("a",)),) + assert requests[right] == (StrictPartitioningRequest(("x",)),) -def test_select_drops_order_hint_on_non_column_output() -> None: +def test_select_drops_order_request_on_non_column_output() -> None: scan = make_scan("a") select = Select( {"x": I64}, @@ -195,17 +197,17 @@ def test_select_drops_order_hint_on_non_column_output() -> None: ) sort = make_sort(select, "x") - hints = collect_partitioning_hints(sort) + requests = collect_partitioning_requests(sort) - assert hints[select] == ( - OrderPartitioningHint( + assert requests[select] == ( + OrderPartitioningRequest( (NamedOrderKey("x", descending=False, nulls_last=False),) ), ) - assert scan not in hints + assert scan not in requests -def test_select_drops_strict_hint_on_non_column_output() -> None: +def test_select_drops_strict_request_on_non_column_output() -> None: scan = make_scan("a") select = Select( {"x": I64}, @@ -223,11 +225,11 @@ def test_select_drops_strict_hint_on_non_column_output() -> None: right, ) - hints = collect_partitioning_hints(join) + requests = collect_partitioning_requests(join) - assert hints[select] == (StrictPartitioningHint(("x",)),) - assert hints[right] == (StrictPartitioningHint(("x",)),) - assert scan not in hints + assert requests[select] == (StrictPartitioningRequest(("x",)),) + assert requests[right] == (StrictPartitioningRequest(("x",)),) + assert scan not in requests def test_hint_sorted_keeps_declared_order_with_compatible_downstream_sort() -> None: @@ -235,10 +237,10 @@ def test_hint_sorted_keeps_declared_order_with_compatible_downstream_sort() -> N hint_sorted = make_hint_sorted(scan, "a") sort = make_sort(hint_sorted, "a") - hints = collect_partitioning_hints(sort) + requests = collect_partitioning_requests(sort) - assert hints[scan] == ( - OrderPartitioningHint( + assert requests[scan] == ( + OrderPartitioningRequest( (NamedOrderKey("a", descending=False, nulls_last=False),) ), ) @@ -249,10 +251,10 @@ def test_hint_sorted_keeps_declared_order_with_extended_downstream_sort() -> Non hint_sorted = make_hint_sorted(scan, "a") sort = make_sort(hint_sorted, "a", "b") - hints = collect_partitioning_hints(sort) + requests = collect_partitioning_requests(sort) - assert hints[scan] == ( - OrderPartitioningHint( + assert requests[scan] == ( + OrderPartitioningRequest( (NamedOrderKey("a", descending=False, nulls_last=False),) ), ) @@ -263,14 +265,16 @@ def test_hint_sorted_keeps_declared_order_with_incompatible_downstream_sort() -> hint_sorted = make_hint_sorted(scan, "a", descending=(True,)) sort = make_sort(hint_sorted, "a") - hints = collect_partitioning_hints(sort) + requests = collect_partitioning_requests(sort) - assert hints[scan] == ( - OrderPartitioningHint((NamedOrderKey("a", descending=True, nulls_last=False),)), + assert requests[scan] == ( + OrderPartitioningRequest( + (NamedOrderKey("a", descending=True, nulls_last=False),) + ), ) -def test_groupby_remaps_order_partition_hint() -> None: +def test_groupby_remaps_order_partition_request() -> None: scan = make_scan("a", "b") groupby = make_groupby( Select( @@ -283,17 +287,17 @@ def test_groupby_remaps_order_partition_hint() -> None: ) sort = make_sort(groupby, "key") - hints = collect_partitioning_hints(sort) + requests = collect_partitioning_requests(sort) - assert hints[scan] == ( - OrderPartitioningHint( + assert requests[scan] == ( + OrderPartitioningRequest( (NamedOrderKey("a", descending=False, nulls_last=False),), strict_key_count=1, ), ) -def test_fanout_keeps_more_specific_compatible_order_hint() -> None: +def test_fanout_keeps_more_specific_compatible_order_request() -> None: scan = make_scan("a", "b") root = Union( scan.schema, @@ -303,10 +307,10 @@ def test_fanout_keeps_more_specific_compatible_order_hint() -> None: make_sort(scan, "a", "b"), ) - hints = collect_partitioning_hints(root) + requests = collect_partitioning_requests(root) - assert hints[scan] == ( - OrderPartitioningHint( + assert requests[scan] == ( + OrderPartitioningRequest( ( NamedOrderKey("a", descending=False, nulls_last=False), NamedOrderKey("b", descending=False, nulls_last=False), @@ -315,7 +319,7 @@ def test_fanout_keeps_more_specific_compatible_order_hint() -> None: ) -def test_fanout_marks_compatible_order_hint_as_strict() -> None: +def test_fanout_marks_compatible_order_request_as_strict() -> None: scan = make_scan("a", "b") right = make_scan("a", "right_value") join = Join( @@ -334,10 +338,10 @@ def test_fanout_marks_compatible_order_hint_as_strict() -> None: join, ) - hints = collect_partitioning_hints(root) + requests = collect_partitioning_requests(root) - assert hints[scan] == ( - OrderPartitioningHint( + assert requests[scan] == ( + OrderPartitioningRequest( ( NamedOrderKey("a", descending=False, nulls_last=False), NamedOrderKey("b", descending=False, nulls_last=False), @@ -345,10 +349,10 @@ def test_fanout_marks_compatible_order_hint_as_strict() -> None: strict_key_count=1, ), ) - assert hints[right] == (StrictPartitioningHint(("a",)),) + assert requests[right] == (StrictPartitioningRequest(("a",)),) -def test_fanout_merges_compatible_strict_hints() -> None: +def test_fanout_merges_compatible_strict_requests() -> None: scan = make_scan("a", "b") root = Union( scan.schema, @@ -358,9 +362,9 @@ def test_fanout_merges_compatible_strict_hints() -> None: make_groupby(scan, "a", "b"), ) - hints = collect_partitioning_hints(root) + requests = collect_partitioning_requests(root) - assert hints[scan] == (StrictPartitioningHint(("a",)),) + assert requests[scan] == (StrictPartitioningRequest(("a",)),) def test_fanout_keeps_incompatible_strict_candidates() -> None: @@ -373,11 +377,11 @@ def test_fanout_keeps_incompatible_strict_candidates() -> None: make_groupby(scan, "b"), ) - hints = collect_partitioning_hints(root) + requests = collect_partitioning_requests(root) - assert set(hints[scan]) == { - StrictPartitioningHint(("a",)), - StrictPartitioningHint(("b",)), + assert set(requests[scan]) == { + StrictPartitioningRequest(("a",)), + StrictPartitioningRequest(("b",)), } @@ -391,17 +395,17 @@ def test_fanout_keeps_incompatible_order_and_strict_candidates() -> None: make_groupby(scan, "b"), ) - hints = collect_partitioning_hints(root) + requests = collect_partitioning_requests(root) - assert set(hints[scan]) == { - OrderPartitioningHint( + assert set(requests[scan]) == { + OrderPartitioningRequest( (NamedOrderKey("a", descending=False, nulls_last=False),) ), - StrictPartitioningHint(("b",)), + StrictPartitioningRequest(("b",)), } -def test_compatible_hint_merging_is_not_directional() -> None: +def test_compatible_request_merging_is_not_directional() -> None: scan = make_scan("a", "b") root = Union( scan.schema, @@ -410,7 +414,9 @@ def test_compatible_hint_merging_is_not_directional() -> None: make_groupby(scan, "a", "b"), make_groupby(scan, "a"), ) - assert collect_partitioning_hints(root)[scan] == (StrictPartitioningHint(("a",)),) + assert collect_partitioning_requests(root)[scan] == ( + StrictPartitioningRequest(("a",)), + ) scan = make_scan("a", "b") root = Union( @@ -420,9 +426,9 @@ def test_compatible_hint_merging_is_not_directional() -> None: make_groupby(scan, "b"), make_sort(scan, "a"), ) - assert set(collect_partitioning_hints(root)[scan]) == { - StrictPartitioningHint(("b",)), - OrderPartitioningHint( + assert set(collect_partitioning_requests(root)[scan]) == { + StrictPartitioningRequest(("b",)), + OrderPartitioningRequest( (NamedOrderKey("a", descending=False, nulls_last=False),) ), } @@ -435,8 +441,8 @@ def test_compatible_hint_merging_is_not_directional() -> None: make_sort(scan, "a", "b"), make_sort(scan, "a"), ) - assert collect_partitioning_hints(root)[scan] == ( - OrderPartitioningHint( + assert collect_partitioning_requests(root)[scan] == ( + OrderPartitioningRequest( ( NamedOrderKey("a", descending=False, nulls_last=False), NamedOrderKey("b", descending=False, nulls_last=False), @@ -445,7 +451,7 @@ def test_compatible_hint_merging_is_not_directional() -> None: ) -def test_conflicting_fanout_keeps_candidate_hints() -> None: +def test_conflicting_fanout_keeps_candidate_requests() -> None: scan = make_scan("a", "b") root = Union( scan.schema, @@ -455,13 +461,13 @@ def test_conflicting_fanout_keeps_candidate_hints() -> None: make_sort(scan, "b"), ) - hints = collect_partitioning_hints(root) + requests = collect_partitioning_requests(root) - assert set(hints[scan]) == { - OrderPartitioningHint( + assert set(requests[scan]) == { + OrderPartitioningRequest( (NamedOrderKey("a", descending=False, nulls_last=False),) ), - OrderPartitioningHint( + OrderPartitioningRequest( (NamedOrderKey("b", descending=False, nulls_last=False),) ), } @@ -478,14 +484,14 @@ def test_repeated_fanout_candidate_is_merged() -> None: make_sort(scan, "a"), ) - hints = collect_partitioning_hints(root) + requests = collect_partitioning_requests(root) - assert len(hints[scan]) == 2 - assert set(hints[scan]) == { - OrderPartitioningHint( + assert len(requests[scan]) == 2 + assert set(requests[scan]) == { + OrderPartitioningRequest( (NamedOrderKey("a", descending=False, nulls_last=False),) ), - OrderPartitioningHint( + OrderPartitioningRequest( (NamedOrderKey("b", descending=False, nulls_last=False),) ), } From ccaed779f2a72392200177a7074ba8d937953edb Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 21 Aug 2026 09:45:07 -0700 Subject: [PATCH 11/12] address cr comment --- .../streaming/test_partitioning_requests.py | 30 +++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/tests/streaming/test_partitioning_requests.py b/python/cudf_polars/tests/streaming/test_partitioning_requests.py index 7088e67ddea8..77d5256f4086 100644 --- a/python/cudf_polars/tests/streaming/test_partitioning_requests.py +++ b/python/cudf_polars/tests/streaming/test_partitioning_requests.py @@ -80,12 +80,12 @@ def make_hint_sorted( ) -def make_groupby(child: IR, *names: str) -> GroupBy: +def make_groupby(child: IR, *names: str, maintain_order: bool = False) -> GroupBy: return GroupBy( dict.fromkeys(names, I64), tuple(named_col(name) for name in names), (), - maintain_order=False, + maintain_order=maintain_order, zlice=None, df=child, ) @@ -187,6 +187,32 @@ def test_select_remaps_strict_partition_request() -> None: assert requests[right] == (StrictPartitioningRequest(("x",)),) +def test_cross_join_does_not_create_strict_partition_request() -> None: + left = make_scan("a") + right = make_scan("x") + join = Join( + {"a": I64, "x": I64}, + (), + (), + ("Cross", False, None, "_right", True, "none"), + left, + right, + ) + + requests = collect_partitioning_requests(join) + + assert requests == {} + + +def test_maintain_order_groupby_does_not_create_strict_partition_request() -> None: + scan = make_scan("a") + groupby = make_groupby(scan, "a", maintain_order=True) + + requests = collect_partitioning_requests(groupby) + + assert requests == {} + + def test_select_drops_order_request_on_non_column_output() -> None: scan = make_scan("a") select = Select( From 3833b48f4a342e61be143efb564b698e9abea0a0 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 21 Aug 2026 09:52:00 -0700 Subject: [PATCH 12/12] tweak language again --- .../streaming/partitioning_requests.py | 23 +++++++++++-------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/partitioning_requests.py b/python/cudf_polars/cudf_polars/streaming/partitioning_requests.py index db948e67a494..aaf6e79ff358 100644 --- a/python/cudf_polars/cudf_polars/streaming/partitioning_requests.py +++ b/python/cudf_polars/cudf_polars/streaming/partitioning_requests.py @@ -4,9 +4,14 @@ Downstream partitioning requests for streaming actor-graph construction. A partitioning request is attached to an IR node when a downstream consumer -may benefit if that node produces data with a specific layout. These requests -do not describe or guarantee the current layout of the node's output; actual -runtime layout metadata is tracked separately in ``ChannelMetadata``. +may benefit from that node producing a specific partitioning. These requests +use "partitioning" in the same broad sense as ``ChannelMetadata.Partitioning``: +rows may be strictly partitioned by equality keys, ordered by key values, or +both. + +Requests are planning-time information. They do not describe or guarantee the +actual partitioning of the node's output. Runtime partitioning metadata is +tracked separately in ``ChannelMetadata``. """ from __future__ import annotations @@ -67,12 +72,12 @@ def collect_partitioning_requests( ir: IR, ) -> dict[IR, tuple[PartitioningRequest, ...]]: """ - Collect downstream layout requests for each IR node. + Collect downstream partitioning requests for each IR node. - The returned mapping answers "which layouts could make downstream consumers - cheaper if this node produced them?" A request is therefore aspirational: - it is not evidence that the data is currently sorted, hash partitioned, or - otherwise laid out that way. + The returned mapping answers "which partitionings could make downstream + consumers cheaper if this node produced them?" A request is therefore + aspirational: it is not evidence that the data is currently sorted, hash + partitioned, or otherwise partitioned that way. """ requests: dict[IR, tuple[PartitioningRequest, ...]] = {} # Reverse post-order ensures every downstream consumer is processed before @@ -86,7 +91,7 @@ def collect_partitioning_requests( def _direct_child_requests(ir: IR) -> list[tuple[IR, PartitioningRequest]]: - """Create child requests implied by operators that consume a layout.""" + """Create child requests implied by partitioning-aware operators.""" if isinstance(ir, Sort): names = _column_names(ir.by) if names is not None: