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..ef94763bb668 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,54 @@ def do_evaluate( return value +class SortedAgg(Expr): + """ + ``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") + + 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 cannot be represented as a pylibcudf " + "aggregation request" + ) + + 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..950ca5a1bc90 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -50,12 +50,13 @@ 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, 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, @@ -2187,52 +2188,30 @@ def do_evaluate( column_order=[k.order for k in keys], null_precedence=[k.null_order for k in keys], ) - requests = [] - names = [] - cast_to_schema = [] + requests: list[expr.NamedExpr] = [] + sorted_requests: list[expr.NamedExpr] = [] 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() - ) + if isinstance(value, expr.SortedAgg): + sorted_requests.append(request) 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, 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 - ) - for name, column, should_cast in zip( - names, - itertools.chain.from_iterable(t.columns() for t in raw_tables), - cast_to_schema, - strict=True, + requests.append(request) + group_keys, results = cls._evaluate_aggregation_requests( + schema, keys, grouper, requests, df + ) + group_keys, sorted_results = cls._evaluate_sorted_aggregations( + sorted_requests, keys, df, target_group_keys=group_keys + ) + if group_keys is None: + group_keys, _ = grouper.aggregate([], stream=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) @@ -2288,6 +2267,232 @@ def do_evaluate( ] return DataFrame(broadcasted, stream=df.stream).slice(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_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_aggregations( + cls, + 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 target_group_keys, [] + + 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 = 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] + + 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, 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 + ) + 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_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 + ] + 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 4b66ebb1846a..3b68c05ab503 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(nested_agg for _, nested_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/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 2d504ee112af..8a4762147ce0 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, @@ -40,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 @@ -128,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]: @@ -251,6 +261,46 @@ 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.""" + 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 +339,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": @@ -436,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 9bfb363147ce..20dc92944c99 100644 --- a/python/cudf_polars/tests/streaming/test_groupby.py +++ b/python/cudf_polars/tests/streaming/test_groupby.py @@ -173,6 +173,92 @@ 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) + + +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_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"), + ) + 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 326d39421265..5924aaabf6be 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,64 @@ 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], + "tie": [1, 1, 1, 1, 1, 1, 1, 1], + "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"), + 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() + 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..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 ( @@ -424,6 +424,141 @@ 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_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: + 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) + + +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(