From e63aad5fe5512125580320de86192713c4ec8445 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 5 Aug 2026 07:15:17 -0700 Subject: [PATCH 01/13] groupby-sortby first/last support --- python/cudf_polars/cudf_polars/dsl/expr.py | 3 +- .../dsl/expressions/aggregation.py | 38 ++++- python/cudf_polars/cudf_polars/dsl/ir.py | 154 ++++++++++++++++-- .../cudf_polars/dsl/utils/aggregations.py | 26 +++ .../cudf_polars/streaming/groupby.py | 48 ++++++ .../tests/streaming/test_groupby.py | 20 +++ .../cudf_polars/tests/streaming/test_spmd.py | 55 ++++++- python/cudf_polars/tests/test_groupby.py | 80 +++++++++ 8 files changed, 401 insertions(+), 23 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/expr.py b/python/cudf_polars/cudf_polars/dsl/expr.py index 16285a6c195c..14dd8d286957 100644 --- a/python/cudf_polars/cudf_polars/dsl/expr.py +++ b/python/cudf_polars/cudf_polars/dsl/expr.py @@ -13,7 +13,7 @@ from __future__ import annotations -from cudf_polars.dsl.expressions.aggregation import Agg, Item +from cudf_polars.dsl.expressions.aggregation import Agg, Item, SortedAgg from cudf_polars.dsl.expressions.base import ( Col, ColRef, @@ -60,6 +60,7 @@ "Slice", "Sort", "SortBy", + "SortedAgg", "StringFunction", "StructFunction", "TemporalFunction", diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py b/python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py index 1f9fe15ae107..8825dd36bfd7 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py @@ -24,7 +24,7 @@ from cudf_polars.containers import DataFrame, DataType -__all__ = ["Agg", "Item"] +__all__ = ["Agg", "Item", "SortedAgg"] class Item(Expr): @@ -73,6 +73,42 @@ def do_evaluate( return value +class SortedAgg(Expr): + """Grouped ``first``/``last`` aggregation ordered by one or more expressions.""" + + __slots__ = ("name", "options") + _non_child = ("dtype", "name", "options") + + def __init__( + self, + dtype: DataType, + name: str, + options: tuple[Any, ...], + value: Expr, + *by: Expr, + ) -> None: + self.dtype = dtype + self.name = name + stable, nulls_last, descending = options + self.options = (stable, tuple(nulls_last), tuple(descending)) + self.children = (value, *by) + self.is_pointwise = False + if name not in {"first", "last"}: + raise NotImplementedError(f"Sorted aggregation {name=}") + if not by: + raise NotImplementedError( + "Sorted aggregation requires order-by expressions" + ) + if len(self.options[1]) != len(by) or len(self.options[2]) != len(by): + raise NotImplementedError( + "Sorted aggregation requires one null/descending option per order key" + ) + + @property + def agg_request(self) -> plc.aggregation.Aggregation: # noqa: D102 + raise NotImplementedError("Sorted aggregation is handled by GroupBy") + + class Agg(Expr): __slots__ = ("context", "name", "op", "options", "request") _non_child = ("dtype", "name", "options", "context") diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 559777f43c7d..fe8382ad7a3a 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -2155,6 +2155,101 @@ def __init__( self.zlice, ) + @staticmethod + def _align_to_group_keys( + group_keys: plc.Table, + result_group_keys: plc.Table, + result: Column, + stream: Any, + ) -> Column: + """Align a grouped result to the group-key order used by this GroupBy.""" + left_order, right_order = plc.join.inner_join( + group_keys, + result_group_keys, + plc.types.NullEquality.EQUAL, + stream=stream, + ) + (right_order,) = plc.sorting.sort_by_key( + plc.Table([right_order]), + plc.Table([left_order]), + [plc.types.Order.ASCENDING], + [plc.types.NullOrder.AFTER], + stream=stream, + ).columns() + (ordered_result,) = plc.copying.gather( + plc.Table([result.obj]), + right_order, + plc.copying.OutOfBoundsPolicy.DONT_CHECK, + stream=stream, + ).columns() + return Column(ordered_result, name=result.name, dtype=result.dtype) + + @staticmethod + def _evaluate_sorted_aggregation( + sorted_agg: expr.SortedAgg, + name: str, + keys: Sequence[Column], + df: DataFrame, + ) -> tuple[plc.Table, Column]: + """Evaluate a grouped first/last aggregation with an explicit order.""" + value_expr, *by_exprs = sorted_agg.children + value, *by = broadcast( + *( + child.evaluate(df, context=ExecutionContext.GROUPBY) + for child in (value_expr, *by_exprs) + ), + target_length=keys[0].size, + stream=df.stream, + ) + stable, nulls_last, descending = sorted_agg.options + key_order = [plc.types.Order.ASCENDING] * len(keys) + key_null_order = [plc.types.NullOrder.BEFORE] * len(keys) + by_order = [ + plc.types.Order.DESCENDING if is_descending else plc.types.Order.ASCENDING + for is_descending in descending + ] + by_null_order = [ + plc.types.NullOrder.AFTER if is_nulls_last else plc.types.NullOrder.BEFORE + for is_nulls_last in nulls_last + ] + do_sort = plc.sorting.stable_sort_by_key if stable else plc.sorting.sort_by_key + sorted_table = do_sort( + plc.Table([*(key.obj for key in keys), value.obj]), + plc.Table([*(key.obj for key in keys), *(col.obj for col in by)]), + [*key_order, *by_order], + [*key_null_order, *by_null_order], + stream=df.stream, + ) + sorted_keys = plc.Table(sorted_table.columns()[: len(keys)]) + sorted_value = sorted_table.columns()[len(keys)] + grouper = plc.groupby.GroupBy( + sorted_keys, + null_handling=plc.types.NullPolicy.INCLUDE, + keys_are_sorted=plc.types.Sorted.YES, + column_order=key_order, + null_precedence=key_null_order, + ) + offset = 0 if sorted_agg.name == "first" else -1 + group_keys, raw_tables = grouper.aggregate( + [ + plc.groupby.GroupByRequest( + sorted_value, + [ + plc.aggregation.nth_element( + offset, + null_handling=plc.types.NullPolicy.INCLUDE, + ) + ], + ) + ], + stream=df.stream, + ) + return group_keys, Column( + raw_tables[0].columns()[0], + name=name, + dtype=sorted_agg.dtype, + ) + @classmethod @log_do_evaluate @nvtx_annotate_cudf_polars(message="GroupBy") @@ -2188,13 +2283,17 @@ def do_evaluate( null_precedence=[k.null_order for k in keys], ) requests = [] - names = [] + sorted_requests = [] + names: list[str | None] = [] cast_to_schema = [] for request in agg_requests: should_cast = False name = request.name value = request.value - if isinstance(value, expr.Len): + if isinstance(value, expr.SortedAgg): + sorted_requests.append(request) + continue + elif isinstance(value, expr.Len): # A count aggregation, we need a column so use a key column col = keys[0].obj elif isinstance(value, expr.Agg): @@ -2210,29 +2309,48 @@ def do_evaluate( ) else: # Anything else, we pre-evaluate - column = value.evaluate(df, context=ExecutionContext.GROUPBY) - if column.size != keys[0].size: - column = broadcast( - column, target_length=keys[0].size, stream=df.stream + evaluated = value.evaluate(df, context=ExecutionContext.GROUPBY) + if evaluated.size != keys[0].size: + evaluated = broadcast( + evaluated, target_length=keys[0].size, stream=df.stream )[0] - col = column.obj + col = evaluated.obj requests.append(plc.groupby.GroupByRequest(col, [value.agg_request])) names.append(name) cast_to_schema.append(should_cast) + if not requests: + requests.append( + plc.groupby.GroupByRequest( + keys[0].obj, + [plc.aggregation.count(null_handling=plc.types.NullPolicy.INCLUDE)], + ) + ) + names.append(None) + cast_to_schema.append(False) group_keys, raw_tables = grouper.aggregate(requests, stream=df.stream) - results = [ - Column(column, name=name, dtype=schema[name]) - if not should_cast - else Column(column, name=name, dtype=schema[name]).astype( - schema[name], stream=df.stream + results = [] + for result_name, plc_column, should_cast in zip( + names, + itertools.chain.from_iterable(t.columns() for t in raw_tables), + cast_to_schema, + strict=True, + ): + if result_name is None: + continue + result = Column(plc_column, name=result_name, dtype=schema[result_name]) + if should_cast: + result = result.astype(schema[result_name], stream=df.stream) + results.append(result) + for request in sorted_requests: + assert isinstance(request.value, expr.SortedAgg) + result_group_keys, result = cls._evaluate_sorted_aggregation( + request.value, request.name, keys, df ) - for name, column, should_cast in zip( - names, - itertools.chain.from_iterable(t.columns() for t in raw_tables), - cast_to_schema, - strict=True, + results.append( + cls._align_to_group_keys( + group_keys, result_group_keys, result, df.stream + ) ) - ] result_keys = [ Column(grouped_key, name=key.name, dtype=key.dtype) for key, grouped_key in zip(keys, group_keys.columns(), strict=True) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py b/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py index 4b66ebb1846a..eafe12442a0b 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py @@ -218,6 +218,32 @@ def decompose_single_agg( child = agg.children[0] else: (child,) = agg.children + if ( + context == ExecutionContext.GROUPBY + and agg.name in {"first", "last"} + and isinstance(child, expr.SortBy) + ): + for sort_child in child.children: + child_aggs, _ = decompose_single_agg( + expr.NamedExpr(next(name_generator), sort_child), + name_generator, + is_top=False, + context=context, + ) + if any(has_agg for _, has_agg in child_aggs): + raise NotImplementedError( + "Nested aggs in sorted groupby aggregation not supported" + ) + return [ + ( + named_expr.reconstruct( + expr.SortedAgg( + agg.dtype, agg.name, child.options, *child.children + ) + ), + True, + ) + ], named_expr.reconstruct(expr.Col(agg.dtype, name)) # Fuse drop_nulls().n_unique() into nunique(null_handling=EXCLUDE) # rather than materializing a filtered intermediate column. if ( diff --git a/python/cudf_polars/cudf_polars/streaming/groupby.py b/python/cudf_polars/cudf_polars/streaming/groupby.py index 2d504ee112af..9cacda2dc115 100644 --- a/python/cudf_polars/cudf_polars/streaming/groupby.py +++ b/python/cudf_polars/cudf_polars/streaming/groupby.py @@ -22,6 +22,7 @@ Len, Literal, NamedExpr, + SortedAgg, StructFunction, Ternary, UnaryFunction, @@ -251,6 +252,51 @@ def _decompose_std_var( return selection, aggregations, reductions, False +def _decompose_sorted_agg( + name: str, expr: SortedAgg, *, names: Generator[str, None, None] +) -> tuple[NamedExpr, list[NamedExpr], list[NamedExpr], bool]: + """Carry the selected payload and order keys through grouped reductions.""" + stable, _, _ = expr.options + if stable: + raise NotImplementedError( + "stable sorted groupby aggregation requires an input-order tie-breaker" + ) + value, *by = expr.children + by_names = [f"{next(names)}__sort_key" for _ in by] + by_cols = [ + Col(order_by.dtype, by_name) + for order_by, by_name in zip(by, by_names, strict=True) + ] + + selection = NamedExpr(name, Col(expr.dtype, name)) + aggregations = [ + NamedExpr(name, SortedAgg(expr.dtype, expr.name, expr.options, value, *by)), + *( + NamedExpr( + by_name, + SortedAgg(order_by.dtype, expr.name, expr.options, order_by, *by), + ) + for order_by, by_name in zip(by, by_names, strict=True) + ), + ] + reductions = [ + NamedExpr( + name, + SortedAgg( + expr.dtype, expr.name, expr.options, Col(expr.dtype, name), *by_cols + ), + ), + *( + NamedExpr( + by_name, + SortedAgg(order_by.dtype, expr.name, expr.options, order_by, *by_cols), + ) + for order_by, by_name in zip(by_cols, by_names, strict=True) + ), + ] + return selection, aggregations, reductions, False + + def decompose( name: str, expr: Expr, *, names: Generator[str, None, None] ) -> tuple[NamedExpr, list[NamedExpr], list[NamedExpr], bool]: @@ -289,6 +335,8 @@ def decompose( ) ] return selection, aggregation, reduction, False + if isinstance(expr, SortedAgg): + return _decompose_sorted_agg(name, expr, names=names) if isinstance(expr, Agg): if (aggfunc := _GB_AGG_REDUCTIONS.get(expr.name)) is not None: if expr.name == "count": diff --git a/python/cudf_polars/tests/streaming/test_groupby.py b/python/cudf_polars/tests/streaming/test_groupby.py index 9bfb363147ce..1397f922967c 100644 --- a/python/cudf_polars/tests/streaming/test_groupby.py +++ b/python/cudf_polars/tests/streaming/test_groupby.py @@ -173,6 +173,26 @@ def test_groupby_agg(df, streaming_engine, op, keys): assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) +def test_groupby_sort_by_first_last(streaming_engine_factory): + streaming_engine = streaming_engine_factory( + StreamingOptions(max_rows_per_partition=2, fallback_mode="raise"), + ) + df = pl.LazyFrame( + { + "g": ["B", "A", "C", "A", "B", "C", "A", "B"], + "idx": [2, 3, 2, 1, 1, 1, 2, 3], + "val": [40, 30, 60, 10, 30, 50, 20, 50], + } + ) + + q = df.group_by("g").agg( + pl.col("val").sum().alias("volume"), + pl.col("val").sort_by("idx").first().alias("open"), + pl.col("val").sort_by("idx").last().alias("close"), + ) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) + + @pytest.mark.parametrize("ddof", [0, 2, 50]) @pytest.mark.parametrize("agg", ["std", "var"]) def test_groupby_std_var_ddof(df, engine, agg, ddof): diff --git a/python/cudf_polars/tests/streaming/test_spmd.py b/python/cudf_polars/tests/streaming/test_spmd.py index 326d39421265..dc287ce5e192 100644 --- a/python/cudf_polars/tests/streaming/test_spmd.py +++ b/python/cudf_polars/tests/streaming/test_spmd.py @@ -7,6 +7,7 @@ import os import uuid from itertools import pairwise +from pathlib import Path from typing import TYPE_CHECKING from unittest.mock import patch @@ -14,13 +15,14 @@ import polars as pl from polars import polars as plrs # type: ignore[attr-defined] +from polars.testing import assert_frame_equal import rmm.mr from rapidsmpf.bootstrap import is_running_with_rrun from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor import cudf_polars.quent -from cudf_polars.engine.core import _find_memory_error +from cudf_polars.engine.core import _find_memory_error, all_gather_host_data from cudf_polars.engine.hardware_binding import HardwareBindingPolicy from cudf_polars.engine.options import StreamingOptions from cudf_polars.engine.spmd import ( @@ -33,8 +35,6 @@ from cudf_polars.utils.config import MemoryResourceConfig if TYPE_CHECKING: - from pathlib import Path - from rapidsmpf.communicator.communicator import Communicator pytestmark = pytest.mark.spmd @@ -795,6 +795,55 @@ def test_over_nonscalar_duplicated_input( ) +def test_groupby_sort_by_first_last_multirank( + comm: Communicator, tmp_path: Path +) -> None: + with SPMDEngine( + comm=comm, + executor_options={ + "max_rows_per_partition": 2, + "dynamic_planning": {}, + "fallback_mode": "raise", + }, + ) as engine: + source_path = tmp_path / "groupby-sort-by.parquet" + if engine.rank == 0: + make_partitioned_source( + pl.DataFrame( + { + "g": ["B", "A", "C", "A", "B", "C", "A", "B"], + "idx": [2, 3, 2, 1, 1, 1, 2, 3], + "val": [40, 30, 60, 10, 30, 50, 20, 50], + } + ), + source_path, + "parquet", + row_group_size=2, + ) + data = str(source_path).encode() + else: + data = b"" + + with reserve_op_id() as op_id: + source_paths = all_gather_host_data( + engine.comm, engine.context.br(), op_id, data + ) + source_path = Path(source_paths[0].decode()) + + q = ( + pl.scan_parquet(source_path) + .group_by("g") + .agg( + pl.col("val").sum().alias("volume"), + pl.col("val").sort_by("idx").first().alias("open"), + pl.col("val").sort_by("idx").last().alias("close"), + ) + ) + expected = q.collect() + got = q.collect(engine=engine) + assert_frame_equal(expected, got, check_row_order=False) + + def test_find_memory_error() -> None: err = MemoryError("oom") assert _find_memory_error(err) is err diff --git a/python/cudf_polars/tests/test_groupby.py b/python/cudf_polars/tests/test_groupby.py index 4030727e54e0..a50e37af52b5 100644 --- a/python/cudf_polars/tests/test_groupby.py +++ b/python/cudf_polars/tests/test_groupby.py @@ -424,6 +424,86 @@ def test_groupby_null_keys(engine: pl.GPUEngine, maintain_order): assert_gpu_result_equal(q, engine=engine) +def test_groupby_sort_by_first_last_in_memory(in_memory_engine: pl.GPUEngine) -> None: + df = pl.LazyFrame( + { + "g": ["B", "A", "C", "A", "B", "C"], + "idx": [2, 2, 2, 1, 1, 1], + "val": [40, 20, 60, 10, 30, 50], + } + ) + + q = df.group_by("g", maintain_order=True).agg( + pl.col("val").sum().alias("volume"), + pl.col("val").sort_by("idx").first().alias("open"), + pl.col("val").sort_by("idx").last().alias("close"), + ) + assert_gpu_result_equal(q, engine=in_memory_engine, check_row_order=True) + + +def test_groupby_sort_by_first_pointwise_in_memory( + in_memory_engine: pl.GPUEngine, +) -> None: + df = pl.LazyFrame( + { + "g": ["A", "A", "A", "B", "B"], + "idx": [2, None, 2, None, 1], + "seq": [1, 2, 3, 1, 2], + "val": [10, 20, 30, 40, 50], + } + ) + + q = ( + df.group_by("g") + .agg( + (pl.col("val") * 2) + .sort_by( + pl.col("idx").fill_null(99), + pl.col("seq"), + descending=[False, True], + nulls_last=[False, False], + ) + .first() + .alias("picked") + ) + .sort("g") + ) + assert_gpu_result_equal(q, engine=in_memory_engine) + + +def test_groupby_sort_by_first_stable_in_memory( + in_memory_engine: pl.GPUEngine, +) -> None: + df = pl.LazyFrame( + { + "g": ["B", "A", "A", "A", "B"], + "idx": [1, 1, 1, 1, 1], + "val": [50, 10, 20, 30, 60], + } + ) + + q = df.group_by("g", maintain_order=True).agg( + pl.col("val").sort_by("idx", maintain_order=True).first().alias("first_tie"), + pl.col("val").sort_by("idx", maintain_order=True).last().alias("last_tie"), + ) + assert_gpu_result_equal(q, engine=in_memory_engine) + + +def test_groupby_sort_by_drop_nulls_first_raises( + in_memory_engine: pl.GPUEngine, +) -> None: + df = pl.LazyFrame( + { + "g": ["A", "A", "B"], + "idx": [2, 1, 1], + "val": [None, 10, 20], + } + ) + + q = df.group_by("g").agg(pl.col("val").sort_by("idx").drop_nulls().first()) + assert_ir_translation_raises(q, in_memory_engine, NotImplementedError) + + @pytest.mark.xfail(reason="https://github.com/pola-rs/polars/issues/17513") def test_groupby_minmax_with_nan(engine: pl.GPUEngine): df = pl.LazyFrame( From 0f851dd81d4ce152b32bdb79390d87fe4f2ac460 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 5 Aug 2026 07:35:58 -0700 Subject: [PATCH 02/13] update comments --- .../cudf_polars/cudf_polars/dsl/expressions/aggregation.py | 7 +++++-- python/cudf_polars/cudf_polars/streaming/groupby.py | 3 ++- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py b/python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py index 8825dd36bfd7..7c49155ca512 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py @@ -74,7 +74,7 @@ def do_evaluate( class SortedAgg(Expr): - """Grouped ``first``/``last`` aggregation ordered by one or more expressions.""" + """``first``/``last`` aggregation ordered by one or more expressions.""" __slots__ = ("name", "options") _non_child = ("dtype", "name", "options") @@ -106,7 +106,10 @@ def __init__( @property def agg_request(self) -> plc.aggregation.Aggregation: # noqa: D102 - raise NotImplementedError("Sorted aggregation is handled by GroupBy") + raise NotImplementedError( + "Sorted aggregation cannot be represented as a pylibcudf " + "aggregation request" + ) class Agg(Expr): diff --git a/python/cudf_polars/cudf_polars/streaming/groupby.py b/python/cudf_polars/cudf_polars/streaming/groupby.py index 9cacda2dc115..588c40870aa7 100644 --- a/python/cudf_polars/cudf_polars/streaming/groupby.py +++ b/python/cudf_polars/cudf_polars/streaming/groupby.py @@ -259,7 +259,8 @@ def _decompose_sorted_agg( stable, _, _ = expr.options if stable: raise NotImplementedError( - "stable sorted groupby aggregation requires an input-order tie-breaker" + "group_by sort_by(..., maintain_order=True).first/last is not " + "supported for multiple partitions" ) value, *by = expr.children by_names = [f"{next(names)}__sort_key" for _ in by] From 592f487b4e337a0bbe77291881cc46eeb4e7423b Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 5 Aug 2026 09:00:14 -0700 Subject: [PATCH 03/13] elaborate in the docstring --- .../cudf_polars/dsl/expressions/aggregation.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py b/python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py index 7c49155ca512..ef94763bb668 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/aggregation.py @@ -74,7 +74,16 @@ def do_evaluate( class SortedAgg(Expr): - """``first``/``last`` aggregation ordered by one or more expressions.""" + """ + ``first``/``last`` aggregation ordered by one or more expressions. + + Notes + ----- + This expression is used by GroupBy infrastructure for ordered + first/last aggregations. The ordering may depend on columns other than + the value being aggregated, so this cannot be evaluated independently + nor can it be represented as a plain :class:`Agg` variant. + """ __slots__ = ("name", "options") _non_child = ("dtype", "name", "options") From 96356d0526c12e62f396c9a40e0090251099f21b Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 5 Aug 2026 12:15:38 -0700 Subject: [PATCH 04/13] self review and cleanup attempt --- python/cudf_polars/cudf_polars/dsl/ir.py | 71 ++++++++++++++---------- 1 file changed, 43 insertions(+), 28 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index fe8382ad7a3a..3cb3242dc2d9 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -2184,6 +2184,24 @@ def _align_to_group_keys( ).columns() return Column(ordered_result, name=result.name, dtype=result.dtype) + @staticmethod + def _collect_group_keys( + grouper: plc.groupby.GroupBy, + key: Column, + stream: Any, + ) -> plc.Table: + """Collect group keys when the GroupBy has no aggregation requests.""" + group_keys, _ = grouper.aggregate( + [ + plc.groupby.GroupByRequest( + key.obj, + [plc.aggregation.count(null_handling=plc.types.NullPolicy.INCLUDE)], + ) + ], + stream=stream, + ) + return group_keys + @staticmethod def _evaluate_sorted_aggregation( sorted_agg: expr.SortedAgg, @@ -2284,7 +2302,7 @@ def do_evaluate( ) requests = [] sorted_requests = [] - names: list[str | None] = [] + names: list[str] = [] cast_to_schema = [] for request in agg_requests: should_cast = False @@ -2318,39 +2336,36 @@ def do_evaluate( requests.append(plc.groupby.GroupByRequest(col, [value.agg_request])) names.append(name) cast_to_schema.append(should_cast) - if not requests: - requests.append( - plc.groupby.GroupByRequest( - keys[0].obj, - [plc.aggregation.count(null_handling=plc.types.NullPolicy.INCLUDE)], - ) - ) - names.append(None) - cast_to_schema.append(False) - group_keys, raw_tables = grouper.aggregate(requests, stream=df.stream) - results = [] - for result_name, plc_column, should_cast in zip( - names, - itertools.chain.from_iterable(t.columns() for t in raw_tables), - cast_to_schema, - strict=True, - ): - if result_name is None: - continue - result = Column(plc_column, name=result_name, dtype=schema[result_name]) - if should_cast: - result = result.astype(schema[result_name], stream=df.stream) - results.append(result) + group_keys: plc.Table | None = None + results: list[Column] = [] + if requests: + group_keys, raw_tables = grouper.aggregate(requests, stream=df.stream) + for result_name, plc_column, should_cast in zip( + names, + itertools.chain.from_iterable(t.columns() for t in raw_tables), + cast_to_schema, + strict=True, + ): + result = Column(plc_column, name=result_name, dtype=schema[result_name]) + if should_cast: + result = result.astype(schema[result_name], stream=df.stream) + results.append(result) for request in sorted_requests: assert isinstance(request.value, expr.SortedAgg) result_group_keys, result = cls._evaluate_sorted_aggregation( request.value, request.name, keys, df ) - results.append( - cls._align_to_group_keys( - group_keys, result_group_keys, result, df.stream + if group_keys is None: + group_keys = result_group_keys + results.append(result) + else: + results.append( + cls._align_to_group_keys( + group_keys, result_group_keys, result, df.stream + ) ) - ) + if group_keys is None: + group_keys = cls._collect_group_keys(grouper, keys[0], df.stream) result_keys = [ Column(grouped_key, name=key.name, dtype=key.dtype) for key, grouped_key in zip(keys, group_keys.columns(), strict=True) From 44f7c9c0347a00b921bb03317a7c5e87084148a3 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 5 Aug 2026 12:41:09 -0700 Subject: [PATCH 05/13] reduce diff --- python/cudf_polars/cudf_polars/dsl/ir.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 3cb3242dc2d9..2461fd63f090 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -2327,12 +2327,12 @@ def do_evaluate( ) else: # Anything else, we pre-evaluate - evaluated = value.evaluate(df, context=ExecutionContext.GROUPBY) - if evaluated.size != keys[0].size: - evaluated = broadcast( - evaluated, target_length=keys[0].size, stream=df.stream + column = value.evaluate(df, context=ExecutionContext.GROUPBY) + if column.size != keys[0].size: + column = broadcast( + column, target_length=keys[0].size, stream=df.stream )[0] - col = evaluated.obj + col = column.obj requests.append(plc.groupby.GroupByRequest(col, [value.agg_request])) names.append(name) cast_to_schema.append(should_cast) From b805bc472190bd987b20bd6194f28ed91833ff9c Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 5 Aug 2026 13:06:07 -0700 Subject: [PATCH 06/13] experimental refactor --- python/cudf_polars/cudf_polars/dsl/ir.py | 165 ++++++++++++++--------- 1 file changed, 105 insertions(+), 60 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 2461fd63f090..c0dca5610b3e 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -2203,7 +2203,76 @@ def _collect_group_keys( return group_keys @staticmethod + def _temporary_value_name(name: str, keys: Sequence[Column]) -> str: + """Return a temporary value name that cannot shadow a group key.""" + used = {key.name for key in keys} + value_name = name + while value_name in used: + value_name = f"_{value_name}" + return value_name + + @staticmethod + def _evaluate_aggregation_requests( + schema: Schema, + keys: Sequence[Column], + grouper: plc.groupby.GroupBy, + agg_requests: Sequence[expr.NamedExpr], + df: DataFrame, + ) -> tuple[plc.Table | None, list[Column]]: + """Evaluate ordinary grouped aggregation requests.""" + requests = [] + names: list[str] = [] + cast_to_schema = [] + for request in agg_requests: + should_cast = False + name = request.name + value = request.value + if isinstance(value, expr.Len): + # A count aggregation, we need a column so use a key column + col = keys[0].obj + elif isinstance(value, expr.Agg): + if value.name == "quantile": + child = value.children[0] + else: + (child,) = value.children + # libcudf will return int64 when summing integers + # but the schema may be a lower bit width + col = child.evaluate(df, context=ExecutionContext.GROUPBY).obj + should_cast = value.name == "sum" and plc.traits.is_integral_not_bool( + col.type() + ) + else: + # Anything else, we pre-evaluate + column = value.evaluate(df, context=ExecutionContext.GROUPBY) + if column.size != keys[0].size: + column = broadcast( + column, target_length=keys[0].size, stream=df.stream + )[0] + col = column.obj + requests.append(plc.groupby.GroupByRequest(col, [value.agg_request])) + names.append(name) + cast_to_schema.append(should_cast) + + if not requests: + return None, [] + + group_keys, raw_tables = grouper.aggregate(requests, stream=df.stream) + results = [] + for result_name, plc_column, should_cast in zip( + names, + itertools.chain.from_iterable(t.columns() for t in raw_tables), + cast_to_schema, + strict=True, + ): + result = Column(plc_column, name=result_name, dtype=schema[result_name]) + if should_cast: + result = result.astype(schema[result_name], stream=df.stream) + results.append(result) + return group_keys, results + + @classmethod def _evaluate_sorted_aggregation( + cls, sorted_agg: expr.SortedAgg, name: str, keys: Sequence[Column], @@ -2239,7 +2308,21 @@ def _evaluate_sorted_aggregation( stream=df.stream, ) sorted_keys = plc.Table(sorted_table.columns()[: len(keys)]) - sorted_value = sorted_table.columns()[len(keys)] + sorted_key_columns = [ + Column(column, name=key.name, dtype=key.dtype) + for key, column in zip(keys, sorted_keys.columns(), strict=True) + ] + value_name = cls._temporary_value_name(name, keys) + sorted_value = Column( + sorted_table.columns()[len(keys)], + name=value_name, + dtype=sorted_agg.dtype, + ) + sorted_df = DataFrame( + [*sorted_key_columns, sorted_value], + stream=df.stream, + num_rows=sorted_table.num_rows(), + ) grouper = plc.groupby.GroupBy( sorted_keys, null_handling=plc.types.NullPolicy.INCLUDE, @@ -2247,26 +2330,27 @@ def _evaluate_sorted_aggregation( column_order=key_order, null_precedence=key_null_order, ) - offset = 0 if sorted_agg.name == "first" else -1 - group_keys, raw_tables = grouper.aggregate( + group_keys, results = cls._evaluate_aggregation_requests( + {name: sorted_agg.dtype}, + sorted_key_columns, + grouper, [ - plc.groupby.GroupByRequest( - sorted_value, - [ - plc.aggregation.nth_element( - offset, - null_handling=plc.types.NullPolicy.INCLUDE, - ) - ], + expr.NamedExpr( + name, + expr.Agg( + sorted_agg.dtype, + sorted_agg.name, + None, + ExecutionContext.GROUPBY, + expr.Col(sorted_agg.dtype, value_name), + ), ) ], - stream=df.stream, - ) - return group_keys, Column( - raw_tables[0].columns()[0], - name=name, - dtype=sorted_agg.dtype, + sorted_df, ) + assert group_keys is not None + (result,) = results + return group_keys, result @classmethod @log_do_evaluate @@ -2302,54 +2386,15 @@ def do_evaluate( ) requests = [] sorted_requests = [] - names: list[str] = [] - cast_to_schema = [] for request in agg_requests: - should_cast = False - name = request.name value = request.value if isinstance(value, expr.SortedAgg): sorted_requests.append(request) - continue - elif isinstance(value, expr.Len): - # A count aggregation, we need a column so use a key column - col = keys[0].obj - elif isinstance(value, expr.Agg): - if value.name == "quantile": - child = value.children[0] - else: - (child,) = value.children - # libcudf will return int64 when summing integers - # but the schema may be a lower bit width - col = child.evaluate(df, context=ExecutionContext.GROUPBY).obj - should_cast = value.name == "sum" and plc.traits.is_integral_not_bool( - col.type() - ) else: - # Anything else, we pre-evaluate - column = value.evaluate(df, context=ExecutionContext.GROUPBY) - if column.size != keys[0].size: - column = broadcast( - column, target_length=keys[0].size, stream=df.stream - )[0] - col = column.obj - requests.append(plc.groupby.GroupByRequest(col, [value.agg_request])) - names.append(name) - cast_to_schema.append(should_cast) - group_keys: plc.Table | None = None - results: list[Column] = [] - if requests: - group_keys, raw_tables = grouper.aggregate(requests, stream=df.stream) - for result_name, plc_column, should_cast in zip( - names, - itertools.chain.from_iterable(t.columns() for t in raw_tables), - cast_to_schema, - strict=True, - ): - result = Column(plc_column, name=result_name, dtype=schema[result_name]) - if should_cast: - result = result.astype(schema[result_name], stream=df.stream) - results.append(result) + requests.append(request) + group_keys, results = cls._evaluate_aggregation_requests( + schema, keys, grouper, requests, df + ) for request in sorted_requests: assert isinstance(request.value, expr.SortedAgg) result_group_keys, result = cls._evaluate_sorted_aggregation( From 3ac1db209dfb251fe436c4184b8c567661702c9e Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 5 Aug 2026 13:44:57 -0700 Subject: [PATCH 07/13] reuse existing utils --- python/cudf_polars/cudf_polars/dsl/ir.py | 252 ++++++++++++++--------- 1 file changed, 157 insertions(+), 95 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index c0dca5610b3e..2f1ebeda72fd 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -50,6 +50,7 @@ from cudf_polars.dsl.nodebase import Node from cudf_polars.dsl.to_ast import _DECIMAL_IDS, to_ast, to_parquet_filter from cudf_polars.dsl.tracing import log_do_evaluate, nvtx_annotate_cudf_polars +from cudf_polars.dsl.utils.naming import unique_names from cudf_polars.dsl.utils.reshape import broadcast from cudf_polars.dsl.utils.windows import ( offsets_to_windows, @@ -2202,15 +2203,6 @@ def _collect_group_keys( ) return group_keys - @staticmethod - def _temporary_value_name(name: str, keys: Sequence[Column]) -> str: - """Return a temporary value name that cannot shadow a group key.""" - used = {key.name for key in keys} - value_name = name - while value_name in used: - value_name = f"_{value_name}" - return value_name - @staticmethod def _evaluate_aggregation_requests( schema: Schema, @@ -2271,86 +2263,149 @@ def _evaluate_aggregation_requests( return group_keys, results @classmethod - def _evaluate_sorted_aggregation( + def _evaluate_sorted_aggregations( cls, - sorted_agg: expr.SortedAgg, - name: str, + sorted_requests: Sequence[expr.NamedExpr], keys: Sequence[Column], df: DataFrame, - ) -> tuple[plc.Table, Column]: - """Evaluate a grouped first/last aggregation with an explicit order.""" - value_expr, *by_exprs = sorted_agg.children - value, *by = broadcast( - *( - child.evaluate(df, context=ExecutionContext.GROUPBY) - for child in (value_expr, *by_exprs) - ), - target_length=keys[0].size, - stream=df.stream, - ) - stable, nulls_last, descending = sorted_agg.options + ) -> tuple[plc.Table | None, list[Column]]: + """Evaluate grouped first/last aggregations with explicit ordering.""" + if not sorted_requests: + return None, [] + + request_groups: dict[ + tuple[Any, tuple[expr.Expr, ...]], list[expr.NamedExpr] + ] = {} + for request in sorted_requests: + assert isinstance(request.value, expr.SortedAgg) + by_exprs = request.value.children[1:] + request_groups.setdefault( + (request.value.options, tuple(by_exprs)), [] + ).append(request) + + common_group_keys: plc.Table | None = None + results_by_name: dict[str, Column] = {} key_order = [plc.types.Order.ASCENDING] * len(keys) key_null_order = [plc.types.NullOrder.BEFORE] * len(keys) - by_order = [ - plc.types.Order.DESCENDING if is_descending else plc.types.Order.ASCENDING - for is_descending in descending - ] - by_null_order = [ - plc.types.NullOrder.AFTER if is_nulls_last else plc.types.NullOrder.BEFORE - for is_nulls_last in nulls_last - ] - do_sort = plc.sorting.stable_sort_by_key if stable else plc.sorting.sort_by_key - sorted_table = do_sort( - plc.Table([*(key.obj for key in keys), value.obj]), - plc.Table([*(key.obj for key in keys), *(col.obj for col in by)]), - [*key_order, *by_order], - [*key_null_order, *by_null_order], - stream=df.stream, - ) - sorted_keys = plc.Table(sorted_table.columns()[: len(keys)]) - sorted_key_columns = [ - Column(column, name=key.name, dtype=key.dtype) - for key, column in zip(keys, sorted_keys.columns(), strict=True) - ] - value_name = cls._temporary_value_name(name, keys) - sorted_value = Column( - sorted_table.columns()[len(keys)], - name=value_name, - dtype=sorted_agg.dtype, - ) - sorted_df = DataFrame( - [*sorted_key_columns, sorted_value], - stream=df.stream, - num_rows=sorted_table.num_rows(), - ) - grouper = plc.groupby.GroupBy( - sorted_keys, - null_handling=plc.types.NullPolicy.INCLUDE, - keys_are_sorted=plc.types.Sorted.YES, - column_order=key_order, - null_precedence=key_null_order, - ) - group_keys, results = cls._evaluate_aggregation_requests( - {name: sorted_agg.dtype}, - sorted_key_columns, - grouper, - [ - expr.NamedExpr( - name, - expr.Agg( - sorted_agg.dtype, - sorted_agg.name, - None, - ExecutionContext.GROUPBY, - expr.Col(sorted_agg.dtype, value_name), - ), + + for group in request_groups.values(): + first_agg = group[0].value + assert isinstance(first_agg, expr.SortedAgg) + value_exprs = [] + for request in group: + assert isinstance(request.value, expr.SortedAgg) + value_exprs.append(request.value.children[0]) + by_exprs = first_agg.children[1:] + columns = broadcast( + *( + child.evaluate(df, context=ExecutionContext.GROUPBY) + for child in (*value_exprs, *by_exprs) + ), + target_length=keys[0].size, + stream=df.stream, + ) + values = columns[: len(value_exprs)] + by = columns[len(value_exprs) :] + stable, nulls_last, descending = first_agg.options + by_order = [ + plc.types.Order.DESCENDING + if is_descending + else plc.types.Order.ASCENDING + for is_descending in descending + ] + by_null_order = [ + plc.types.NullOrder.AFTER + if is_nulls_last + else plc.types.NullOrder.BEFORE + for is_nulls_last in nulls_last + ] + do_sort = ( + plc.sorting.stable_sort_by_key if stable else plc.sorting.sort_by_key + ) + sorted_table = do_sort( + plc.Table( + [*(key.obj for key in keys), *(value.obj for value in values)] + ), + plc.Table([*(key.obj for key in keys), *(col.obj for col in by)]), + [*key_order, *by_order], + [*key_null_order, *by_null_order], + stream=df.stream, + ) + sorted_keys = plc.Table(sorted_table.columns()[: len(keys)]) + sorted_key_columns = [ + Column(column, name=key.name, dtype=key.dtype) + for key, column in zip(keys, sorted_keys.columns(), strict=True) + ] + value_names = unique_names( + ( + *(key.name for key in keys if key.name is not None), + *(request.name for request in group), ) - ], - sorted_df, - ) - assert group_keys is not None - (result,) = results - return group_keys, result + ) + sorted_values = [] + ordinary_requests = [] + schema = {} + for request, column in zip( + group, + sorted_table.columns()[len(keys) :], + strict=True, + ): + sorted_agg = request.value + assert isinstance(sorted_agg, expr.SortedAgg) + value_name = next(value_names) + sorted_values.append( + Column(column, name=value_name, dtype=sorted_agg.dtype) + ) + ordinary_requests.append( + expr.NamedExpr( + request.name, + expr.Agg( + sorted_agg.dtype, + sorted_agg.name, + None, + ExecutionContext.GROUPBY, + expr.Col(sorted_agg.dtype, value_name), + ), + ) + ) + schema[request.name] = sorted_agg.dtype + sorted_df = DataFrame( + [*sorted_key_columns, *sorted_values], + stream=df.stream, + num_rows=sorted_table.num_rows(), + ) + grouper = plc.groupby.GroupBy( + sorted_keys, + null_handling=plc.types.NullPolicy.INCLUDE, + keys_are_sorted=plc.types.Sorted.YES, + column_order=key_order, + null_precedence=key_null_order, + ) + group_keys, results = cls._evaluate_aggregation_requests( + schema, + sorted_key_columns, + grouper, + ordinary_requests, + sorted_df, + ) + assert group_keys is not None + if common_group_keys is None: + common_group_keys = group_keys + else: + results = [ + cls._align_to_group_keys( + common_group_keys, group_keys, result, df.stream + ) + for result in results + ] + results_by_name.update( + (request.name, result) + for request, result in zip(group, results, strict=True) + ) + + return common_group_keys, [ + results_by_name[request.name] for request in sorted_requests + ] @classmethod @log_do_evaluate @@ -2384,8 +2439,8 @@ def do_evaluate( column_order=[k.order for k in keys], null_precedence=[k.null_order for k in keys], ) - requests = [] - sorted_requests = [] + requests: list[expr.NamedExpr] = [] + sorted_requests: list[expr.NamedExpr] = [] for request in agg_requests: value = request.value if isinstance(value, expr.SortedAgg): @@ -2395,22 +2450,29 @@ def do_evaluate( group_keys, results = cls._evaluate_aggregation_requests( schema, keys, grouper, requests, df ) - for request in sorted_requests: - assert isinstance(request.value, expr.SortedAgg) - result_group_keys, result = cls._evaluate_sorted_aggregation( - request.value, request.name, keys, df - ) + sorted_group_keys, sorted_results = cls._evaluate_sorted_aggregations( + sorted_requests, keys, df + ) + if sorted_group_keys is not None: if group_keys is None: - group_keys = result_group_keys - results.append(result) + group_keys = sorted_group_keys else: - results.append( + sorted_results = [ cls._align_to_group_keys( - group_keys, result_group_keys, result, df.stream + group_keys, sorted_group_keys, result, df.stream ) - ) + for result in sorted_results + ] if group_keys is None: group_keys = cls._collect_group_keys(grouper, keys[0], df.stream) + results_by_name = { + request.name: result + for request, result in itertools.chain( + zip(requests, results, strict=True), + zip(sorted_requests, sorted_results, strict=True), + ) + } + results = [results_by_name[request.name] for request in agg_requests] result_keys = [ Column(grouped_key, name=key.name, dtype=key.dtype) for key, grouped_key in zip(keys, group_keys.columns(), strict=True) From 2ac42dfa4c356e1864a677f3e0c9be35a9fd9332 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 7 Aug 2026 18:39:34 -0700 Subject: [PATCH 08/13] support maintain_order=True --- python/cudf_polars/cudf_polars/dsl/ir.py | 21 +++------- .../streaming/actor_graph/groupby.py | 7 +++- .../cudf_polars/streaming/groupby.py | 6 --- .../tests/streaming/test_groupby.py | 42 +++++++++++++++++++ .../cudf_polars/tests/streaming/test_spmd.py | 9 ++++ python/cudf_polars/tests/test_groupby.py | 22 ++++++++++ 6 files changed, 84 insertions(+), 23 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 2f1ebeda72fd..ebe939b56f66 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -56,7 +56,7 @@ offsets_to_windows, range_window_bounds, ) -from cudf_polars.utils import dtypes +from cudf_polars.utils import dtypes, sorting from cudf_polars.utils.cuda_stream import ( get_cuda_stream, stream_ordered_after, @@ -2285,8 +2285,8 @@ def _evaluate_sorted_aggregations( common_group_keys: plc.Table | None = None results_by_name: dict[str, Column] = {} - key_order = [plc.types.Order.ASCENDING] * len(keys) - key_null_order = [plc.types.NullOrder.BEFORE] * len(keys) + key_order = [key.order for key in keys] + key_null_order = [key.null_order for key in keys] for group in request_groups.values(): first_agg = group[0].value @@ -2307,18 +2307,9 @@ def _evaluate_sorted_aggregations( values = columns[: len(value_exprs)] by = columns[len(value_exprs) :] stable, nulls_last, descending = first_agg.options - by_order = [ - plc.types.Order.DESCENDING - if is_descending - else plc.types.Order.ASCENDING - for is_descending in descending - ] - by_null_order = [ - plc.types.NullOrder.AFTER - if is_nulls_last - else plc.types.NullOrder.BEFORE - for is_nulls_last in nulls_last - ] + by_order, by_null_order = sorting.sort_order( + descending, nulls_last=nulls_last, num_keys=len(by_exprs) + ) do_sort = ( plc.sorting.stable_sort_by_key if stable else plc.sorting.sort_by_key ) 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 f0c4979fef6f..a9a878e29d75 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py @@ -19,7 +19,7 @@ from rapidsmpf.streaming.core.context import Context from cudf_polars.containers import DataType -from cudf_polars.dsl.expr import Col, NamedExpr +from cudf_polars.dsl.expr import Col, NamedExpr, SortedAgg from cudf_polars.dsl.ir import IR, Distinct, GroupBy, Select from cudf_polars.dsl.utils.naming import names_to_indices, unique_names from cudf_polars.streaming.actor_graph.collectives.shuffle import ShuffleManager @@ -498,7 +498,10 @@ def _key_indices( def _maintain_order(ir: GroupBy | Distinct) -> bool: if isinstance(ir, GroupBy): - return ir.maintain_order + return ir.maintain_order or any( + isinstance(ne.value, SortedAgg) and ne.value.options[0] + for ne in ir.agg_requests + ) else: return ir.stable or ir.keep in ( plc.stream_compaction.DuplicateKeepOption.KEEP_FIRST, diff --git a/python/cudf_polars/cudf_polars/streaming/groupby.py b/python/cudf_polars/cudf_polars/streaming/groupby.py index 588c40870aa7..5a3f4c0fb865 100644 --- a/python/cudf_polars/cudf_polars/streaming/groupby.py +++ b/python/cudf_polars/cudf_polars/streaming/groupby.py @@ -256,12 +256,6 @@ def _decompose_sorted_agg( name: str, expr: SortedAgg, *, names: Generator[str, None, None] ) -> tuple[NamedExpr, list[NamedExpr], list[NamedExpr], bool]: """Carry the selected payload and order keys through grouped reductions.""" - stable, _, _ = expr.options - if stable: - raise NotImplementedError( - "group_by sort_by(..., maintain_order=True).first/last is not " - "supported for multiple partitions" - ) value, *by = expr.children by_names = [f"{next(names)}__sort_key" for _ in by] by_cols = [ diff --git a/python/cudf_polars/tests/streaming/test_groupby.py b/python/cudf_polars/tests/streaming/test_groupby.py index 1397f922967c..b50f3351ad3b 100644 --- a/python/cudf_polars/tests/streaming/test_groupby.py +++ b/python/cudf_polars/tests/streaming/test_groupby.py @@ -193,6 +193,48 @@ def test_groupby_sort_by_first_last(streaming_engine_factory): assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) +def test_groupby_sort_by_first_last_stable_ties(streaming_engine_factory): + streaming_engine = streaming_engine_factory( + StreamingOptions(max_rows_per_partition=2, fallback_mode="raise"), + ) + df = pl.LazyFrame( + { + "g": ["A", "B", "A", "B", "A", "B"], + "idx": [1, 1, 1, 1, 1, 1], + "val": [10, 40, 20, 50, 30, 60], + } + ) + + q = df.group_by("g").agg( + pl.col("val").sort_by("idx", maintain_order=True).first().alias("first_tie"), + pl.col("val").sort_by("idx", maintain_order=True).last().alias("last_tie"), + ) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) + + +def test_groupby_sort_by_preserves_sorted_key_order(streaming_engine_factory): + streaming_engine = streaming_engine_factory( + StreamingOptions(max_rows_per_partition=2, fallback_mode="raise"), + ) + df = pl.LazyFrame( + { + "g": ["A", "A", "B", "B", "C", "C"], + "idx": [1, 2, 1, 2, 1, 2], + "val": [10, 20, 30, 40, 50, 60], + } + ) + + q = ( + df.sort("g", descending=True) + .group_by("g", maintain_order=True) + .agg( + pl.col("val").sort_by("idx").first().alias("open"), + pl.col("val").sort_by("idx").last().alias("close"), + ) + ) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=True) + + @pytest.mark.parametrize("ddof", [0, 2, 50]) @pytest.mark.parametrize("agg", ["std", "var"]) def test_groupby_std_var_ddof(df, engine, agg, ddof): diff --git a/python/cudf_polars/tests/streaming/test_spmd.py b/python/cudf_polars/tests/streaming/test_spmd.py index dc287ce5e192..5924aaabf6be 100644 --- a/python/cudf_polars/tests/streaming/test_spmd.py +++ b/python/cudf_polars/tests/streaming/test_spmd.py @@ -813,6 +813,7 @@ def test_groupby_sort_by_first_last_multirank( { "g": ["B", "A", "C", "A", "B", "C", "A", "B"], "idx": [2, 3, 2, 1, 1, 1, 2, 3], + "tie": [1, 1, 1, 1, 1, 1, 1, 1], "val": [40, 30, 60, 10, 30, 50, 20, 50], } ), @@ -837,6 +838,14 @@ def test_groupby_sort_by_first_last_multirank( pl.col("val").sum().alias("volume"), pl.col("val").sort_by("idx").first().alias("open"), pl.col("val").sort_by("idx").last().alias("close"), + pl.col("val") + .sort_by("tie", maintain_order=True) + .first() + .alias("first_tie"), + pl.col("val") + .sort_by("tie", maintain_order=True) + .last() + .alias("last_tie"), ) ) expected = q.collect() diff --git a/python/cudf_polars/tests/test_groupby.py b/python/cudf_polars/tests/test_groupby.py index a50e37af52b5..8c59c2616706 100644 --- a/python/cudf_polars/tests/test_groupby.py +++ b/python/cudf_polars/tests/test_groupby.py @@ -489,6 +489,28 @@ def test_groupby_sort_by_first_stable_in_memory( assert_gpu_result_equal(q, engine=in_memory_engine) +def test_groupby_sort_by_preserves_sorted_key_order_in_memory( + in_memory_engine: pl.GPUEngine, +) -> None: + df = pl.LazyFrame( + { + "g": ["A", "A", "B", "B", "C", "C"], + "idx": [1, 2, 1, 2, 1, 2], + "val": [10, 20, 30, 40, 50, 60], + } + ) + + q = ( + df.sort("g", descending=True) + .group_by("g", maintain_order=True) + .agg( + pl.col("val").sort_by("idx").first().alias("open"), + pl.col("val").sort_by("idx").last().alias("close"), + ) + ) + assert_gpu_result_equal(q, engine=in_memory_engine, check_row_order=True) + + def test_groupby_sort_by_drop_nulls_first_raises( in_memory_engine: pl.GPUEngine, ) -> None: From e7159ea685dc7ab271f8dc803a5769332e7bce5b Mon Sep 17 00:00:00 2001 From: rjzamora Date: Mon, 10 Aug 2026 06:30:49 -0700 Subject: [PATCH 09/13] code cov --- python/cudf_polars/tests/test_groupby.py | 35 +++++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/tests/test_groupby.py b/python/cudf_polars/tests/test_groupby.py index 8c59c2616706..80763ac0d9ff 100644 --- a/python/cudf_polars/tests/test_groupby.py +++ b/python/cudf_polars/tests/test_groupby.py @@ -13,7 +13,7 @@ from cudf_polars import Translator from cudf_polars.containers import DataType -from cudf_polars.dsl.expressions.aggregation import Agg +from cudf_polars.dsl.expressions.aggregation import Agg, SortedAgg from cudf_polars.dsl.expressions.base import Col, ExecutionContext, NamedExpr from cudf_polars.dsl.utils.aggregations import decompose_single_agg from cudf_polars.testing.asserts import ( @@ -526,6 +526,39 @@ def test_groupby_sort_by_drop_nulls_first_raises( assert_ir_translation_raises(q, in_memory_engine, NotImplementedError) +def test_groupby_sort_by_nested_agg_order_by_raises( + in_memory_engine: pl.GPUEngine, +) -> None: + df = pl.LazyFrame( + { + "g": ["A", "A", "B"], + "idx": [2, 1, 1], + "val": [30, 10, 20], + } + ) + + q = df.group_by("g").agg(pl.col("val").sort_by(pl.col("idx").max()).first()) + assert_ir_translation_raises(q, in_memory_engine, NotImplementedError) + + +def test_sorted_agg_validation() -> None: + dtype = DataType(pl.Int64()) + value = Col(dtype, "value") + by = Col(dtype, "by") + options = (False, (False,), (False,)) + + with pytest.raises(NotImplementedError, match="name='sum'"): + SortedAgg(dtype, "sum", options, value, by) + with pytest.raises(NotImplementedError, match="requires order-by expressions"): + SortedAgg(dtype, "first", (False, (), ()), value) + with pytest.raises(NotImplementedError, match="one null/descending option"): + SortedAgg(dtype, "first", (False, (False, False), (False,)), value, by) + + sorted_agg = SortedAgg(dtype, "first", options, value, by) + with pytest.raises(NotImplementedError, match="pylibcudf aggregation request"): + _ = sorted_agg.agg_request + + @pytest.mark.xfail(reason="https://github.com/pola-rs/polars/issues/17513") def test_groupby_minmax_with_nan(engine: pl.GPUEngine): df = pl.LazyFrame( From 6675be676dd3344978a2b2d8141cb81a311be79a Mon Sep 17 00:00:00 2001 From: rjzamora Date: Mon, 10 Aug 2026 15:07:50 -0700 Subject: [PATCH 10/13] handle need-preshuffle case --- .../cudf_polars/streaming/groupby.py | 21 +++++++++++++++- .../tests/streaming/test_groupby.py | 24 +++++++++++++++++++ 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/streaming/groupby.py b/python/cudf_polars/cudf_polars/streaming/groupby.py index 5a3f4c0fb865..8a4762147ce0 100644 --- a/python/cudf_polars/cudf_polars/streaming/groupby.py +++ b/python/cudf_polars/cudf_polars/streaming/groupby.py @@ -41,7 +41,7 @@ ) if TYPE_CHECKING: - from collections.abc import Generator, MutableMapping + from collections.abc import Generator, Iterable, MutableMapping from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import IR @@ -129,6 +129,15 @@ def combine( ) +def _has_stable_sorted_agg(exprs: Iterable[NamedExpr]) -> bool: + """Return True if any expression requires stable sorted aggregation.""" + return any( + isinstance(expr, SortedAgg) and expr.options[0] + for ne in exprs + for expr in traversal([ne.value]) + ) + + def _decompose_std_var( name: str, expr: Agg, *, names: Generator[str, None, None] ) -> tuple[NamedExpr, list[NamedExpr], list[NamedExpr], bool]: @@ -479,6 +488,16 @@ def _( # Preshuffle ir.child if needed if need_preshuffle: + if _has_stable_sorted_agg(ir.agg_requests): + return _lower_ir_fallback( + ir, + rec, + msg=( + "Streaming group_by does not yet preserve input-order ties " + "for sort_by(..., maintain_order=True) when another " + "aggregation requires repartitioning." + ), + ) child = Shuffle( child.schema, ir.keys, diff --git a/python/cudf_polars/tests/streaming/test_groupby.py b/python/cudf_polars/tests/streaming/test_groupby.py index b50f3351ad3b..20dc92944c99 100644 --- a/python/cudf_polars/tests/streaming/test_groupby.py +++ b/python/cudf_polars/tests/streaming/test_groupby.py @@ -212,6 +212,30 @@ def test_groupby_sort_by_first_last_stable_ties(streaming_engine_factory): assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) +def test_groupby_sort_by_stable_ties_with_preshuffle_fallback( + streaming_engine_factory, +): + streaming_engine = streaming_engine_factory( + StreamingOptions(max_rows_per_partition=2, fallback_mode="raise"), + ) + df = pl.LazyFrame( + { + "g": ["A", "B", "A", "B", "A", "B", "A", "B"], + "idx": [1, 1, 1, 1, 1, 1, 1, 1], + "val": [10, 100, 20, 200, 30, 300, 40, 400], + "u": [1, 1, 2, 2, 3, 3, 4, 4], + } + ) + + q = df.group_by("g").agg( + pl.col("u").n_unique().alias("nu"), + pl.col("val").sort_by("idx", maintain_order=True).first().alias("first_tie"), + pl.col("val").sort_by("idx", maintain_order=True).last().alias("last_tie"), + ) + with pytest.raises(NotImplementedError, match="input-order ties"): + q.collect(engine=streaming_engine) + + def test_groupby_sort_by_preserves_sorted_key_order(streaming_engine_factory): streaming_engine = streaming_engine_factory( StreamingOptions(max_rows_per_partition=2, fallback_mode="raise"), From 7b7545bc8681561f2c97ebb7e8f7dd05114b46ea Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 11 Aug 2026 06:29:30 -0700 Subject: [PATCH 11/13] reorder methods --- python/cudf_polars/cudf_polars/dsl/ir.py | 242 +++++++++--------- .../cudf_polars/dsl/utils/aggregations.py | 2 +- 2 files changed, 122 insertions(+), 122 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index ebe939b56f66..098027c01d5f 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -2156,6 +2156,127 @@ def __init__( self.zlice, ) + @classmethod + @log_do_evaluate + @nvtx_annotate_cudf_polars(message="GroupBy") + def do_evaluate( + cls, + schema: Schema, + keys_in: Sequence[expr.NamedExpr], + agg_requests: Sequence[expr.NamedExpr], + maintain_order: bool, # noqa: FBT001 + zlice: Zlice | None, + df: DataFrame, + *, + context: IRExecutionContext, + ) -> DataFrame: + """Evaluate and return a dataframe.""" + keys = broadcast( + *(k.evaluate(df) for k in keys_in), + target_length=df.num_rows, + stream=df.stream, + ) + keys_are_sorted = ( + plc.types.Sorted.YES + if all(k.is_sorted for k in keys) + else plc.types.Sorted.NO + ) + grouper = plc.groupby.GroupBy( + plc.Table([k.obj for k in keys]), + null_handling=plc.types.NullPolicy.INCLUDE, + keys_are_sorted=keys_are_sorted, + column_order=[k.order for k in keys], + null_precedence=[k.null_order for k in keys], + ) + requests: list[expr.NamedExpr] = [] + sorted_requests: list[expr.NamedExpr] = [] + for request in agg_requests: + value = request.value + if isinstance(value, expr.SortedAgg): + sorted_requests.append(request) + else: + requests.append(request) + group_keys, results = cls._evaluate_aggregation_requests( + schema, keys, grouper, requests, df + ) + sorted_group_keys, sorted_results = cls._evaluate_sorted_aggregations( + sorted_requests, keys, df + ) + if sorted_group_keys is not None: + if group_keys is None: + group_keys = sorted_group_keys + else: + sorted_results = [ + cls._align_to_group_keys( + group_keys, sorted_group_keys, result, df.stream + ) + for result in sorted_results + ] + if group_keys is None: + group_keys = cls._collect_group_keys(grouper, keys[0], df.stream) + results_by_name = { + request.name: result + for request, result in itertools.chain( + zip(requests, results, strict=True), + zip(sorted_requests, sorted_results, strict=True), + ) + } + results = [results_by_name[request.name] for request in agg_requests] + result_keys = [ + Column(grouped_key, name=key.name, dtype=key.dtype) + for key, grouped_key in zip(keys, group_keys.columns(), strict=True) + ] + if keys_are_sorted: + result_keys = [ + col.sorted_like(key) for col, key in zip(result_keys, keys, strict=True) + ] + broadcasted = broadcast(*result_keys, *results, stream=df.stream) + # Handle order preservation of groups + if maintain_order and not keys_are_sorted: + # The order we want + want = plc.stream_compaction.stable_distinct( + plc.Table([k.obj for k in keys]), + list(range(group_keys.num_columns())), + plc.stream_compaction.DuplicateKeepOption.KEEP_FIRST, + plc.types.NullEquality.EQUAL, + plc.types.NanEquality.ALL_EQUAL, + stream=df.stream, + ) + # The order we have + have = plc.Table([key.obj for key in broadcasted[: len(keys)]]) + + # We know an inner join is OK because by construction + # want and have are permutations of each other. + left_order, right_order = plc.join.inner_join( + want, have, plc.types.NullEquality.EQUAL, stream=df.stream + ) + # Now left_order is an arbitrary permutation of the ordering we + # want, and right_order is a matching permutation of the ordering + # we have. To get to the original ordering, we need + # left_order == iota(nrows), with right_order permuted + # appropriately. This can be obtained by sorting + # right_order by left_order. + (right_order,) = plc.sorting.sort_by_key( + plc.Table([right_order]), + plc.Table([left_order]), + [plc.types.Order.ASCENDING], + [plc.types.NullOrder.AFTER], + stream=df.stream, + ).columns() + ordered_table = plc.copying.gather( + plc.Table([col.obj for col in broadcasted]), + right_order, + plc.copying.OutOfBoundsPolicy.DONT_CHECK, + stream=df.stream, + ) + broadcasted = [ + Column(reordered, name=old.name, dtype=old.dtype) + for reordered, old in zip( + ordered_table.columns(), broadcasted, strict=True + ) + ] + return DataFrame(broadcasted, stream=df.stream).slice(zlice) + @staticmethod def _align_to_group_keys( group_keys: plc.Table, @@ -2398,127 +2519,6 @@ def _evaluate_sorted_aggregations( results_by_name[request.name] for request in sorted_requests ] - @classmethod - @log_do_evaluate - @nvtx_annotate_cudf_polars(message="GroupBy") - def do_evaluate( - cls, - schema: Schema, - keys_in: Sequence[expr.NamedExpr], - agg_requests: Sequence[expr.NamedExpr], - maintain_order: bool, # noqa: FBT001 - zlice: Zlice | None, - df: DataFrame, - *, - context: IRExecutionContext, - ) -> DataFrame: - """Evaluate and return a dataframe.""" - keys = broadcast( - *(k.evaluate(df) for k in keys_in), - target_length=df.num_rows, - stream=df.stream, - ) - keys_are_sorted = ( - plc.types.Sorted.YES - if all(k.is_sorted for k in keys) - else plc.types.Sorted.NO - ) - grouper = plc.groupby.GroupBy( - plc.Table([k.obj for k in keys]), - null_handling=plc.types.NullPolicy.INCLUDE, - keys_are_sorted=keys_are_sorted, - column_order=[k.order for k in keys], - null_precedence=[k.null_order for k in keys], - ) - requests: list[expr.NamedExpr] = [] - sorted_requests: list[expr.NamedExpr] = [] - for request in agg_requests: - value = request.value - if isinstance(value, expr.SortedAgg): - sorted_requests.append(request) - else: - requests.append(request) - group_keys, results = cls._evaluate_aggregation_requests( - schema, keys, grouper, requests, df - ) - sorted_group_keys, sorted_results = cls._evaluate_sorted_aggregations( - sorted_requests, keys, df - ) - if sorted_group_keys is not None: - if group_keys is None: - group_keys = sorted_group_keys - else: - sorted_results = [ - cls._align_to_group_keys( - group_keys, sorted_group_keys, result, df.stream - ) - for result in sorted_results - ] - if group_keys is None: - group_keys = cls._collect_group_keys(grouper, keys[0], df.stream) - results_by_name = { - request.name: result - for request, result in itertools.chain( - zip(requests, results, strict=True), - zip(sorted_requests, sorted_results, strict=True), - ) - } - results = [results_by_name[request.name] for request in agg_requests] - result_keys = [ - Column(grouped_key, name=key.name, dtype=key.dtype) - for key, grouped_key in zip(keys, group_keys.columns(), strict=True) - ] - if keys_are_sorted: - result_keys = [ - col.sorted_like(key) for col, key in zip(result_keys, keys, strict=True) - ] - broadcasted = broadcast(*result_keys, *results, stream=df.stream) - # Handle order preservation of groups - if maintain_order and not keys_are_sorted: - # The order we want - want = plc.stream_compaction.stable_distinct( - plc.Table([k.obj for k in keys]), - list(range(group_keys.num_columns())), - plc.stream_compaction.DuplicateKeepOption.KEEP_FIRST, - plc.types.NullEquality.EQUAL, - plc.types.NanEquality.ALL_EQUAL, - stream=df.stream, - ) - # The order we have - have = plc.Table([key.obj for key in broadcasted[: len(keys)]]) - - # We know an inner join is OK because by construction - # want and have are permutations of each other. - left_order, right_order = plc.join.inner_join( - want, have, plc.types.NullEquality.EQUAL, stream=df.stream - ) - # Now left_order is an arbitrary permutation of the ordering we - # want, and right_order is a matching permutation of the ordering - # we have. To get to the original ordering, we need - # left_order == iota(nrows), with right_order permuted - # appropriately. This can be obtained by sorting - # right_order by left_order. - (right_order,) = plc.sorting.sort_by_key( - plc.Table([right_order]), - plc.Table([left_order]), - [plc.types.Order.ASCENDING], - [plc.types.NullOrder.AFTER], - stream=df.stream, - ).columns() - ordered_table = plc.copying.gather( - plc.Table([col.obj for col in broadcasted]), - right_order, - plc.copying.OutOfBoundsPolicy.DONT_CHECK, - stream=df.stream, - ) - broadcasted = [ - Column(reordered, name=old.name, dtype=old.dtype) - for reordered, old in zip( - ordered_table.columns(), broadcasted, strict=True - ) - ] - return DataFrame(broadcasted, stream=df.stream).slice(zlice) - def _strip_predicate_casts(node: expr.Expr) -> expr.Expr: if isinstance(node, expr.Cast): diff --git a/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py b/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py index eafe12442a0b..3b68c05ab503 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/aggregations.py @@ -230,7 +230,7 @@ def decompose_single_agg( is_top=False, context=context, ) - if any(has_agg for _, has_agg in child_aggs): + if any(nested_agg for _, nested_agg in child_aggs): raise NotImplementedError( "Nested aggs in sorted groupby aggregation not supported" ) From 97b09578038cf84742ca23592732592799f98d77 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 11 Aug 2026 06:42:17 -0700 Subject: [PATCH 12/13] move group alignment into _evaluate_sorted_aggregations --- python/cudf_polars/cudf_polars/dsl/ir.py | 20 ++++++-------------- 1 file changed, 6 insertions(+), 14 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 098027c01d5f..6b5860f56831 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -2199,19 +2199,9 @@ def do_evaluate( group_keys, results = cls._evaluate_aggregation_requests( schema, keys, grouper, requests, df ) - sorted_group_keys, sorted_results = cls._evaluate_sorted_aggregations( - sorted_requests, keys, df + group_keys, sorted_results = cls._evaluate_sorted_aggregations( + sorted_requests, keys, df, target_group_keys=group_keys ) - if sorted_group_keys is not None: - if group_keys is None: - group_keys = sorted_group_keys - else: - sorted_results = [ - cls._align_to_group_keys( - group_keys, sorted_group_keys, result, df.stream - ) - for result in sorted_results - ] if group_keys is None: group_keys = cls._collect_group_keys(grouper, keys[0], df.stream) results_by_name = { @@ -2389,10 +2379,12 @@ def _evaluate_sorted_aggregations( sorted_requests: Sequence[expr.NamedExpr], keys: Sequence[Column], df: DataFrame, + *, + target_group_keys: plc.Table | None = None, ) -> tuple[plc.Table | None, list[Column]]: """Evaluate grouped first/last aggregations with explicit ordering.""" if not sorted_requests: - return None, [] + return target_group_keys, [] request_groups: dict[ tuple[Any, tuple[expr.Expr, ...]], list[expr.NamedExpr] @@ -2404,7 +2396,7 @@ def _evaluate_sorted_aggregations( (request.value.options, tuple(by_exprs)), [] ).append(request) - common_group_keys: plc.Table | None = None + common_group_keys = target_group_keys results_by_name: dict[str, Column] = {} key_order = [key.order for key in keys] key_null_order = [key.null_order for key in keys] From 31e3662af9113a66aa36f1664105ad7eaa2a3f67 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 11 Aug 2026 07:04:45 -0700 Subject: [PATCH 13/13] remove _collect_group_keys --- python/cudf_polars/cudf_polars/dsl/ir.py | 20 +------------------- 1 file changed, 1 insertion(+), 19 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 6b5860f56831..950ca5a1bc90 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -2203,7 +2203,7 @@ def do_evaluate( sorted_requests, keys, df, target_group_keys=group_keys ) if group_keys is None: - group_keys = cls._collect_group_keys(grouper, keys[0], df.stream) + group_keys, _ = grouper.aggregate([], stream=df.stream) results_by_name = { request.name: result for request, result in itertools.chain( @@ -2296,24 +2296,6 @@ def _align_to_group_keys( ).columns() return Column(ordered_result, name=result.name, dtype=result.dtype) - @staticmethod - def _collect_group_keys( - grouper: plc.groupby.GroupBy, - key: Column, - stream: Any, - ) -> plc.Table: - """Collect group keys when the GroupBy has no aggregation requests.""" - group_keys, _ = grouper.aggregate( - [ - plc.groupby.GroupByRequest( - key.obj, - [plc.aggregation.count(null_handling=plc.types.NullPolicy.INCLUDE)], - ) - ], - stream=stream, - ) - return group_keys - @staticmethod def _evaluate_aggregation_requests( schema: Schema,