diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py b/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py index 98532df2a878..f21451e702ee 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py @@ -615,13 +615,17 @@ def _grouped_window_scan_setup( local = self._sorted_grouper(by_cols_for_scan) return order_index, by_cols_for_scan, local + # TODO: this is an ordered left-join that drops the join keys. + # Rename it and replace the manual scatter+gather with the Join IR's + # _reorder_maps helper (lifted to a shared utility) so the streaming + # and in-memory over paths share the same primitive. def _broadcast_agg_results( self, by_tbl: plc.Table, group_keys_tbl: plc.Table, - value_tbls: list[plc.Table], - names: list[str], - dtypes: list[DataType], + value_tbl: plc.Table, + names: Sequence[str], + dtypes: Sequence[DataType], stream: Stream, ) -> list[Column]: # We do a left-join between the input keys to group-keys @@ -647,7 +651,7 @@ def _broadcast_agg_results( # Broadcast each scalar aggregated result back to row-shape using # the aligned mapping between row indices and group indices. - out_cols = (t.columns()[0] for t in value_tbls) + out_cols = value_tbl.columns() return [ Column( plc.copying.gather( @@ -773,7 +777,7 @@ def do_evaluate( # noqa: D102 broadcasted_cols = self._broadcast_agg_results( by_tbl, group_keys_tbl, - value_tables, + plc.Table([col for t in value_tables for col in t.columns()]), out_names, out_dtypes, df.stream, @@ -807,7 +811,7 @@ def do_evaluate( # noqa: D102 self._broadcast_agg_results( by_tbl, group_keys_tbl_local, - value_tables_local, + plc.Table([col for t in value_tables_local for col in t.columns()]), out_names, out_dtypes, df.stream, diff --git a/python/cudf_polars/cudf_polars/dsl/translate.py b/python/cudf_polars/cudf_polars/dsl/translate.py index 47afd9b7e58e..8f68b259954c 100644 --- a/python/cudf_polars/cudf_polars/dsl/translate.py +++ b/python/cudf_polars/cudf_polars/dsl/translate.py @@ -38,6 +38,8 @@ ) if TYPE_CHECKING: + from collections.abc import Generator + from polars import GPUEngine from cudf_polars.typing import NodeTraverser @@ -90,6 +92,7 @@ def __init__(self, visitor: NodeTraverser, engine: GPUEngine): self.errors: list[Exception] = [] self._cache_nodes: dict[int, ir.Cache] = {} self._expr_context: ExecutionContext = ExecutionContext.FRAME + self._internal_name_gen: Generator[str, None, None] | None = None def translate_ir(self, *, n: int | None = None) -> ir.IR: """ @@ -247,6 +250,24 @@ def __exit__(self, *args: Any) -> None: self.translator._expr_context = self._prev +class set_internal_name_gen(AbstractContextManager[None]): + """Share one internal-name generator across sibling expression translations.""" + + __slots__ = ("_prev", "schema", "translator") + + def __init__(self, translator: Translator, schema: Schema) -> None: + self.translator = translator + self.schema = schema + self._prev: Generator[str, None, None] | None = None + + def __enter__(self) -> None: + self._prev = self.translator._internal_name_gen + self.translator._internal_name_gen = unique_names(self.schema) + + def __exit__(self, *args: Any) -> None: + self.translator._internal_name_gen = self._prev + + @singledispatch def _translate_ir(node: Any, translator: Translator, schema: Schema) -> ir.IR: raise NotImplementedError( @@ -362,9 +383,11 @@ def _( def _(node: plrs._ir_nodes.Select, translator: Translator, schema: Schema) -> ir.IR: with set_node(translator.visitor, node.input): inp = translator.translate_ir(n=None) - exprs = [ - translate_named_expr(translator, n=e, schema=inp.schema) for e in node.expr - ] + with set_internal_name_gen(translator, inp.schema): + exprs = [ + translate_named_expr(translator, n=e, schema=inp.schema) + for e in node.expr + ] return ir.Select(schema, exprs, node.should_broadcast, inp) @@ -478,9 +501,11 @@ def _(node: plrs._ir_nodes.Join, translator: Translator, schema: Schema) -> ir.I def _(node: plrs._ir_nodes.HStack, translator: Translator, schema: Schema) -> ir.IR: with set_node(translator.visitor, node.input): inp = translator.translate_ir(n=None) - exprs = [ - translate_named_expr(translator, n=e, schema=inp.schema) for e in node.exprs - ] + with set_internal_name_gen(translator, inp.schema): + exprs = [ + translate_named_expr(translator, n=e, schema=inp.schema) + for e in node.exprs + ] return ir.HStack(schema, exprs, node.should_broadcast, inp) @@ -830,7 +855,7 @@ def _( # pl.col("a").rolling(...) with set_expr_context(translator, ExecutionContext.ROLLING): agg = translator.translate_expr(n=node.function, schema=schema) - name_generator = unique_names(schema) + name_generator = translator._internal_name_gen or unique_names(schema) aggs, named_post_agg = decompose_single_agg( expr.NamedExpr(next(name_generator), agg), name_generator, @@ -875,7 +900,7 @@ def _( # not exposed until polars 1.39. with set_expr_context(translator, ExecutionContext.WINDOW): agg = translator.translate_expr(n=node.function, schema=schema) - name_gen = unique_names(schema) + name_gen = translator._internal_name_gen or unique_names(schema) aggs, post = decompose_single_agg( expr.NamedExpr(next(name_gen), agg), name_gen, diff --git a/python/cudf_polars/cudf_polars/dsl/utils/naming.py b/python/cudf_polars/cudf_polars/dsl/utils/naming.py index 65eedbb14955..018477f91775 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/naming.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/naming.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 """Name generation utilities.""" @@ -7,11 +7,15 @@ from typing import TYPE_CHECKING +from cudf_polars.dsl.expr import NamedExpr + if TYPE_CHECKING: from collections.abc import Generator, Iterable + from cudf_polars.typing import Schema + -__all__ = ["unique_names"] +__all__ = ["names_to_indices", "unique_names"] def unique_names(names: Iterable[str]) -> Generator[str, None, None]: @@ -32,3 +36,28 @@ def unique_names(names: Iterable[str]) -> Generator[str, None, None]: while True: yield f"{prefix}{i}" i += 1 + + +def names_to_indices( + names: tuple[str | NamedExpr, ...], schema: Schema +) -> tuple[int, ...]: + """ + Return column indices for the given names in schema order. + + Accepts either column names (str) or NamedExpr, so it can be used with + e.g. ir.left_on, ir.right_on as well as plain name tuples. + + Parameters + ---------- + names + The names to get indices for. + schema + The schema to get indices from. + + Returns + ------- + The column indices for each name in schema order. + """ + keys = list(schema.keys()) + str_names = [n.name if isinstance(n, NamedExpr) else n for n in names] + return tuple(keys.index(n) for n in str_names) diff --git a/python/cudf_polars/cudf_polars/experimental/expressions.py b/python/cudf_polars/cudf_polars/experimental/expressions.py index f578e3b0e071..ef94da02dea1 100644 --- a/python/cudf_polars/cudf_polars/experimental/expressions.py +++ b/python/cudf_polars/cudf_polars/experimental/expressions.py @@ -44,6 +44,7 @@ from cudf_polars.dsl.expressions.base import Col, ExecutionContext, NamedExpr from cudf_polars.dsl.expressions.binaryop import BinOp from cudf_polars.dsl.expressions.literal import Literal +from cudf_polars.dsl.expressions.rolling import GroupedWindow from cudf_polars.dsl.expressions.ternary import Ternary from cudf_polars.dsl.expressions.unary import Cast, Len, UnaryFunction from cudf_polars.dsl.ir import Distinct, Empty, HConcat, Select @@ -51,6 +52,7 @@ CachingVisitor, ) from cudf_polars.experimental.base import PartitionInfo +from cudf_polars.experimental.over import _decompose_grouped_window_node from cudf_polars.experimental.repartition import Repartition from cudf_polars.experimental.utils import _dynamic_planning_on @@ -462,6 +464,10 @@ def _decompose_expr_node( ) (expr,) = columns return expr, input_ir, partition_info + elif isinstance(expr, GroupedWindow) and _dynamic_planning_on(config_options): + return _decompose_grouped_window_node( + expr, input_ir, partition_info, config_options, names=names + ) else: # This is an un-supported expression - raise. raise NotImplementedError( diff --git a/python/cudf_polars/cudf_polars/experimental/over.py b/python/cudf_polars/cudf_polars/experimental/over.py new file mode 100644 index 000000000000..abc90d358100 --- /dev/null +++ b/python/cudf_polars/cudf_polars/experimental/over.py @@ -0,0 +1,337 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Over IR node for streaming window expressions.""" + +from __future__ import annotations + +import itertools +from collections import defaultdict +from typing import TYPE_CHECKING, ClassVar, cast + +from cudf_polars.dsl.expr import Agg, Col, Len, NamedExpr +from cudf_polars.dsl.ir import IR, GroupBy, Select +from cudf_polars.dsl.utils.naming import names_to_indices, unique_names +from cudf_polars.experimental.groupby import combine, decompose + +if TYPE_CHECKING: + from collections.abc import Generator, MutableMapping + + from cudf_polars.containers import DataFrame + from cudf_polars.dsl.expr import GroupedWindow + from cudf_polars.dsl.expressions.base import Expr + from cudf_polars.dsl.ir import IRExecutionContext + from cudf_polars.experimental.base import PartitionInfo + from cudf_polars.typing import Schema + from cudf_polars.utils.config import ConfigOptions + + +# Aggregations whose partial results can be combined. +_DECOMPOSABLE_AGG_NAMES: frozenset[str] = frozenset( + ("sum", "count", "mean", "min", "max", "std", "var") +) + + +def _build_over_groupby_irs( + gw_nodes: tuple[GroupedWindow, ...], + child_ir: IR, +) -> tuple[GroupBy, GroupBy, Select]: + """ + Build piecewise, reduction, and selection GroupBy IRs. + + Parameters + ---------- + gw_nodes + Top-level GroupedWindow nodes sharing the same partition-by keys; + all must be scalar (Agg/Len only in named_aggs). + child_ir + Input IR feeding the Over node; defines the schema seen by the + per-chunk piecewise GroupBy. + + Returns + ------- + piecewise_ir + GroupBy IR that computes partial aggregates per chunk. + reduction_ir + GroupBy IR that reduces partial aggregates to a single result. + agg_select_ir + Select IR applied on top of the reduction. Carries any post- + aggregation expressions (e.g. division for mean); for fully + pass-through aggregations it is a Select of plain ``Col`` refs + of the same shape as the reduction output. + """ + gw = gw_nodes[0] + by_exprs = cast("list[Col]", list(gw.children[: gw.by_count])) + key_named_exprs = [NamedExpr(e.name, e) for e in by_exprs] + key_schema = {e.name: child_ir.schema[e.name] for e in by_exprs} + + all_scalar_named: list[NamedExpr] = [] + seen: set[str] = set() + for gw_node in gw_nodes: + reductions, unary_ops = gw_node._split_named_expr() + assert not any(unary_ops.values()), "unary window ops not allowed here" + for ne in reductions: + if ne.name in seen: + continue + all_scalar_named.append(ne) + seen.add(ne.name) + + name_gen = unique_names(child_ir.schema.keys()) + decompositions = [ + decompose(ne.name, ne.value, names=name_gen) for ne in all_scalar_named + ] + selection_exprs, piecewise_exprs, reduction_exprs, need_preshuffle = combine( + *decompositions + ) + assert not need_preshuffle, ( + "Scalar AllGather path does not support aggregations requiring pre-shuffle" + ) + + pwise_schema = dict(key_schema) | { + ne.name: ne.value.dtype for ne in piecewise_exprs + } + piecewise_ir = GroupBy( + pwise_schema, + key_named_exprs, + piecewise_exprs, + False, # noqa: FBT003 + None, + child_ir, + ) + + reduction_key_exprs = [ + NamedExpr(ne.name, Col(pwise_schema[ne.name], ne.name)) + for ne in key_named_exprs + ] + reduction_schema = { + ne.name: ne.value.dtype + for ne in itertools.chain(reduction_key_exprs, reduction_exprs) + } + reduction_ir = GroupBy( + reduction_schema, + reduction_key_exprs, + reduction_exprs, + False, # noqa: FBT003 + None, + piecewise_ir, + ) + + select_key_exprs = [ + NamedExpr(ne.name, Col(reduction_schema[ne.name], ne.name)) + for ne in key_named_exprs + ] + select_schema = { + ne.name: ne.value.dtype + for ne in itertools.chain(select_key_exprs, selection_exprs) + } + agg_select_ir = Select( + select_schema, + [*select_key_exprs, *selection_exprs], + False, # noqa: FBT003 + reduction_ir, + ) + + return piecewise_ir, reduction_ir, agg_select_ir + + +class Over(IR): + """Window over() IR node for the streaming runtime.""" + + __slots__ = ("exprs", "is_scalar", "key_indices") + _non_child: ClassVar[tuple[str, ...]] = ( + "schema", + "key_indices", + "is_scalar", + "exprs", + ) + _n_non_child_args: ClassVar[int] = 1 + key_indices: tuple[int, ...] + is_scalar: bool + exprs: tuple[NamedExpr, ...] + + def __init__( + self, + schema: Schema, + key_indices: tuple[int, ...], + is_scalar: bool, # noqa: FBT001 + exprs: tuple[NamedExpr, ...], + input_ir: IR, + ): + assert len(key_indices) > 0, "Over node requires at least one partition-by key" + self.schema = schema + self.key_indices = key_indices + self.is_scalar = is_scalar + self.exprs = exprs + self._non_child_args = (exprs,) + self.children = (input_ir,) + + @classmethod + def do_evaluate( + cls, + exprs: tuple[NamedExpr, ...], + df: DataFrame, + *, + context: IRExecutionContext, + ) -> DataFrame: + """Evaluate window expressions against df.""" + # At evaluation time Over is just a Select with should_broadcast=True; + # the window-specific work lives in the GroupedWindow expressions. + return Select.do_evaluate(exprs, True, df, context=context) # noqa: FBT003 + + +def _is_scalar_grouped_window(expr: GroupedWindow) -> bool: + """Return True if this GroupedWindow can use the scalar broadcast path.""" + reductions, unary_ops = expr._split_named_expr() + if any(unary_ops.values()): + return False + if not all(isinstance(c, Col) for c in expr.children[: expr.by_count]): + return False + return all( + isinstance(ne.value, Len) + or (isinstance(ne.value, Agg) and ne.value.name in _DECOMPOSABLE_AGG_NAMES) + for ne in reductions + ) + + +def _extract_over_shuffle_indices( + expr: GroupedWindow, child_schema: Schema +) -> tuple[int, ...] | None: + """ + Return partition-by column indices in ``child_schema``, or None. + + Returns None when any partition-by expression is not a plain column + reference (the multi-partition path only supports Col keys today). + """ + by_children = expr.children[: expr.by_count] + if not all(isinstance(c, Col) for c in by_children): + return None + return names_to_indices( + tuple(cast("Col", c).name for c in by_children), child_schema + ) + + +def _decompose_grouped_window_node( + expr: GroupedWindow, + input_ir: IR, + partition_info: MutableMapping[IR, PartitionInfo], + config_options: ConfigOptions, + *, + names: Generator[str, None, None], +) -> tuple[Expr, IR, MutableMapping[IR, PartitionInfo]]: + """ + Build an Over IR node wrapping a single GroupedWindow expression. + + Every GroupedWindow becomes its own Over here; co-keyed Overs are + fused together later by select fusion so the actor evaluates all + window expressions in one pass. + + Returns + ------- + Expr + A ``Col`` referencing the Over node's output column, suitable + for substitution into the enclosing expression. + IR + The new ``Over`` IR node. + MutableMapping[IR, PartitionInfo] + ``partition_info`` augmented with an entry for the new node. + """ + indices = _extract_over_shuffle_indices(expr, input_ir.schema) + if indices is None: + # TODO: support non-Col partition-by keys on the multi-partition + # paths. Today the hash shuffle layer rejects expression keys, and + # the scalar-aggregation broadcast path builds its piecewise + # groupby from Col by-children directly. Supporting expression + # keys would require lowering them to columns in the input first. + raise NotImplementedError( + "GroupedWindow with non-Col partition-by keys " + "is not supported for multiple partitions." + ) + is_scalar = _is_scalar_grouped_window(expr) + col_name = next(names) + over_node = Over( + {col_name: expr.dtype}, + indices, + is_scalar, + (NamedExpr(col_name, expr),), + input_ir, + ) + partition_info[over_node] = partition_info[input_ir] + return Col(expr.dtype, col_name), over_node, partition_info + + +def _fuse_over_nodes( + selections: list[Select], + partition_info: MutableMapping[IR, PartitionInfo], +) -> tuple[list[Select], MutableMapping[IR, PartitionInfo]]: + """ + Fuse per-expression Over nodes that share the same grouping key. + + Selects sharing the Over's input IR are absorbed into the merged Over + so the actor produces the full output schema in one shuffle pass. The + grouping key is ``(key_indices, is_scalar, input_ir)``. + + Returns + ------- + list[Select] + The rewritten selections: one merged ``Select`` per Over group, + followed by any selections that were neither part of an Over + group nor absorbed into one. + MutableMapping[IR, PartitionInfo] + ``partition_info`` augmented with entries for the merged Over + nodes and merged Select nodes introduced by the rewrite. + """ + over_groups: defaultdict[ + tuple[tuple[int, ...], bool, IR], list[tuple[Select, Over]] + ] = defaultdict(list) + passthrough: list[Select] = [] + + for sel in selections: + child = sel.children[0] + if isinstance(child, Over): + input_ir = child.children[0] + over_groups[(child.key_indices, child.is_scalar, input_ir)].append( + (sel, child) + ) + else: + passthrough.append(sel) + + if not over_groups: + return selections, partition_info + + result: list[Select] = [] + for (key_indices, is_scalar, input_ir), group in over_groups.items(): + pi = partition_info[group[0][1]] + + absorbed: list[Select] = [] + remaining: list[Select] = [] + for s in passthrough: + (absorbed if s.children[0] == input_ir else remaining).append(s) + passthrough = remaining + + over_exprs = tuple( + itertools.chain( + *(s.exprs for s in absorbed), + *(over.exprs for _, over in group), + ) + ) + merged_over = Over( + {ne.name: ne.value.dtype for ne in over_exprs}, + key_indices, + is_scalar, + over_exprs, + input_ir, + ) + partition_info[merged_over] = pi + this_group = {*absorbed, *(sel for sel, _ in group)} + outer_exprs = tuple( + itertools.chain.from_iterable( + s.exprs for s in selections if s in this_group + ) + ) + outer_schema = {ne.name: ne.value.dtype for ne in outer_exprs} + + merged_sel = Select(outer_schema, outer_exprs, True, merged_over) # noqa: FBT003 + partition_info[merged_sel] = pi + result.append(merged_sel) + + result.extend(passthrough) + return result, partition_info diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/__init__.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/__init__.py index 2ea9af50c1ba..7eedbf12bd4e 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/__init__.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/__init__.py @@ -13,6 +13,7 @@ import cudf_polars.experimental.rapidsmpf.groupby import cudf_polars.experimental.rapidsmpf.io import cudf_polars.experimental.rapidsmpf.join +import cudf_polars.experimental.rapidsmpf.over import cudf_polars.experimental.rapidsmpf.repartition import cudf_polars.experimental.rapidsmpf.union # noqa: F401 diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/common.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/common.py index 0cb26c3689ff..fa5cb995ac55 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/common.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/common.py @@ -14,6 +14,7 @@ from cudf_polars.dsl.traversal import traversal from cudf_polars.experimental.io import StreamingSink from cudf_polars.experimental.join import Join +from cudf_polars.experimental.over import Over from cudf_polars.experimental.repartition import Repartition from cudf_polars.experimental.shuffle import Shuffle @@ -101,6 +102,7 @@ def __init__( Sort, GroupBy, Distinct, + Over, ) self.collective_nodes: list[IR] = [ @@ -150,6 +152,16 @@ def __enter__(self) -> dict[IR, list[int]]: _get_new_collective_id(), _get_new_collective_id(), ] + elif isinstance(node, Over) and not node.is_scalar: + # Non-scalar Over needs 2 IDs: one for the size AllGather + + # forward shuffle (the AllGather completes before the forward + # shuffle starts, so they can share), and a separate ID for + # the return shuffle (which overlaps with the forward shuffle + # during extract+insert). + self.collective_id_map[node] = [ + _get_new_collective_id(), + _get_new_collective_id(), + ] else: self.collective_id_map[node] = [_get_new_collective_id()] diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/sort.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/sort.py index 324534bafdf2..3bcc2517c4b0 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/sort.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/sort.py @@ -4,7 +4,6 @@ from __future__ import annotations -from collections import deque from typing import TYPE_CHECKING from rapidsmpf.shuffler import PartitionAssignment @@ -25,7 +24,7 @@ from cudf_polars.containers import DataFrame, DataType from cudf_polars.dsl.expr import Col, NamedExpr from cudf_polars.dsl.ir import Empty, Sort -from cudf_polars.dsl.utils.naming import unique_names +from cudf_polars.dsl.utils.naming import names_to_indices, unique_names from cudf_polars.experimental.rapidsmpf.collectives.allgather import AllGatherManager from cudf_polars.experimental.rapidsmpf.collectives.shuffle import ShuffleManager from cudf_polars.experimental.rapidsmpf.dispatch import generate_ir_sub_network @@ -35,6 +34,7 @@ ) from cudf_polars.experimental.rapidsmpf.utils import ( ChannelManager, + ChunkStore, NormalizedPartitioning, allgather_reduce, chunk_to_frame, @@ -44,7 +44,6 @@ evaluate_batch, evaluate_chunk, gather_in_task_group, - names_to_indices, process_children, recv_metadata, replay_buffered_channel, @@ -60,8 +59,6 @@ from cudf_polars.utils.cuda_stream import get_joined_cuda_stream if TYPE_CHECKING: - from collections.abc import Generator - from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.context import Context @@ -73,23 +70,6 @@ from cudf_polars.utils.config import StreamingExecutor -class ChunkStore: - """Ordered spillable buffer for TableChunk messages.""" - - def __init__(self, ctx: Context) -> None: - self._mids: deque[int] = deque() - self._store = ctx.spillable_messages() - - def insert(self, msg: Message) -> None: - """Insert a message into the store.""" - self._mids.append(self._store.insert(msg)) - - def __iter__(self) -> Generator[Message, None, None]: - """Yield messages in insertion order, draining the store.""" - while self._mids: - yield self._store.extract(mid=self._mids.popleft()) - - async def _simple_top_or_bottom_k( context: Context, comm: Communicator, diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py index 539ff10e7b69..8ec9c9faec5d 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py @@ -19,6 +19,7 @@ Union, ) from cudf_polars.dsl.traversal import CachingVisitor, traversal +from cudf_polars.experimental.over import Over from cudf_polars.experimental.rapidsmpf.dispatch import FanoutInfo from cudf_polars.experimental.rapidsmpf.nodes import ( generate_ir_sub_network_wrapper, @@ -169,12 +170,11 @@ def _mark_children_unbounded(node: IR) -> None: for node in traversal([ir]): if node in unbounded: _mark_children_unbounded(node) - elif isinstance(node, Union): - # Union processes children sequentially, so all children - # with multiple consumers need unbounded fanout - _mark_children_unbounded(node) - elif isinstance(node, Join): - # This may be a broadcast join + elif isinstance(node, (Union, Join, Over)): + # Union processes children sequentially; Join may broadcast one + # side; Over buffers (or samples-then-replays) its input before + # producing output. In every case the input source needs + # unbounded fanout so other consumers don't block it. _mark_children_unbounded(node) elif len(node.children) > 1: # Check if this node is doing any broadcasting. diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/groupby.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/groupby.py index 5ca71f8f1a07..243255738e3a 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/groupby.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/groupby.py @@ -15,7 +15,6 @@ from rapidsmpf.streaming.cudf.channel_metadata import ( ChannelMetadata, HashScheme, - Partitioning, ) from rapidsmpf.streaming.cudf.table_chunk import TableChunk @@ -26,7 +25,6 @@ from cudf_polars.dsl.ir import IR, Distinct, GroupBy, Select from cudf_polars.dsl.utils.naming import unique_names from cudf_polars.experimental.groupby import combine, decompose -from cudf_polars.experimental.rapidsmpf.collectives.allgather import AllGatherManager from cudf_polars.experimental.rapidsmpf.collectives.shuffle import ShuffleManager from cudf_polars.experimental.rapidsmpf.dispatch import ( generate_ir_sub_network, @@ -34,6 +32,8 @@ from cudf_polars.experimental.rapidsmpf.utils import ( ChannelManager, NormalizedPartitioning, + _make_hash_shuffle_metadata, + allgather_and_reduce, allgather_reduce, chunkwise_evaluate, empty_table_chunk, @@ -308,26 +308,13 @@ async def _tree_reduce( await send_metadata(ch_out, context, metadata_out) if need_allgather: - allgather = AllGatherManager(context, comm, collective_id) - with allgather.inserting() as inserter: - inserter.insert( - 0, - _enforce_schema( - aggregated, decomposed.reduction_ir.schema, context.br() - ), - ) - - stream = ir_context.get_cuda_stream() - aggregated = await evaluate_chunk( + aggregated = await allgather_and_reduce( context, - TableChunk.from_pylibcudf_table( - await allgather.extract_concatenated(stream), - stream, - exclusive_view=True, - br=context.br(), - ), + comm, + collective_id, + _enforce_schema(aggregated, decomposed.reduction_ir.schema, context.br()), decomposed.reduction_ir, - ir_context=ir_context, + ir_context, ) if decomposed.select_ir is not None: @@ -401,31 +388,9 @@ async def _shuffle_reduce( options = Options(get_environment_variables()) shuffle_comm = single_comm(options, comm.progress_thread) shuffle_context = Context(shuffle_comm.logger, context.br(), options) - shuf_nranks = shuffle_comm.nranks - shuf_rank = shuffle_comm.rank - modulus = max(shuf_nranks, modulus) - - if shuf_nranks == 1: - inter_rank_scheme = ( - None - if metadata_in.partitioning is None - else metadata_in.partitioning.inter_rank - ) - local_scheme = HashScheme( - column_indices=decomposed.output_indices, modulus=modulus - ) - local_output_count = modulus - else: - inter_rank_scheme = HashScheme( - column_indices=decomposed.output_indices, modulus=modulus - ) - local_scheme = "inherit" - local_output_count = (modulus - shuf_rank + shuf_nranks - 1) // shuf_nranks - - metadata_out = ChannelMetadata( - local_count=local_output_count, - partitioning=Partitioning(inter_rank_scheme, local_scheme), - duplicated=metadata_in.duplicated, + modulus = max(shuffle_comm.nranks, modulus) + metadata_out = _make_hash_shuffle_metadata( + shuffle_comm, decomposed.output_indices, modulus, metadata_in ) await send_metadata(ch_out, context, metadata_out) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py index 17d1d924a3fe..6a9815678d04 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py @@ -5,7 +5,7 @@ from __future__ import annotations import asyncio -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal from rapidsmpf.memory.memory_reservation import opaque_memory_usage @@ -30,6 +30,7 @@ from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import IR, Join +from cudf_polars.dsl.utils.naming import names_to_indices from cudf_polars.experimental.rapidsmpf.collectives.allgather import AllGatherManager from cudf_polars.experimental.rapidsmpf.collectives.shuffle import _global_shuffle from cudf_polars.experimental.rapidsmpf.dispatch import ( @@ -39,13 +40,14 @@ from cudf_polars.experimental.rapidsmpf.utils import ( ChannelManager, NormalizedPartitioning, + TableSizeStats, _is_already_partitioned, + _sample_chunks, allgather_reduce, chunk_to_frame, empty_table_chunk, gather_in_task_group, maybe_remap_partitioning, - names_to_indices, process_children, recv_metadata, replay_buffered_channel, @@ -76,20 +78,6 @@ MAX_BROADCAST_ROWS = CUDF_ROW_LIMIT // 2 -@dataclass(frozen=True) -class JoinSideStats: - """Sampled chunks and aggregate size/row stats for one side of a join.""" - - chunks: dict[int, TableChunk] = field(default_factory=dict) - """The sampled chunks, keyed by sequence number.""" - total_size: int = 0 - """The total estimated size of the child table.""" - total_rows: int = 0 - """The total estimated number of rows in the child table.""" - total_chunks: int = 0 - """The total estimated number of chunks in the child table.""" - - @dataclass(frozen=True) class JoinStrategy: """Summary of sampling and strategy selection for a dynamic join.""" @@ -855,10 +843,10 @@ def _num_indices(partitioning: NormalizedPartitioning) -> int: async def _aggregate_estimates( context: Context, comm: Communicator, - left_sample: JoinSideStats, - right_sample: JoinSideStats, + left_sample: TableSizeStats, + right_sample: TableSizeStats, collective_ids: list[int], -) -> tuple[JoinSideStats, JoinSideStats]: +) -> tuple[TableSizeStats, TableSizeStats]: """Aggregate table-size and row estimates across ranks.""" # AllGather size, row, and chunk count estimates across ranks ( @@ -880,13 +868,13 @@ async def _aggregate_estimates( right_sample.total_chunks, ) - new_left_sample = JoinSideStats( + new_left_sample = TableSizeStats( chunks=left_sample.chunks, total_size=left_total, total_rows=left_total_rows, total_chunks=left_total_chunks, ) - new_right_sample = JoinSideStats( + new_right_sample = TableSizeStats( chunks=right_sample.chunks, total_size=right_total, total_rows=right_total_rows, @@ -904,8 +892,8 @@ async def _choose_strategy_from_samples( right_partitioning: NormalizedPartitioning, executor: StreamingExecutor, *, - left_sample: JoinSideStats, - right_sample: JoinSideStats, + left_sample: TableSizeStats, + right_sample: TableSizeStats, chunkwise: bool, tracer: ActorTracer | None, ) -> JoinStrategy: @@ -1043,59 +1031,6 @@ def _modulus(partitioning: NormalizedPartitioning) -> int | None: return max(large, min_shuffle_modulus) -async def _sample_chunks( - context: Context, - ch: Channel[TableChunk], - max_sample_chunks: int, - max_sample_bytes: int, - local_count: int, -) -> JoinSideStats: - """ - Sample chunks from a channel. - - Parameters - ---------- - context - The context. - ch - The channel to sample from. - max_sample_chunks - The maximum number of chunks to sample. - max_sample_bytes - The maximum number of bytes to sample. - local_count - The number of local chunks. - - Returns - ------- - The sampled chunks. - """ - sampled_chunks: dict[int, TableChunk] = {} - total_size = 0 - total_rows = 0 - for _ in range(max_sample_chunks): - msg = await ch.recv(context) - if msg is None: - break - chunk = TableChunk.from_message(msg, br=context.br()).make_available_and_spill( - context.br(), allow_overbooking=True - ) - sampled_chunks[msg.sequence_number] = chunk - total_size += chunk.data_alloc_size() - total_rows += chunk.shape[0] - if total_size >= max_sample_bytes: - break - if sampled_chunks: - total_size = int((total_size / len(sampled_chunks)) * local_count) - total_rows = int((total_rows / len(sampled_chunks)) * local_count) - return JoinSideStats( - chunks=sampled_chunks, - total_size=total_size, - total_rows=total_rows, - total_chunks=local_count, - ) - - async def _choose_strategy( context: Context, comm: Communicator, @@ -1108,7 +1043,7 @@ async def _choose_strategy( collective_ids: list[int], *, tracer: ActorTracer | None, -) -> tuple[JoinSideStats, JoinSideStats, JoinStrategy]: +) -> tuple[TableSizeStats, TableSizeStats, JoinStrategy]: """Sample both sides, aggregate estimates, and choose broadcast vs shuffle.""" nranks = comm.nranks left_partitioning = NormalizedPartitioning.from_keys( @@ -1125,8 +1060,8 @@ async def _choose_strategy( if left_partitioning.is_aligned_with(right_partitioning, context.br()): # We can use a chunkwise join chunkwise = True - left_sample = JoinSideStats(total_chunks=left_metadata.local_count) - right_sample = JoinSideStats(total_chunks=right_metadata.local_count) + left_sample = TableSizeStats(total_chunks=left_metadata.local_count) + right_sample = TableSizeStats(total_chunks=right_metadata.local_count) else: # Need to shuffle or broadcast - Use sampled data to choose a strategy chunkwise = False diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/over.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/over.py new file mode 100644 index 000000000000..1a0e09b8e8e7 --- /dev/null +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/over.py @@ -0,0 +1,810 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +""" +Window ``over()`` actor for the RapidsMPF streaming runtime. + +Implements the ``group_to_rows`` ``WindowMapping`` only: each input row +receives the value computed for its group. Other mappings (``explode``, +``join``) are not supported. + +The actor picks one of three strategies at runtime based on the incoming +channel metadata and the shape of the windowed expressions. + +Chunkwise (already partitioned) + If the channel is already hash-partitioned on the over-keys (or any + prefix of them), every group is fully contained within one rank's + chunks. The window expression is correct on each chunk in isolation + and no cross-rank coordination is needed. + +Scalar broadcast (decomposable aggregations) + When every aggregation is decomposable, partial aggregates can be + combined associatively across ranks. Each rank computes per-chunk + partials, an AllGather collects them, a single reduction yields the + global aggregate per group, and each input chunk has those results + joined back onto its rows by the partition keys. Order is preserved + naturally: input chunks are buffered in receive order and emitted in + the same order after the global aggregate is known. + +Forward + return shuffle (non-decomposable aggregations) + For functions that need every row in a group visible at once, a hash + shuffle on the partition keys co-locates each group on one rank for + evaluation. After evaluation, a second shuffle routes each row back + to the rank that originally received it (output channels are + rank-local, so only the originating rank can emit), and the rows are + reassembled in input order using stamps that travel with the data + through both shuffles. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, ClassVar, cast + +from rapidsmpf.memory.memory_reservation import opaque_memory_usage +from rapidsmpf.shuffler import PartitionAssignment +from rapidsmpf.streaming.core.actor import define_actor +from rapidsmpf.streaming.core.message import Message +from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata +from rapidsmpf.streaming.cudf.table_chunk import ( + TableChunk, + make_table_chunks_available_or_wait, +) + +import polars as pl + +import pylibcudf as plc + +from cudf_polars.containers import Column, DataFrame, DataType +from cudf_polars.dsl.expr import GroupedWindow +from cudf_polars.dsl.expressions.base import ExecutionContext +from cudf_polars.dsl.utils.naming import unique_names +from cudf_polars.dsl.utils.reshape import broadcast +from cudf_polars.experimental.over import Over, _build_over_groupby_irs +from cudf_polars.experimental.rapidsmpf.collectives.shuffle import ( + LocalRepartitioner, + ShuffleManager, +) +from cudf_polars.experimental.rapidsmpf.dispatch import generate_ir_sub_network +from cudf_polars.experimental.rapidsmpf.utils import ( + ChannelManager, + ChunkStore, + NormalizedPartitioning, + _evaluate_chunk_sync, + _sample_chunks, + allgather_and_reduce, + allgather_reduce, + chunk_to_frame, + chunkwise_evaluate, + empty_table_chunk, + evaluate_batch, + evaluate_chunk, + gather_in_task_group, + maybe_remap_partitioning, + process_children, + recv_metadata, + replay_buffered_channel, + send_metadata, + shutdown_on_error, +) + +if TYPE_CHECKING: + from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.memory.buffer_resource import BufferResource + from rapidsmpf.streaming.core.channel import Channel + from rapidsmpf.streaming.core.context import Context + + from rmm.pylibrmm.stream import Stream + + from cudf_polars.dsl.expr import Col + from cudf_polars.dsl.ir import IR, GroupBy, IRExecutionContext, Select + from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator + from cudf_polars.experimental.rapidsmpf.utils import TableSizeStats + + +@dataclass(frozen=True) +class _ScalarOverPlan: + """Pre-computed IR rewrites for the scalar Over path.""" + + key_names: tuple[str, ...] + piecewise_ir: GroupBy + reduction_ir: GroupBy + agg_select_ir: Select + + +def _build_scalar_over_plan(ir: Over) -> _ScalarOverPlan: + """Pre-compute the IR rewrites needed by the scalar Over path.""" + gw_nodes = tuple(ne.value for ne in ir.exprs if isinstance(ne.value, GroupedWindow)) + # Lowering rejects non-Col partition-by keys, so every by-child here is a Col. + by_children = gw_nodes[0].children[: gw_nodes[0].by_count] + key_names = tuple(cast("Col", c).name for c in by_children) + piecewise_ir, reduction_ir, agg_select_ir = _build_over_groupby_irs( + gw_nodes, ir.children[0] + ) + return _ScalarOverPlan( + key_names=key_names, + piecewise_ir=piecewise_ir, + reduction_ir=reduction_ir, + agg_select_ir=agg_select_ir, + ) + + +def _broadcast_gw_sync( + gw: GroupedWindow, + chunk_df: DataFrame, + global_agg_df: DataFrame, + key_names: tuple[str, ...], + stream: Stream, +) -> Any: + """Broadcast the global aggregate for one GroupedWindow back to row positions.""" + by_exprs = gw.children[: gw.by_count] + by_cols = broadcast( + *(b.evaluate(chunk_df) for b in by_exprs), + target_length=chunk_df.num_rows, + stream=stream, + ) + by_tbl = plc.Table([c.obj for c in by_cols]) + group_keys_tbl = global_agg_df.select(key_names).table + + out_names, out_dtypes = zip( + *((ne.name, ne.value.dtype) for ne in gw.named_aggs), strict=True + ) + value_tbl = global_agg_df.select(out_names).table + + broadcasted_cols = gw._broadcast_agg_results( + by_tbl, group_keys_tbl, value_tbl, out_names, out_dtypes, stream + ) + temp_df = DataFrame(broadcasted_cols, stream=stream) + return gw.post.value.evaluate(temp_df, context=ExecutionContext.FRAME) + + +def _evaluate_ir_broadcast_sync( + chunk: TableChunk, + ir: Over, + global_agg_df: DataFrame, + key_names: tuple[str, ...], + ir_context: IRExecutionContext, + br: BufferResource, +) -> TableChunk: + """Map the per-group aggregate onto a chunk's rows to produce its Over output.""" + chunk_df = chunk_to_frame(chunk, ir.children[0]) + # global_agg_df and chunk_df may live on different streams (the former from + # the upstream allgather/reduction on ir_context's stream, the latter from + # the input message). Join them so the broadcast kernels read global_agg_df + # safely. + with ir_context.stream_ordered_after(chunk_df, global_agg_df) as stream: + result_cols = [ + _broadcast_gw_sync( + ne.value, chunk_df, global_agg_df, key_names, stream + ).rename(ne.name) + if isinstance(ne.value, GroupedWindow) + else ne.evaluate(chunk_df, context=ExecutionContext.FRAME) + for ne in ir.exprs + ] + + return TableChunk.from_pylibcudf_table( + plc.Table([c.obj for c in result_cols]), + stream, + exclusive_view=True, + br=br, + ) + + +async def _evaluate_broadcast_chunk( + context: Context, + chunk: TableChunk, + ir: Over, + global_agg_df: DataFrame, + key_names: tuple[str, ...], + ir_context: IRExecutionContext, + global_agg_per_row_size: int, +) -> TableChunk: + """Unspill the chunk and map the per-group aggregate onto its rows.""" + chunk, extra = await make_table_chunks_available_or_wait( + context, + chunk, + reserve_extra=global_agg_per_row_size * chunk.shape[0], + net_memory_delta=0, + ) + with opaque_memory_usage(extra): + return await asyncio.to_thread( + _evaluate_ir_broadcast_sync, + chunk, + ir, + global_agg_df, + key_names, + ir_context, + context.br(), + ) + + +@dataclass(frozen=True) +class OriginStamps: + """ + Stamp column names that ride both shuffles for output reassembly. + + Parameters + ---------- + chunk_index + Column name for a dense rank-local 0..N-1 counter identifying which + input chunk a row came from. + position + Column name for the row's position within its input chunk. + rank + Column name for the originating rank. + """ + + chunk_index: str + position: str + rank: str + + dtype: ClassVar[DataType] = DataType(pl.Int32()) + + @property + def names(self) -> tuple[str, str, str]: + """Stamp column names, in the order they are appended to the table.""" + return (self.chunk_index, self.position, self.rank) + + +def _origin_stamps_for(ir: Over) -> OriginStamps: + """Pick three stamp column names that do not collide with the schema.""" + names = unique_names((*ir.children[0].schema.keys(), *ir.schema.keys())) + return OriginStamps(next(names), next(names), next(names)) + + +def _append_origin_stamps( + chunk: TableChunk, + chunk_index: int, + origin_rank: int, + stream: Stream, + br: Any, +) -> TableChunk: + """Append (chunk_index, position, rank) stamp columns to *chunk*.""" + table = chunk.table_view() + n_rows = table.num_rows() + int32 = plc.types.DataType(plc.TypeId.INT32) + chunk_index_col = plc.Column.from_scalar( + plc.Scalar.from_py(chunk_index, int32, stream=stream), n_rows, stream=stream + ) + rank_col = plc.Column.from_scalar( + plc.Scalar.from_py(origin_rank, int32, stream=stream), n_rows, stream=stream + ) + position_col = plc.filling.sequence( + n_rows, + plc.Scalar.from_py(0, int32, stream=stream), + plc.Scalar.from_py(1, int32, stream=stream), + stream=stream, + ) + return TableChunk.from_pylibcudf_table( + plc.Table([*table.columns(), chunk_index_col, position_col, rank_col]), + stream, + exclusive_view=False, + br=br, + ) + + +def _evaluate_window_with_stamps( + chunk: TableChunk, + ir: Over, + ir_context: IRExecutionContext, + stamps: OriginStamps, +) -> DataFrame: + """Evaluate *ir* on the un-stamped portion of *chunk*; reattach stamps after.""" + child_schema = ir.children[0].schema + stream = ir_context.get_cuda_stream() + columns = chunk.table_view().columns() + n_child = len(child_schema) + + input_df = DataFrame.from_table( + plc.Table(columns[:n_child]), + list(child_schema.keys()), + list(child_schema.values()), + stream, + ) + result = ir.do_evaluate(ir.exprs, input_df, context=ir_context) + stamp_cols = [ + Column(col, dtype=stamps.dtype, name=name) + for col, name in zip(columns[n_child:], stamps.names, strict=True) + ] + return result.with_columns(stamp_cols, stream=stream) + + +def _partition_by_origin_rank( + result: DataFrame, + num_ranks: int, + br: Any, +) -> tuple[TableChunk | None, list[int]]: + """ + Rearrange rows so partition i contains rows whose origin rank is i. + + Returns a chunk with the rank stamp dropped and the per-rank split + indices for direct insertion into the return shuffle. + """ + if result.table.num_rows() == 0: + return None, [] + + stream = result.stream + columns = result.table.columns() + rank_column = columns[-1] + payload = plc.Table(columns[:-1]) + + rearranged, offsets = plc.partitioning.partition( + payload, rank_column, num_ranks, stream=stream + ) + return ( + TableChunk.from_pylibcudf_table(rearranged, stream, exclusive_view=True, br=br), + list(offsets[1:-1]), + ) + + +async def _allgather_and_broadcast( + context: Context, + comm: Communicator, + ir: Over, + ir_context: IRExecutionContext, + ch_in: Channel[TableChunk], + ch_out: Channel[TableChunk], + metadata_in: ChannelMetadata, + tracer: Any, + collective_id: int, + plan: _ScalarOverPlan, +) -> None: + """Compute partial aggregates per chunk, AllGather globally, then broadcast to each chunk.""" + piecewise_ir = plan.piecewise_ir + reduction_ir = plan.reduction_ir + agg_select_ir = plan.agg_select_ir + + buffer = ChunkStore(context) + partial_aggs: list[TableChunk] = [] + + while (msg := await ch_in.recv(context)) is not None: + chunk = TableChunk.from_message(msg, br=context.br()) + chunk, extra = await make_table_chunks_available_or_wait( + context, + chunk, + reserve_extra=chunk.data_alloc_size(), + net_memory_delta=0, + ) + with opaque_memory_usage(extra): + partial = await asyncio.to_thread( + _evaluate_chunk_sync, + chunk, + piecewise_ir, + ir_context, + context.br(), + ) + partial_aggs.append(partial) + buffer.insert(Message(msg.sequence_number, chunk)) + + if partial_aggs: + local_agg = await evaluate_batch( + partial_aggs, context, reduction_ir, ir_context=ir_context + ) + else: + local_agg = empty_table_chunk( + reduction_ir, context, ir_context.get_cuda_stream() + ) + + # AllGather the locally-reduced partials (pre post-aggregation) so a + # single global reduction combines them; the post-aggregation step + # runs once after. + if comm.nranks > 1 and not metadata_in.duplicated: + global_agg = await allgather_and_reduce( + context, comm, collective_id, local_agg, reduction_ir, ir_context + ) + else: + global_agg = local_agg + + global_agg = await evaluate_chunk( + context, global_agg, agg_select_ir, ir_context=ir_context + ) + global_agg_df = chunk_to_frame(global_agg, agg_select_ir) + global_agg_per_row_size = global_agg.data_alloc_size() // max( + 1, global_agg_df.num_rows + ) + + metadata_out = ChannelMetadata( + local_count=metadata_in.local_count, + partitioning=maybe_remap_partitioning(ir, metadata_in.partitioning), + duplicated=metadata_in.duplicated, + ) + await send_metadata(ch_out, context, metadata_out) + + for msg in buffer: + result = await _evaluate_broadcast_chunk( + context, + TableChunk.from_message(msg, br=context.br()), + ir, + global_agg_df, + plan.key_names, + ir_context, + global_agg_per_row_size, + ) + if tracer is not None: + tracer.add_chunk(table=result.table_view()) + await ch_out.send(context, Message(msg.sequence_number, result)) + + await ch_out.drain(context) + + +async def _choose_modulus( + context: Context, + comm: Communicator, + ch_in: Channel[TableChunk], + metadata_in: ChannelMetadata, + collective_id: int, + target_partition_size: int, + sample_chunk_count: int, +) -> tuple[TableSizeStats, int]: + """ + Sample input, AllGather size estimates, and derive the forward-shuffle modulus. + + Returns the sample (whose chunks must be replayed back to the consumer) + and the chosen number of forward-shuffle partitions. + """ + sample = await _sample_chunks( + context, + ch_in, + sample_chunk_count, + target_partition_size, + metadata_in.local_count, + ) + if comm.nranks > 1 and not metadata_in.duplicated: + total_bytes, total_count = await allgather_reduce( + context, comm, collective_id, sample.total_size, sample.total_chunks + ) + else: + total_bytes, total_count = sample.total_size, sample.total_chunks + modulus = min( + max(comm.nranks, total_bytes // max(1, target_partition_size)), + max(1, total_count), + ) + return sample, modulus + + +async def _distribute_by_group( + context: Context, + comm: Communicator, + forward_shuffle: ShuffleManager, + ch_in: Channel[TableChunk], + key_indices: tuple[int, ...], + ir_context: IRExecutionContext, + skip_insert: bool, # noqa: FBT001 +) -> list[int]: + """Stream chunks from *ch_in* into the forward shuffle with origin stamps.""" + # We already have the upstream metadata; signal we don't need the replay + # channel's copy. + await ch_in.shutdown_metadata(context) + + sequence_numbers: list[int] = [] + chunk_index = 0 + async with forward_shuffle.inserting() as inserter: + while (msg := await ch_in.recv(context)) is not None: + chunk = TableChunk.from_message( + msg, br=context.br() + ).make_available_and_spill(context.br(), allow_overbooking=True) + sequence_numbers.append(msg.sequence_number) + if not skip_insert: + # TODO: For duplicated input only rank 0 inserts here, and + # every row is stamped with origin_rank=0, so the return + # shuffle routes all output back to rank 0 and ranks + # 1..nranks-1 sit idle on emit. Slice the duplicated input + # across ranks (e.g. stripe by row index) and stamp each + # slice with its target origin rank to distribute emit work. + stamped = await asyncio.to_thread( + _append_origin_stamps, + chunk, + chunk_index, + comm.rank, + ir_context.get_cuda_stream(), + context.br(), + ) + inserter.insert_hash(stamped, key_indices) + chunk_index += 1 + return sequence_numbers + + +async def _evaluate_and_route_to_origin( + context: Context, + ir: Over, + ir_context: IRExecutionContext, + forward_shuffle: ShuffleManager, + return_shuffle: ShuffleManager, + num_ranks: int, + stamps: OriginStamps, +) -> None: + """Window-evaluate each local forward partition, then ship rows back to their origin.""" + async with return_shuffle.inserting() as inserter: + for partition_id in forward_shuffle.local_partitions(): + stream = ir_context.get_cuda_stream() + extracted = forward_shuffle.extract_chunk(partition_id, stream) + if extracted.num_rows() == 0: + continue + partition = TableChunk.from_pylibcudf_table( + extracted, stream, exclusive_view=True, br=context.br() + ) + evaluated = await asyncio.to_thread( + _evaluate_window_with_stamps, partition, ir, ir_context, stamps + ) + routed, splits = await asyncio.to_thread( + _partition_by_origin_rank, evaluated, num_ranks, context.br() + ) + if routed is not None: + inserter.insert_split(routed, splits) + + +async def _reassemble_input_chunks( + context: Context, + ch_out: Channel[TableChunk], + ir_context: IRExecutionContext, + return_shuffle: ShuffleManager, + sequence_numbers: list[int], + ir: Over, + tracer: Any, +) -> None: + """Emit one output chunk per input chunk, in original order.""" + n_chunks = len(sequence_numbers) + if n_chunks == 0: + return + + n_exprs = len(ir.exprs) + chunk_index_column = n_exprs + + # TODO: thread ir_context through repartition_by_index so each + # PackedData piece moves on its own pool stream rather than sharing one. + local = LocalRepartitioner(return_shuffle, local_count=n_chunks) + await local.repartition_by_index( + partition_col=chunk_index_column, stream=ir_context.get_cuda_stream() + ) + + for chunk_index, sequence_number in zip( + local.local_partitions(), sequence_numbers, strict=True + ): + # Distinct stream per chunk so downstream work on different + # chunks can overlap on the GPU. + stream = ir_context.get_cuda_stream() + tbl = local.extract_chunk(chunk_index, stream) + if tbl.num_rows() == 0: + chunk = empty_table_chunk(ir, context, stream) + else: + sorted_tbl = plc.sorting.stable_sort_by_key( + tbl, + plc.Table([tbl.columns()[n_exprs]]), + [plc.types.Order.ASCENDING], + [plc.types.NullOrder.AFTER], + stream=stream, + ) + chunk = TableChunk.from_pylibcudf_table( + plc.Table(sorted_tbl.columns()[:n_exprs]), + stream, + exclusive_view=True, + br=context.br(), + ) + if tracer is not None: + tracer.add_chunk(table=chunk.table_view()) + await ch_out.send(context, Message(sequence_number, chunk)) + + +async def _shuffle_and_reassemble( + context: Context, + comm: Communicator, + ir: Over, + ir_context: IRExecutionContext, + ch_in: Channel[TableChunk], + ch_out: Channel[TableChunk], + metadata_in: ChannelMetadata, + tracer: Any, + size_collective_id: int, + forward_shuffle_collective_id: int, + return_shuffle_collective_id: int, + target_partition_size: int, + sample_chunk_count: int, +) -> None: + """Hash-shuffle by partition keys, evaluate, then route rows back to their origin rank.""" + stamps = _origin_stamps_for(ir) + + metadata_out = ChannelMetadata( + local_count=metadata_in.local_count, + partitioning=maybe_remap_partitioning(ir, metadata_in.partitioning), + duplicated=False, + ) + await send_metadata(ch_out, context, metadata_out) + + skip_insert = metadata_in.duplicated and comm.rank != 0 + + sample, forward_modulus = await _choose_modulus( + context, + comm, + ch_in, + metadata_in, + size_collective_id, + target_partition_size, + sample_chunk_count, + ) + + forward_shuffle = ShuffleManager( + context, comm, forward_modulus, forward_shuffle_collective_id + ) + return_shuffle = ShuffleManager( + context, + comm, + comm.nranks, + return_shuffle_collective_id, + partition_assignment=PartitionAssignment.CONTIGUOUS, + ) + + ch_replay = context.create_channel() + sequence_numbers, _ = await gather_in_task_group( + _distribute_by_group( + context, + comm, + forward_shuffle, + ch_replay, + ir.key_indices, + ir_context, + skip_insert, + ), + replay_buffered_channel( + context, ch_replay, ch_in, sample.chunks, metadata_in, trace_ir=ir + ), + ) + + await _evaluate_and_route_to_origin( + context, + ir, + ir_context, + forward_shuffle, + return_shuffle, + comm.nranks, + stamps, + ) + await _reassemble_input_chunks( + context, ch_out, ir_context, return_shuffle, sequence_numbers, ir, tracer + ) + + await ch_out.drain(context) + + +@define_actor() +async def over_actor( + context: Context, + comm: Communicator, + ir: Over, + ir_context: IRExecutionContext, + ch_out: Channel[TableChunk], + ch_in: Channel[TableChunk], + collective_ids: list[int], + target_partition_size: int, + sample_chunk_count: int, + scalar_plan: _ScalarOverPlan | None, +) -> None: + """ + Streaming actor for window ``over()`` expressions. + + Parameters + ---------- + context + The rapidsmpf context. + comm + The communicator. + ir + The Over IR node. + ir_context + The IR execution context. + ch_out + The output channel. + ch_in + The input channel. + collective_ids + Collective IDs reserved for this operation. Scalar Over nodes receive + one ID (AllGather); non-scalar nodes receive two (one shared by the + size AllGather and forward Shuffle, plus a separate one for the + return Shuffle which overlaps with the forward extract). + target_partition_size + Target output partition size in bytes, used to compute the shuffle + modulus for the non-scalar path. + sample_chunk_count + Maximum number of input chunks to sample when estimating the shuffle + modulus on the non-scalar path. + scalar_plan + Pre-computed IR rewrites for the scalar Over path, built at planning + time. ``None`` for non-scalar Over nodes. + """ + async with shutdown_on_error( + context, ch_in, ch_out, trace_ir=ir, ir_context=ir_context + ) as tracer: + metadata_in = await recv_metadata(ch_in, context) + + partitioning = NormalizedPartitioning.from_keys( + metadata_in.partitioning, + comm.nranks, + keys=ir.key_indices, + allow_subset=True, + ) + if partitioning.is_strictly_partitioned(): + metadata_out = ChannelMetadata( + local_count=metadata_in.local_count, + partitioning=maybe_remap_partitioning(ir, metadata_in.partitioning), + duplicated=metadata_in.duplicated, + ) + await chunkwise_evaluate( + context, + ir, + ir_context, + ch_out, + ch_in, + metadata_out, + tracer=tracer, + ) + return + + if ir.is_scalar: + assert scalar_plan is not None + await _allgather_and_broadcast( + context, + comm, + ir, + ir_context, + ch_in, + ch_out, + metadata_in, + tracer, + collective_ids[0], + scalar_plan, + ) + else: + await _shuffle_and_reassemble( + context, + comm, + ir, + ir_context, + ch_in, + ch_out, + metadata_in, + tracer, + # collective_ids[0] is reused for the size AllGather and the + # forward shuffle (sequential, no overlap); collective_ids[1] + # is the return shuffle, which overlaps with forward extract. + size_collective_id=collective_ids[0], + forward_shuffle_collective_id=collective_ids[0], + return_shuffle_collective_id=collective_ids[1], + target_partition_size=target_partition_size, + sample_chunk_count=sample_chunk_count, + ) + + +@generate_ir_sub_network.register(Over) +def _( + ir: Over, rec: SubNetGenerator +) -> tuple[dict[IR, list[Any]], dict[IR, ChannelManager]]: + executor = rec.state["config_options"].executor + actors, channels = process_children(ir, rec) + channels[ir] = ChannelManager(rec.state["context"]) + collective_ids = list(rec.state["collective_id_map"].get(ir, [])) + if not ir.is_scalar and executor.dynamic_planning is None: + raise ValueError( + "Non-scalar over() requires dynamic planning to size the forward " + "shuffle. Enable it via StreamingExecutor(dynamic_planning=...) " + "or the --dynamic-planning CLI flag." + ) + sample_chunk_count = ( + executor.dynamic_planning.sample_chunk_count + if executor.dynamic_planning is not None + else 0 + ) + scalar_plan = _build_scalar_over_plan(ir) if ir.is_scalar else None + actors[ir] = [ + over_actor( + rec.state["context"], + rec.state["comm"], + ir, + rec.state["ir_context"], + channels[ir].reserve_input_slot(), + channels[ir.children[0]].reserve_output_slot(), + collective_ids, + executor.target_partition_size, + sample_chunk_count, + scalar_plan, + ) + ] + return actors, channels diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py index ec24bfab3b7a..c9f39cab068f 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py @@ -71,6 +71,11 @@ def add_chunk(self, *, table: plc.Table | None = None) -> None: table The table to record. """ + # TODO: replace this API with one that takes a TableChunk directly, + # using TableChunk.shape[0] for the row count, and consider providing a + # helper that logs and sends a chunk in one call so the + # ``tracer.add_chunk(...) + ch_out.send(...)`` pattern doesn't have to + # be duplicated across every actor. if table is not None: # pragma: no cover; Covered by rapidsmpf tests self.row_count = (self.row_count or 0) + table.num_rows() self.chunk_count += 1 diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py index 341255b27fdd..d21f6c4f9774 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py @@ -10,8 +10,9 @@ import operator import struct import time +from collections import deque from contextlib import asynccontextmanager -from dataclasses import dataclass +from dataclasses import dataclass, field from functools import reduce from typing import TYPE_CHECKING, Any, Literal, TypeAlias, cast @@ -39,11 +40,20 @@ from cudf_polars.dsl.expr import Col, NamedExpr from cudf_polars.dsl.ir import Cache, Filter, GroupBy, HStack, Join, Projection, Select from cudf_polars.dsl.tracing import Scope +from cudf_polars.dsl.utils.naming import names_to_indices +from cudf_polars.experimental.rapidsmpf.collectives.allgather import AllGatherManager from cudf_polars.experimental.utils import _concat from cudf_polars.utils.dtypes import make_empty_column if TYPE_CHECKING: - from collections.abc import AsyncIterator, Callable, Coroutine, Iterator, Sequence + from collections.abc import ( + AsyncIterator, + Callable, + Coroutine, + Generator, + Iterator, + Sequence, + ) from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.memory.buffer_resource import BufferResource @@ -63,6 +73,23 @@ PartitioningScheme: TypeAlias = InterRankScheme | Literal["inherit"] +class ChunkStore: + """Ordered spillable buffer for TableChunk messages.""" + + def __init__(self, ctx: Context) -> None: + self._mids: deque[int] = deque() + self._store = ctx.spillable_messages() + + def insert(self, msg: Message) -> None: + """Insert a message into the store.""" + self._mids.append(self._store.insert(msg)) + + def __iter__(self) -> Generator[Message, None, None]: + """Yield messages in insertion order, draining the store.""" + while self._mids: + yield self._store.extract(mid=self._mids.popleft()) + + @contextlib.contextmanager def set_memory_resource(mr: rmm.mr.DeviceMemoryResource) -> Iterator[None]: """ @@ -354,6 +381,54 @@ async def recv_metadata(ch: Channel[TableChunk], ctx: Context) -> ChannelMetadat return ChannelMetadata.from_message(msg) +def _make_hash_shuffle_metadata( + comm: Communicator, + key_indices: tuple[int, ...], + modulus: int, + metadata_in: ChannelMetadata, +) -> ChannelMetadata: + """ + Build output ChannelMetadata for a hash shuffle by key_indices. + + Parameters + ---------- + comm + The communicator. + key_indices + Column indices to hash-partition on. + modulus + Number of output partitions (must be >= comm.nranks). + metadata_in + Input channel metadata (used for duplicated flag and, on a + single-rank run, to preserve the existing inter-rank scheme). + + Returns + ------- + ChannelMetadata + Ready to pass to send_metadata. + """ + nranks = comm.nranks + if nranks == 1: + inter_rank_scheme = ( + None + if metadata_in.partitioning is None + else metadata_in.partitioning.inter_rank + ) + local_scheme: HashScheme | str = HashScheme( + column_indices=key_indices, modulus=modulus + ) + local_output_count = modulus + else: + inter_rank_scheme = HashScheme(column_indices=key_indices, modulus=modulus) + local_scheme = "inherit" + local_output_count = (modulus - comm.rank + nranks - 1) // nranks + return ChannelMetadata( + local_count=local_output_count, + partitioning=Partitioning(inter_rank_scheme, local_scheme), + duplicated=metadata_in.duplicated, + ) + + def _evaluate_chunk_sync( chunk: TableChunk, ir: IR, @@ -434,6 +509,49 @@ async def evaluate_chunk( return chunk +async def allgather_and_reduce( + context: Context, + comm: Communicator, + collective_id: int, + local_chunk: TableChunk, + reduce_ir: IR, + ir_context: IRExecutionContext, +) -> TableChunk: + """ + AllGather ``local_chunk`` across ranks and apply ``reduce_ir`` to the result. + + Parameters + ---------- + context + The rapidsmpf streaming context. + comm + The communicator. + collective_id + Collective operation ID for the AllGather. + local_chunk + The locally-reduced chunk this rank contributes. + reduce_ir + IR node applied to the concatenated AllGather output. + ir_context + The IR execution context. + + Returns + ------- + The chunk produced by evaluating ``reduce_ir`` on the gathered result. + """ + allgather = AllGatherManager(context, comm, collective_id) + with allgather.inserting() as inserter: + inserter.insert(0, local_chunk) + stream = ir_context.get_cuda_stream() + concat_chunk = TableChunk.from_pylibcudf_table( + await allgather.extract_concatenated(stream), + stream, + exclusive_view=True, + br=context.br(), + ) + return await evaluate_chunk(context, concat_chunk, reduce_ir, ir_context=ir_context) + + async def concat_batch( batch: list[TableChunk], context: Context, @@ -610,29 +728,71 @@ def indices_to_names(indices: tuple[int, ...], schema: Schema) -> tuple[str, ... return tuple(keys[i] for i in indices) -def names_to_indices( - names: tuple[str | NamedExpr, ...], schema: Schema -) -> tuple[int, ...]: - """ - Return column indices for the given names in schema order. +@dataclass(frozen=True) +class TableSizeStats: + """Sampled chunks and aggregate size/row stats for a table channel.""" + + chunks: dict[int, TableChunk] = field(default_factory=dict) + """The sampled chunks, keyed by sequence number.""" + total_size: int = 0 + """The total estimated size of the table in bytes.""" + total_rows: int = 0 + """The total estimated number of rows in the table.""" + total_chunks: int = 0 + """The total estimated number of chunks in the table.""" - Accepts either column names (str) or NamedExpr, so it can be used with - e.g. ir.left_on, ir.right_on as well as plain name tuples. + +async def _sample_chunks( + context: Context, + ch: Channel[TableChunk], + max_sample_chunks: int, + max_sample_bytes: int, + local_count: int, +) -> TableSizeStats: + """ + Sample chunks from a channel and extrapolate to a per-rank size estimate. Parameters ---------- - names - The names to get indices for. - schema - The schema to get indices from. + context + The context. + ch + The channel to sample from. + max_sample_chunks + The maximum number of chunks to sample. + max_sample_bytes + The maximum number of bytes to sample. + local_count + The expected number of local chunks (used for extrapolation). Returns ------- - The column indices for each name in schema order. + Sampled chunks and the extrapolated total size/rows for this rank. """ - keys = list(schema.keys()) - str_names = [n.name if isinstance(n, NamedExpr) else n for n in names] - return tuple(keys.index(n) for n in str_names) + sampled_chunks: dict[int, TableChunk] = {} + total_size = 0 + total_rows = 0 + for _ in range(max_sample_chunks): + msg = await ch.recv(context) + if msg is None: + break + chunk = TableChunk.from_message(msg, br=context.br()).make_available_and_spill( + context.br(), allow_overbooking=True + ) + sampled_chunks[msg.sequence_number] = chunk + total_size += chunk.data_alloc_size() + total_rows += chunk.shape[0] + if total_size >= max_sample_bytes: + break + if sampled_chunks: + total_size = int((total_size / len(sampled_chunks)) * local_count) + total_rows = int((total_rows / len(sampled_chunks)) * local_count) + return TableSizeStats( + chunks=sampled_chunks, + total_size=total_size, + total_rows=total_rows, + total_chunks=local_count, + ) async def replay_buffered_channel( diff --git a/python/cudf_polars/cudf_polars/experimental/select.py b/python/cudf_polars/cudf_polars/experimental/select.py index 606741d587d1..87c34db1b1cc 100644 --- a/python/cudf_polars/cudf_polars/experimental/select.py +++ b/python/cudf_polars/cudf_polars/experimental/select.py @@ -20,6 +20,7 @@ decompose_expr_graph, make_expr_decomposer, ) +from cudf_polars.experimental.over import _fuse_over_nodes from cudf_polars.experimental.repartition import Repartition from cudf_polars.experimental.utils import ( _contains_unsupported_fill_strategy, @@ -166,6 +167,7 @@ def decompose_select( # Concatenate partial selections new_ir: Select | HConcat + selections, partition_info = _fuse_over_nodes(selections, partition_info) selections, partition_info = _fuse_simple_reductions( selections, partition_info, diff --git a/python/cudf_polars/tests/experimental/test_rolling.py b/python/cudf_polars/tests/experimental/test_rolling.py index bc1bbbc40eb2..6c7a71ce0b9b 100644 --- a/python/cudf_polars/tests/experimental/test_rolling.py +++ b/python/cudf_polars/tests/experimental/test_rolling.py @@ -49,6 +49,211 @@ def test_rolling_datetime(request, engine): assert_gpu_result_equal(q, engine=engine) +@pytest.mark.parametrize( + "expr", + [ + pl.col("x").sum().over("g"), + pl.len().over("g"), + pl.col("x").sum().over("g", "g2"), + pl.col("x").sum().over("g_null"), + pl.col("x").sum().over("g", order_by="s"), + pl.col("x").rank(method="dense", descending=True).over("g"), + pl.col("x").rank(method="min").over("g", "g2"), + pl.col("x").cum_sum().over("g", order_by="s"), + pl.when((pl.col("x") % 2) == 0) + .then(None) + .otherwise(pl.col("x")) + .fill_null(strategy="forward") + .over("g", order_by="s"), + ], + ids=[ + "single_key_sum", + "len_over", + "multi_key", + "null_keys", + "order_by", + "rank_dense", + "rank_min_multi_key", + "cum_sum_order_by", + "fill_null_forward", + ], +) +def test_over_select(engine, expr): + df = pl.LazyFrame( + { + "g": [1, 1, 2, 2, 2, 1], + "x": [1, 2, 3, 4, 5, 6], + "g2": ["a", "b", "a", "b", "a", "b"], + "g_null": [1, None, 1, None, 2, 1], + "s": [6, 5, 4, 3, 2, 1], + } + ) + assert_gpu_result_equal(df.select(expr), engine=engine, check_row_order=True) + + +def test_over_with_columns(engine): + df = pl.LazyFrame( + { + "g": [1, 1, 2, 2, 2, 1], + "x": [1, 2, 3, 4, 5, 6], + } + ) + assert_gpu_result_equal( + df.with_columns(pl.col("x").sum().over("g")), + engine=engine, + check_row_order=True, + ) + + +def test_over_colliding_internal_agg_names(engine): + df = pl.LazyFrame( + { + "category": ["A", "A", "B", "B", "C"], + "value": [20, 30, 15, 40, 35], + } + ) + q = df.select( + pl.col("category"), + pl.col("value"), + pl.col("value").sum().over("category").alias("cat_sum"), + pl.col("value").mean().over("category").alias("cat_avg"), + ).sort("category", "value") + assert_gpu_result_equal(q, engine=engine, check_row_order=True) + + +@pytest.mark.parametrize( + "expr", + [ + pl.col("x").sum().over(pl.col("g") % 2), + pl.col("x").sum().over("g", pl.col("x") % 2), + ], + ids=["noncol_key", "mixed_col_and_expr_key"], +) +def test_over_noncol_key_fallback(request, streaming_engine_factory, expr) -> None: + # Non-Col and mixed Col/expr partition-by keys are not yet supported for + # multi-partition streaming and should fall back to single-partition. + engine = streaming_engine_factory( + StreamingOptions(max_rows_per_partition=2, fallback_mode="warn"), + ) + if not isinstance(engine, SPMDEngine): + # On Dask/Ray the fallback warning fires on worker processes and is + # invisible to ``pytest.warns``. + request.applymarker( + pytest.mark.xfail( + reason="https://github.com/rapidsai/cudf/issues/22405", + strict=False, + ) + ) + df = pl.LazyFrame( + { + "g": [1, 1, 2, 2, 2, 1], + "x": [1, 2, 3, 4, 5, 6], + } + ) + with pytest.warns(UserWarning, match=r"not supported for multiple partitions"): + assert_gpu_result_equal(df.select(expr), engine=engine) + + +def test_over_mixed_keys(streaming_engine_factory) -> None: + # Multiple over expressions with different partition-by keys are decomposed + # into separate Over nodes (one per key group) and combined with HConcat. + engine = streaming_engine_factory( + StreamingOptions(max_rows_per_partition=2, fallback_mode="warn"), + ) + df = pl.LazyFrame( + { + "g": [1, 1, 2, 2, 2, 1], + "g2": ["a", "b", "a", "b", "a", "b"], + "x": [1, 2, 3, 4, 5, 6], + } + ) + q = df.select( + pl.col("x").sum().over("g").alias("s_g"), + pl.col("x").sum().over("g2").alias("s_g2"), + ) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) + + +@pytest.mark.parametrize( + "expr", + [ + pl.col("x").sum().over("g"), + pl.len().over("g"), + pl.col("x").rank(method="dense").over("g"), + pl.col("x").cum_sum().over("g", order_by="s"), + ], + ids=["scalar_sum", "scalar_len", "nonscalar_rank", "nonscalar_cum_sum"], +) +@pytest.mark.parametrize("max_rows_per_partition", [1, 2]) +def test_over_many_partitions( + streaming_engine_factory, max_rows_per_partition, expr +) -> None: + # Small max_rows_per_partition forces many chunks, exercising the AllGather + # (scalar broadcast) and sort-and-split (non-scalar) paths across many + # partitions. Two values cover both single-row chunks and multi-row chunks + # so the within-chunk position sort is also exercised. + engine = streaming_engine_factory( + StreamingOptions( + max_rows_per_partition=max_rows_per_partition, fallback_mode="warn" + ), + ) + df = pl.LazyFrame( + { + "g": [1, 1, 2, 2, 2, 1], + "x": [1, 2, 3, 4, 5, 6], + "s": [6, 5, 4, 3, 2, 1], + } + ) + assert_gpu_result_equal(df.select(expr), engine=engine, check_row_order=True) + + +@pytest.mark.parametrize( + "expr", + [ + pl.col("x").sum().over("g"), + pl.col("x").rank(method="dense").over("g"), + ], + ids=["scalar_sum", "nonscalar_rank"], +) +def test_over_empty_input(streaming_engine_factory, expr) -> None: + engine = streaming_engine_factory( + StreamingOptions(max_rows_per_partition=2, fallback_mode="warn"), + ) + df = pl.LazyFrame( + { + "g": pl.Series([], dtype=pl.Int64), + "x": pl.Series([], dtype=pl.Int64), + } + ) + assert_gpu_result_equal(df.select(expr), engine=engine, check_row_order=True) + + +@pytest.mark.parametrize( + "expr", + [ + pl.col("x").sum().over("g"), + pl.col("x").rank(method="dense").over("g"), + ], + ids=["scalar_sum", "nonscalar_rank"], +) +def test_over_already_partitioned(streaming_engine_factory, expr) -> None: + # broadcast_limit=0 disables broadcast joins entirely. Therefore, we should + # already be shuffled on "g" after the join. The over("g") actor should + # detect this and evaluate chunkwise without a second shuffle. + engine = streaming_engine_factory( + StreamingOptions( + max_rows_per_partition=3, broadcast_limit=0, fallback_mode="warn" + ), + ) + left = pl.LazyFrame({"g": [1, 1, 2, 2, 2, 1], "x": [1, 2, 3, 4, 5, 6]}) + right = pl.LazyFrame({"g": [1, 2], "y": [10, 20]}) + assert_gpu_result_equal( + left.join(right, on="g").with_columns(expr), + engine=engine, + check_row_order=False, + ) + + def test_over_in_filter_unsupported(request, streaming_engine_factory) -> None: engine = streaming_engine_factory( StreamingOptions(max_rows_per_partition=1, fallback_mode="warn"), diff --git a/python/cudf_polars/tests/experimental/test_spmd.py b/python/cudf_polars/tests/experimental/test_spmd.py index fabaeedbc780..a61268cf3876 100644 --- a/python/cudf_polars/tests/experimental/test_spmd.py +++ b/python/cudf_polars/tests/experimental/test_spmd.py @@ -393,3 +393,146 @@ def test_reset_rejects_construction_time_engine_options( ) with pytest.raises(ValueError, match="memory_resource_config"): engine._reset(engine_options={"memory_resource_config": None}) + + +# Group keys probed with num_partitions=2, nranks=2, ROUND_ROBIN: +# _SAME_RANK_KEYS[r] hashes to partition r: data stays on its origin rank. +# _CROSS_RANK_KEYS[r] hashes to partition 1-r: data is fully shuffled away. +# num_partitions=2 = max(nranks=2, local_count=1). local_count=1 requires +# max_rows_per_partition >= the number of rows per rank (3 here), so we use 4. +_SAME_RANK_KEYS = [ + 0, + 3, +] # g=0 hashes to partition 0 (rank 0); g=3 hashes to partition 1 (rank 1) +_CROSS_RANK_KEYS = [ + 3, + 0, +] # g=3 hashes to partition 1 (rank 1); g=0 hashes to partition 0 (rank 0) + + +@pytest.mark.parametrize( + "expr,is_scalar", + [ + (pl.col("x").sum().over("g").alias("result"), True), + (pl.col("x").rank(method="dense").over("g").alias("result"), False), + ], + ids=["scalar_sum", "nonscalar_rank"], +) +@pytest.mark.parametrize( + "cross_rank", + [False, True], + ids=["same_rank", "cross_rank"], +) +def test_over_multirank( + request: pytest.FixtureRequest, + comm: Communicator, + expr: pl.Expr, + is_scalar: bool, # noqa: FBT001 + cross_rank: bool, # noqa: FBT001 +) -> None: + """over() correctness in multi-rank SPMD mode, same-rank and cross-rank cases. + + same_rank: group keys hash to the origin rank's own partition (happy path). + cross_rank: group keys hash to the other rank's partition, exercising the + bug where row_idx spaces are rank-local so Phase 2 fills the wrong + accumulated slots and each rank receives the other rank's data. + + max_rows_per_partition=4 keeps all 3 rows in one chunk (local_count=1), + so num_partitions=max(nranks=2, 1)=2, matching the probed key assignments. + """ + with SPMDEngine( + comm=comm, + executor_options={"max_rows_per_partition": 4, "dynamic_planning": {}}, + ) as engine: + rank = engine.rank + nranks = engine.nranks + if nranks != 2: + request.applymarker( + pytest.mark.skip( + reason="key assignments are probed for exactly 2 ranks" + ) + ) + keys = _CROSS_RANK_KEYS if cross_rank else _SAME_RANK_KEYS + g = keys[rank] + xs = [rank * 3 + 1, rank * 3 + 2, rank * 3 + 3] + lf = pl.LazyFrame({"g": [g, g, g], "x": xs}) + local_result = lf.select(pl.col("g"), pl.col("x"), expr).collect(engine=engine) + + # Each rank must get back its OWN rows (not another rank's). + assert local_result["g"].unique().to_list() == [g], ( + f"rank {rank}: expected only group {g} in output, " + f"got {local_result['g'].unique().to_list()}" + ) + + with reserve_op_id() as op_id: + global_result = allgather_polars_dataframe( + engine=engine, local_df=local_result, op_id=op_id + ) + + assert global_result.shape == (3 * nranks, 3) + for r in range(nranks): + grp_g = keys[r] + grp = global_result.filter(pl.col("g") == grp_g).sort("x") + assert grp.shape == (3, 3), f"rank {r} group has wrong row count" + expected_xs = [r * 3 + 1, r * 3 + 2, r * 3 + 3] + assert grp["x"].to_list() == expected_xs + if is_scalar: + assert grp["result"].to_list() == [sum(expected_xs)] * 3 + else: + assert grp["result"].to_list() == [1, 2, 3] + + +def test_over_nonscalar_duplicated_input( + request: pytest.FixtureRequest, + comm: Communicator, +) -> None: + """Non-scalar over() on duplicated=True input produces correct row count and values. + + group_by() AllGathers its result onto all ranks (duplicated=True). The + non-scalar over() path must output duplicated=False and only insert rows on + rank 0, otherwise all ranks insert the same rows (N-fold overcounting) and + the downstream Repartition skips AllGather. + + max_rows_per_partition=10 keeps all 3 rows in one chunk (local_count=1), + so modulus=max(nranks=2, 1)=2, matching the _SAME_RANK_KEYS assignments. + """ + with SPMDEngine( + comm=comm, + executor_options={"max_rows_per_partition": 10, "dynamic_planning": {}}, + ) as engine: + rank = engine.rank + nranks = engine.nranks + if nranks != 2: + request.applymarker( + pytest.mark.skip( + reason="key assignments are probed for exactly 2 ranks" + ) + ) + + coarse_g = _SAME_RANK_KEYS[rank] + fine_gs = [rank * 3 + 1, rank * 3 + 2, rank * 3 + 3] + xs = [rank * 30 + 10, rank * 30 + 20, rank * 30 + 30] + lf = pl.LazyFrame({"fine_g": fine_gs, "coarse_g": [coarse_g] * 3, "x": xs}) + local_result = ( + lf.group_by("fine_g", "coarse_g") + .agg(pl.col("x").first()) + .with_columns( + pl.col("x").rank(method="dense").over("coarse_g").alias("rank_x") + ) + .collect(engine=engine) + ) + + with reserve_op_id() as op_id: + global_result = allgather_polars_dataframe( + engine=engine, local_df=local_result, op_id=op_id + ) + + assert global_result.shape == (3 * nranks, 4) + for r in range(nranks): + cg = _SAME_RANK_KEYS[r] + grp = global_result.filter(pl.col("coarse_g") == cg).sort("x") + assert grp.shape == (3, 4), f"coarse_g={cg}: wrong row count" + assert grp["rank_x"].to_list() == [1, 2, 3], ( + f"coarse_g={cg}: expected dense ranks [1, 2, 3] " + f"but got {grp['rank_x'].to_list()}" + )