Repository: ibis-project/ibis

Issue:
BUG: Reduction UDF returning `np.array` hits error when used in grouped aggregation
```
import numpy as np
import pandas as pd

import ibis
import ibis.expr.datatypes as dt
import ibis.expr.operations as ops
import ibis.expr.types as ir
import ibis.udf.vectorized as udf

### Setup

df = pd.DataFrame({
    'key_col': pd.Series([1, 1, 2, 2,]),
    'int_col': pd.Series([100, 200, 300,]),
})

ibis_pandas = ibis.pandas.connect({'test': df})
t = ibis_pandas.table('test')


### The interesting part

@udf.reduction(
    input_type=[dt.int64],
    output_type=dt.Array(dt.int64)
)
def collect_udf(col):
    return np.array(col)

expr = t.groupby('key_col').aggregate(
    collected=collect_udf(t.int_col)
)

result = ibis_pandas.execute(expr)
```

```
...
~/ibis/ibis/backends/pandas/aggcontext.py in agg(self, grouped_data, function, *args, **kwargs)
    367             )
    368 
--> 369         return grouped_data.agg(wrap_for_agg(function, args, kwargs))
    370 
    371 

~/miniconda3/envs/ibis-dev/lib/python3.9/site-packages/pandas/core/groupby/generic.py in aggregate(self, func, engine, engine_kwargs, *args, **kwargs)
    261                 # TODO: KeyError is raised in _python_agg_general,
    262                 #  see test_groupby.test_basic
--> 263                 result = self._aggregate_named(func, *args, **kwargs)
    264 
    265             index = Index(sorted(result), name=self.grouper.names[0])

~/miniconda3/envs/ibis-dev/lib/python3.9/site-packages/pandas/core/groupby/generic.py in _aggregate_named(self, func, *args, **kwargs)
    482             if not initialized:
    483                 # We only do this validation on the first iteration
--> 484                 libreduction.check_result_array(output, 0)
    485                 initialized = True
    486             result[name] = output

pandas/_libs/reduction.pyx in pandas._libs.reduction.check_result_array()

pandas/_libs/reduction.pyx in pandas._libs.reduction.check_result_array()

ValueError: Must produce aggregated value
```


Discussion:
Looks like this is due to some Pandas internals: https://stackoverflow.com/a/16975602

The solution may just be to work around it in some way

Resolving patch:
diff --git a/ibis/backends/pandas/aggcontext.py b/ibis/backends/pandas/aggcontext.py
--- a/ibis/backends/pandas/aggcontext.py
+++ b/ibis/backends/pandas/aggcontext.py
@@ -363,7 +363,21 @@ def agg(self, grouped_data, function, *args, **kwargs):
                 'Object {} is not callable or a string'.format(function)
             )
 
-        return grouped_data.agg(wrap_for_agg(function, args, kwargs))
+        if isinstance(grouped_data, pd.core.groupby.generic.SeriesGroupBy):
+            # `SeriesGroupBy.agg` does not allow np.arrays to be returned
+            # from UDFs. To avoid `SeriesGroupBy.agg`, we will us
+            # `Series.agg` manually on each group. (#2768)
+            aggs = {}
+            for k, v in grouped_data:
+                aggs[k] = v.agg(wrap_for_agg(function, args, kwargs))
+                grouped_col_name = v.name
+            return (
+                pd.Series(aggs)
+                .rename(grouped_col_name)
+                .rename_axis(grouped_data.grouper.names)
+            )
+        else:
+            return grouped_data.agg(wrap_for_agg(function, args, kwargs))
 
 
 class Transform(AggregationContext):


Test patch:
diff --git a/ibis/backends/dask/tests/test_udf.py b/ibis/backends/dask/tests/test_udf.py
--- a/ibis/backends/dask/tests/test_udf.py
+++ b/ibis/backends/dask/tests/test_udf.py
@@ -449,6 +449,7 @@ def qs(request):
 
 
 def test_array_return_type_reduction(con, t, df, qs):
