[FEA] UDF aggregation on GroupBy

[din14970]

Is your feature request related to a problem? Please describe.
Not sure if this is a feature request or a bug. Sorry if this a duplicate, browsed through the issues but couldn't seem to find a related issues. It seems UDFs are not supported with agg on GroupBy:

df.groupby("groups")["values"].agg(my_udf)

Perhaps naively I would say this should behave in the same way as a rolling window UDF, where a callable that takes in a group and returns a single value. However, if I pass a custom function, instead I get:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[61], line 11
---> 11 grp["value_1"].agg(geometric_mean)

File [cudf/core/groupby/groupby.py:3429], in SeriesGroupBy.agg(self, func, engine, engine_kwargs, *args, **kwargs)
   3428 def agg(self, func, *args, engine=None, engine_kwargs=None, **kwargs):
-> 3429     result = super().agg(
   3430         func, *args, engine=engine, engine_kwargs=engine_kwargs, **kwargs
   3431     )
   3433     # downcast the result to a Series:
   3434     if result._num_columns:

File [cudf/utils/performance_tracking.py:51], in _performance_tracking.<locals>.wrapper(*args, **kwargs)
     43 if nvtx.enabled():
     44     stack.enter_context(
     45         nvtx.annotate(
     46             message=func.__qualname__,
   (...)     49         )
     50     )
---> 51 return func(*args, **kwargs)

File [cudf/core/groupby/groupby.py:1012], in GroupBy.agg(self, func, engine, engine_kwargs, *args, **kwargs)
   1002 orig_dtypes = tuple(c.dtype for c in columns)
   1004 # Note: When there are no key columns, the below produces
   1005 # an Index with float64 dtype, while Pandas returns
   1006 # an Index with int64 dtype.
   1007 # (GH: 6945)
   1008 (
   1009     result_columns,
   1010     grouped_key_cols,
   1011     included_aggregations,
-> 1012 ) = self._aggregate(columns, normalized_aggs)
   1014 result_index = self.grouping.keys._from_columns_like_self(
   1015     grouped_key_cols,
   1016 )
   1018 multilevel = _is_multi_agg(func)

File [cudf/core/groupby/groupby.py:830], in GroupBy._aggregate(self, values, aggregations)
    826 if _is_unsupported_agg_for_type(col.dtype, str_agg):
    827     raise TypeError(
    828         f"{col.dtype} type does not support {agg} operations"
    829     )
--> 830 agg_obj = aggregation.make_aggregation(agg)
    831 if (
    832     valid_aggregations == "ALL"
    833     or agg_obj.kind in valid_aggregations
    834 ):
    835     included_aggregations_i.append((agg, agg_obj.kind))

File [cudf/core/_internals/aggregation.py:287], in make_aggregation(op, kwargs)
    285         return Aggregation.from_udf(op, **kwargs)
    286     else:
--> 287         return op(Aggregation)
    288 raise TypeError(f"Unknown aggregation {op}")

Cell In[61], line 4, in geometric_mean(values)
      2 n = values.size
      3 product = 1
----> 4 for value in values:
      5     product *= value
      6 return (product) ** (1/n)

TypeError: 'type' object is not iterable

Clearly the last bit of my error is specific to my UDF, but I would expect an error more in line with "we don't support this at the moment". This is confusing because the agg docstring clearly states it accepts "callables", but it is not entirely clear what is an acceptable callable.

I had had a quick look at aggregation.make_aggregation, which also seems to be used for the rolling window UDFs, and it seems this supports UDFs when an output data type is supplied, which is not done in

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

Line 830 in f4e35ca

agg_obj = aggregation.make_aggregation(agg)

.

Basically, I'm wondering whether this is an oversight or there is a good reason this doesn't work.

Describe the solution you'd like

For UDFs to work in GroupBy.transform and GroupBy.agg in the same way as Rolling.apply. Alteratively, if not possible for good reasons, a more constructive error message should be raised that UDFs are not supported, or the docs should clarify what kind of callables are supported.

Describe alternatives you've considered

I'm exploring the groupby functionality specifically for something I'm writing up about cudf and bumped into this issue. Of course one can always write a custom CUDA kernel.

Additional context
I'm on stable version 25.10

[brandon-b-miller]

Hi @din14970 ,

There's some limited support for UDFs through groupby.apply, see caveats at https://docs.rapids.ai/api/cudf/stable/user_guide/guide-to-udfs/. Right now we're missing a few bits and pieces needed to compute a geometric mean but it might be doable with a few small PRs. Can you share a bit more detail about your use case? Would support for float/int 32/64 data be sufficient?

[din14970]

Hey @brandon-b-miller ,

As mentioned, this is mainly a dummy example I came up with to test a UDF with agg for some teaching material I'm developing on cudf. For the geometric mean specifically a simple workaround is:

aggregated = grp["value_1"].agg(["prod", "count"])
aggregated["prod"] ** (1/aggregated["count"])

So a custom aggregation method for it is not very urgent.

I'll have a look at apply instead to see what's possible, but I guess looping over elements is out of the question here? It's confusing what is and isn't allowed in these functions.

[din14970]

It's confusing what is and isn't allowed in these functions.

I take that back, it's clearly documented here: https://docs.rapids.ai/api/cudf/stable/user_guide/api_docs/api/cudf.core.groupby.dataframegroupby.apply/#cudf.core.groupby.DataFrameGroupBy.apply.

I do think this part of the docs should be modified:

callable: a function that accepts a Series/DataFrame and performs a supported operation on it.

Clearly that is not the case, the feature set is much more restricted.