Support skipna in groupby reductions (first/last, sum/prod/mean/median/min/max, idxmin/idxmax)

python/cudf/cudf/core/_internals/aggregation.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
@@ -152,15 +152,25 @@ def unique(cls) -> Self:
         )
 
     @classmethod
-    def first(cls) -> Self:
+    def first(cls, skipna: bool = True) -> Self:
         return cls(
-            plc.aggregation.nth_element(0, plc.types.NullPolicy.EXCLUDE)
+            plc.aggregation.nth_element(
+                0,
+                plc.types.NullPolicy.EXCLUDE
+                if skipna
+                else plc.types.NullPolicy.INCLUDE,
+            )
         )
 
     @classmethod
-    def last(cls) -> Self:
+    def last(cls, skipna: bool = True) -> Self:
         return cls(
-            plc.aggregation.nth_element(-1, plc.types.NullPolicy.EXCLUDE)
+            plc.aggregation.nth_element(
+                -1,
+                plc.types.NullPolicy.EXCLUDE
+                if skipna
+                else plc.types.NullPolicy.INCLUDE,
+            )
         )
 
     @classmethod

python/cudf/cudf/core/groupby/groupby.py

@@ -78,6 +78,20 @@
 # different dtypes. These strings must be elements of the AggregationKind enum.
 # The libcudf infrastructure exists for "COLLECT" support on
 # categoricals, but the dtype support in python does not.
+# Reductions whose result for a group becomes null when that group contains
+# any null value and ``skipna=False`` (libcudf otherwise always drops nulls).
+_NULL_PROPAGATING_REDUCTIONS = {
+    "sum",
+    "prod",
+    "product",
+    "mean",
+    "median",
+    "var",
+    "std",
+    "min",
+    "max",
+}
+
 _CATEGORICAL_AGGS = {"COUNT", "NUNIQUE", "SIZE", "UNIQUE"}
 _STRING_AGGS = {
     "COLLECT",
@@ -1318,7 +1332,9 @@ def agg(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs):
 
         return result
 
-    def _wrap_idxmin_idxmax(self, result: DataFrame | Series, *, skipna: bool):
+    def _wrap_idxmin_idxmax(
+        self, result: DataFrame | Series, *, skipna: bool, how: str
+    ):
         # libcudf's idxmin/idxmax return the integer row-position of the
         # min/max element within each group (null if the group's values were
         # all NA). pandas instead returns the *label* of that row taken from
@@ -1327,6 +1343,11 @@ def _wrap_idxmin_idxmax(self, result: DataFrame | Series, *, skipna: bool):
         from cudf.core.multiindex import MultiIndex
         from cudf.core.series import Series
 
+        if not skipna:
+            # pandas does not support positional idxmin/idxmax with
+            # skipna=False (it cannot represent "the label of a NA").
+            raise ValueError(f"{how} with skipna=False")
+
         key_names = set(self.grouping.names)
         if result.ndim == 2:
             value_items = [
@@ -1413,7 +1434,44 @@ def _reduce(
                 skipna=kwargs.get("skipna", True), min_count=min_count
             )
 
-        result = self.agg(op)
+        skipna = kwargs.get("skipna", True)
+        agg_op: str | _FirstLastAggSpec = op
+        if op in {"first", "last"} and not skipna:
+            # ``first``/``last`` default to dropping nulls (skipna=True). With
+            # ``skipna=False`` the actual first/last element of each group is
+            # returned even when it is null, matching pandas.
+            agg_op = _FirstLastAggSpec(op, skipna=False)
+
+        result = self.agg(agg_op)
+        if op in _NULL_PROPAGATING_REDUCTIONS and not skipna:
+            # libcudf reductions always drop nulls. With ``skipna=False`` a
+            # group containing any null in a column yields a null result for
+            # that (group, column), matching pandas. A (group, column) is
+            # all-non-null when its non-null count equals the group size
+            # (``size()`` is used instead of the ``size`` aggregation because
+            # the latter is unsupported for string columns).
+            from cudf.core.dataframe import DataFrame
+
+            non_null_counts = self.agg("count")
+            group_sizes = self.size()
+            if isinstance(group_sizes, DataFrame):
+                # With ``as_index=False`` the per-group counts are returned as
+                # the "size" column of a DataFrame; reduce it to a Series so it
+                # aligns with each value column below.
+                group_sizes = group_sizes["size"]
+            if isinstance(result, DataFrame):
+                # ``as_index=False`` keeps the grouping keys as columns of
+                # ``result``; they must never be nulled out, so mask only the
+                # value columns.
+                key_names = set(self.grouping.names)
+                for name in result._column_names:
+                    if name in key_names:
+                        continue
+                    result[name] = result[name].where(
+                        non_null_counts[name] == group_sizes, None
+                    )
+            else:
+                result = result.where(non_null_counts == group_sizes, None)
         if min_count and min_count > 0:
             counts = self.agg("count")
             result = result.where(counts >= min_count, None)
@@ -3523,7 +3581,7 @@ def idxmin(
         **kwargs: Any,
     ) -> DataFrame:
         result = self._reduce("idxmin", numeric_only=numeric_only)
-        return self._wrap_idxmin_idxmax(result, skipna=skipna)
+        return self._wrap_idxmin_idxmax(result, skipna=skipna, how="idxmin")
 
     def idxmax(
         self,
@@ -3533,7 +3591,7 @@ def idxmax(
         **kwargs: Any,
     ) -> DataFrame:
         result = self._reduce("idxmax", numeric_only=numeric_only)
-        return self._wrap_idxmin_idxmax(result, skipna=skipna)
+        return self._wrap_idxmin_idxmax(result, skipna=skipna, how="idxmax")
 
     def value_counts(
         self,
@@ -3848,13 +3906,13 @@ def idxmin(
         self, skipna: bool = True, min_count: int = 0, **kwargs: Any
     ) -> Series:
         result = self._reduce("idxmin")
-        return self._wrap_idxmin_idxmax(result, skipna=skipna)
+        return self._wrap_idxmin_idxmax(result, skipna=skipna, how="idxmin")
 
     def idxmax(
         self, skipna: bool = True, min_count: int = 0, **kwargs: Any
     ) -> Series:
         result = self._reduce("idxmax")
-        return self._wrap_idxmin_idxmax(result, skipna=skipna)
+        return self._wrap_idxmin_idxmax(result, skipna=skipna, how="idxmax")
 
     @property
     def dtype(self) -> pd.Series:
@@ -4166,6 +4224,28 @@ def copy(self, deep=True):
         return out
 
 
+class _FirstLastAggSpec:
+    """Callable aggregation spec for groupby ``first``/``last``.
+
+    Lets :meth:`GroupBy._reduce` thread ``skipna`` to
+    :meth:`Aggregation.first`/:meth:`Aggregation.last` through
+    ``make_aggregation``'s callable path. ``__str__``/``__name__`` report the
+    op name so aggregation-validity checks and result-column naming behave
+    exactly as they do for the plain ``"first"``/``"last"`` string specs.
+    """
+
+    def __init__(self, op: str, skipna: bool) -> None:
+        self._op = op
+        self._skipna = skipna
+        self.__name__ = op
+
+    def __call__(self, agg):
+        return getattr(agg, self._op)(skipna=self._skipna)
+
+    def __str__(self) -> str:
+        return self._op
+
+
 def _is_multi_agg(aggs):
     """
     Returns True if more than one aggregation is performed