[WIP] Support ordered first/last groupby aggregations in cudf-polars

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",

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")