+    """Tests reduction UDF returning an array."""
     expr = quantiles(t.b, quantiles=qs)
     result = expr.execute()
     expected = df.b.quantile(qs).compute()
@@ -459,6 +460,7 @@ def test_array_return_type_reduction(con, t, df, qs):
     raises=NotImplementedError, reason='TODO - windowing - #2553'
 )
 def test_array_return_type_reduction_window(con, t, df, qs):
+    """Tests reduction UDF returning an array, used over a window."""
     expr = quantiles(t.b, quantiles=qs).over(ibis.window())
     result = expr.execute()
     expected_raw = df.b.quantile(qs).tolist()
@@ -466,6 +468,20 @@ def test_array_return_type_reduction_window(con, t, df, qs):
     tm.assert_series_equal(result, expected)
 
 
+def test_array_return_type_reduction_group_by(con, t, df, qs):
+    """Tests reduction UDF returning an array, used in a grouped agg."""
+    expr = t.groupby(t.key).aggregate(
+        quantiles_col=quantiles(t.b, quantiles=qs)
+    )
+    result = expr.execute()
+
+    df = df.compute()  # Convert to Pandas
+    expected_col = df.groupby(df.key).b.agg(lambda s: s.quantile(qs).tolist())
+    expected = pd.DataFrame({'quantiles_col': expected_col}).reset_index()
+
+    tm.assert_frame_equal(result, expected)
+
+
 def test_elementwise_udf_with_many_args(t2):
     @elementwise(
         input_type=[dt.double] * 16 + [dt.int32] * 8, output_type=dt.double
diff --git a/ibis/backends/pandas/tests/test_udf.py b/ibis/backends/pandas/tests/test_udf.py
--- a/ibis/backends/pandas/tests/test_udf.py
+++ b/ibis/backends/pandas/tests/test_udf.py
@@ -2,7 +2,7 @@
 
 import numpy as np
 import pandas as pd
-import pandas.testing as tm
+import pandas._testing as tm
 import pytest
 
 import ibis
@@ -98,7 +98,7 @@ def zscore(series):
     input_type=[dt.double], output_type=dt.Array(dt.double),
 )
 def quantiles(series, *, quantiles):
-    return list(series.quantile(quantiles))
+    return np.array(series.quantile(quantiles))
 
 
 def test_udf(t, df):
@@ -394,13 +394,15 @@ def qs(request):
 
 
 def test_array_return_type_reduction(con, t, df, qs):
+    """Tests reduction UDF returning an array."""
     expr = quantiles(t.b, quantiles=qs)
     result = expr.execute()
-    expected = df.b.quantile(qs)
-    assert result == expected.tolist()
+    expected = np.array(df.b.quantile(qs))
+    tm.assert_numpy_array_equal(result, expected)
 
 
 def test_array_return_type_reduction_window(con, t, df, qs):
+    """Tests reduction UDF returning an array, used over a window."""
     expr = quantiles(t.b, quantiles=qs).over(ibis.window())
     result = expr.execute()
     expected_raw = df.b.quantile(qs).tolist()
@@ -408,6 +410,23 @@ def test_array_return_type_reduction_window(con, t, df, qs):
     tm.assert_series_equal(result, expected)
 
 
+def test_array_return_type_reduction_group_by(con, t, df, qs):
+    """Tests reduction UDF returning an array, used in a grouped aggregation.
+
+    Getting this use case to succeed required avoiding use of
+    `SeriesGroupBy.agg` in the `Summarize` aggcontext implementation (#2768).
+    """
+    expr = t.groupby(t.key).aggregate(
+        quantiles_col=quantiles(t.b, quantiles=qs)
+    )
+    result = expr.execute()
+
+    expected_col = df.groupby(df.key).b.agg(lambda s: s.quantile(qs).tolist())
+    expected = pd.DataFrame({'quantiles_col': expected_col}).reset_index()
+
+    tm.assert_frame_equal(result, expected)
+
+
 def test_elementwise_udf_with_many_args(t2):
     @udf.elementwise(
         input_type=[dt.double] * 16 + [dt.int32] * 8, output_type=dt.double