feat: bigframes.bigquery.array_agg(SeriesGroupBy|DataFrameGroupby) (#663)

Author: chelsea-lin

bigframes/bigquery/__init__.py

@@ -22,9 +22,13 @@
 
 import typing
 
+import bigframes.constants as constants
+import bigframes.core.groupby as groupby
 import bigframes.operations as ops
+import bigframes.operations.aggregations as agg_ops
 
 if typing.TYPE_CHECKING:
+    import bigframes.dataframe as dataframe
     import bigframes.series as series
 
 
@@ -52,9 +56,66 @@ def array_length(series: series.Series) -> series.Series:
         2    2
         dtype: Int64
 
+    Args:
+        series (bigframes.series.Series):
+                A Series with array columns.
+
     Returns:
         bigframes.series.Series: A Series of integer values indicating
             the length of each element in the Series.
 
     """
     return series._apply_unary_op(ops.len_op)
+
+
+def array_agg(
+    obj: groupby.SeriesGroupBy | groupby.DataFrameGroupBy,
+) -> series.Series | dataframe.DataFrame:
+    """Group data and create arrays from selected columns, omitting NULLs to avoid
+    BigQuery errors (NULLs not allowed in arrays).
+
+    **Examples:**
+
+        >>> import bigframes.pandas as bpd
+        >>> import bigframes.bigquery as bbq
+        >>> import numpy as np
+        >>> bpd.options.display.progress_bar = None
+
+    For a SeriesGroupBy object:
+
+        >>> lst = ['a', 'a', 'b', 'b', 'a']
+        >>> s = bpd.Series([1, 2, 3, 4, np.nan], index=lst)
+        >>> bbq.array_agg(s.groupby(level=0))
+        a    [1. 2.]
+        b    [3. 4.]
+        dtype: list<item: double>[pyarrow]
+
+    For a DataFrameGroupBy object:
+
+        >>> l = [[1, 2, 3], [1, None, 4], [2, 1, 3], [1, 2, 2]]
+        >>> df = bpd.DataFrame(l, columns=["a", "b", "c"])
+        >>> bbq.array_agg(df.groupby(by=["b"]))
+                 a      c
+        b
+        1.0    [2]    [3]
+        2.0  [1 1]  [3 2]
+        <BLANKLINE>
+        [2 rows x 2 columns]
+
+    Args:
+        obj (groupby.SeriesGroupBy | groupby.DataFrameGroupBy):
+                A GroupBy object to be applied the function.
+
+    Returns:
+        bigframes.series.Series | bigframes.dataframe.DataFrame: A Series or
+            DataFrame containing aggregated array columns, and indexed by the
+            original group columns.
+    """
+    if isinstance(obj, groupby.SeriesGroupBy):
+        return obj._aggregate(agg_ops.ArrayAggOp())
+    elif isinstance(obj, groupby.DataFrameGroupBy):
+        return obj._aggregate_all(agg_ops.ArrayAggOp(), numeric_only=False)
+    else:
+        raise ValueError(
+            f"Unsupported type {type(obj)} to apply `array_agg` function. {constants.FEEDBACK_LINK}"
+        )