NVIDIA / NVIDIA/cudf

[FEA] UDF aggregation on GroupBy

Open
#20,592 3 comments 0 reactions 0 assignees View on GitHub
feature request Python
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

**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`:

```python
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..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 .

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

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.