googleapis · chelsea-lin · May 16, 2024 · May 6, 2024 · May 7, 2024 · May 8, 2024
@@ -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}"
+        )