fix order_by syntax

Author: chelsea-lin

bigframes/core/compile/aggregate_compiler.py

@@ -34,11 +34,11 @@
 def compile_aggregate(
     aggregate: ex.Aggregation,
     bindings: typing.Dict[str, ibis_types.Value],
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.Value:
     if isinstance(aggregate, ex.UnaryAggregation):
         input = scalar_compiler.compile_expression(aggregate.arg, bindings=bindings)
-        return compile_unary_agg(aggregate.op, input, agg_order_by=agg_order_by)
+        return compile_unary_agg(aggregate.op, input, order_by=order_by)
     elif isinstance(aggregate, ex.BinaryAggregation):
         left = scalar_compiler.compile_expression(aggregate.left, bindings=bindings)
         right = scalar_compiler.compile_expression(aggregate.right, bindings=bindings)
@@ -76,7 +76,7 @@ def compile_unary_agg(
     op: agg_ops.WindowOp,
     input: ibis_types.Column,
     window: Optional[window_spec.WindowSpec] = None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.Value:
     raise ValueError(f"Can't compile unrecognized operation: {op}")
 
@@ -87,7 +87,7 @@ def constrained_op(
         op,
         column: ibis_types.Column,
         window=None,
-        agg_order_by: typing.Sequence[ibis_types.Value] = [],
+        order_by: typing.Sequence[ibis_types.Value] = [],
     ):
         if column.type().is_boolean():
             column = typing.cast(
@@ -112,7 +112,7 @@ def _(
     op: agg_ops.SumOp,
     column: ibis_types.NumericColumn,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.NumericValue:
     # Will be null if all inputs are null. Pandas defaults to zero sum though.
     bq_sum = _apply_window_if_present(column.sum(), window)
@@ -127,7 +127,7 @@ def _(
     op: agg_ops.MedianOp,
     column: ibis_types.NumericColumn,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.NumericValue:
     # PERCENTILE_CONT has very few allowed windows. For example, "window
     # framing clause is not allowed for analytic function percentile_cont".
@@ -148,7 +148,7 @@ def _(
     op: agg_ops.ApproxQuartilesOp,
     column: ibis_types.NumericColumn,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.NumericValue:
     # PERCENTILE_CONT has very few allowed windows. For example, "window
     # framing clause is not allowed for analytic function percentile_cont".
@@ -168,7 +168,7 @@ def _(
     op: agg_ops.QuantileOp,
     column: ibis_types.NumericColumn,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.NumericValue:
     return _apply_window_if_present(column.quantile(op.q), window)
 
@@ -179,7 +179,7 @@ def _(
     op: agg_ops.MeanOp,
     column: ibis_types.NumericColumn,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.NumericValue:
     return _apply_window_if_present(column.mean(), window)
 
@@ -190,7 +190,7 @@ def _(
     op: agg_ops.ProductOp,
     column: ibis_types.NumericColumn,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.NumericValue:
     # Need to short-circuit as log with zeroes is illegal sql
     is_zero = cast(ibis_types.BooleanColumn, (column == 0))
@@ -229,7 +229,7 @@ def _(
     op: agg_ops.MaxOp,
     column: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.Value:
     return _apply_window_if_present(column.max(), window)
 
@@ -239,7 +239,7 @@ def _(
     op: agg_ops.MinOp,
     column: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.Value:
     return _apply_window_if_present(column.min(), window)
 
@@ -250,7 +250,7 @@ def _(
     op: agg_ops.StdOp,
     x: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.Value:
     return _apply_window_if_present(cast(ibis_types.NumericColumn, x).std(), window)
 
@@ -261,7 +261,7 @@ def _(
     op: agg_ops.VarOp,
     x: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.Value:
     return _apply_window_if_present(cast(ibis_types.NumericColumn, x).var(), window)
 
@@ -272,7 +272,7 @@ def _(
     op: agg_ops.PopVarOp,
     x: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.Value:
     return _apply_window_if_present(
         cast(ibis_types.NumericColumn, x).var(how="pop"), window
@@ -284,7 +284,7 @@ def _(
     op: agg_ops.CountOp,
     column: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.IntegerValue:
     return _apply_window_if_present(column.count(), window)
 
@@ -294,7 +294,7 @@ def _(
     op: agg_ops.ArrayAggOp,
     column: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.ArrayValue:
     # BigQuery doesn't currently support using ARRAY_AGG with both window and aggregate
     # functions simultaneously. Some aggregate functions (or its equivalent syntax)
@@ -310,7 +310,7 @@ def _(
 
     return vendored_ibis_ops.ArrayAggregate(
         column,
-        order_by_columns=agg_order_by,
+        order_by=order_by,
     ).to_expr()
 
 
@@ -319,7 +319,7 @@ def _(
     op: agg_ops.CutOp,
     x: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ):
     out = ibis.case()
     if isinstance(op.bins, int):
@@ -376,7 +376,7 @@ def _(
     self: agg_ops.QcutOp,
     column: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.IntegerValue:
     if isinstance(self.quantiles, int):
         quantiles_ibis = dtypes.literal_to_ibis_scalar(self.quantiles)
@@ -409,7 +409,7 @@ def _(
     op: agg_ops.NuniqueOp,
     column: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.IntegerValue:
     return _apply_window_if_present(column.nunique(), window)
 
@@ -419,7 +419,7 @@ def _(
     op: agg_ops.AnyValueOp,
     column: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.IntegerValue:
     return _apply_window_if_present(column.arbitrary(), window)
 
@@ -429,7 +429,7 @@ def _(
     op: agg_ops.RankOp,
     column: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.IntegerValue:
     # Ibis produces 0-based ranks, while pandas creates 1-based ranks
     return _apply_window_if_present(ibis.rank(), window) + 1
@@ -440,7 +440,7 @@ def _(
     op: agg_ops.DenseRankOp,
     column: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.IntegerValue:
     # Ibis produces 0-based ranks, while pandas creates 1-based ranks
     return _apply_window_if_present(column.dense_rank(), window) + 1
@@ -456,7 +456,7 @@ def _(
     op: agg_ops.FirstNonNullOp,
     column: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.Value:
     return _apply_window_if_present(
         vendored_ibis_ops.FirstNonNullValue(column).to_expr(), window  # type: ignore
@@ -468,7 +468,7 @@ def _(
     op: agg_ops.LastOp,
     column: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.Value:
     return _apply_window_if_present(column.last(), window)
 
@@ -478,7 +478,7 @@ def _(
     op: agg_ops.LastNonNullOp,
     column: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.Value:
     return _apply_window_if_present(
         vendored_ibis_ops.LastNonNullValue(column).to_expr(), window  # type: ignore
@@ -490,7 +490,7 @@ def _(
     op: agg_ops.ShiftOp,
     column: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.Value:
     if op.periods == 0:  # No-op
         return column
@@ -504,7 +504,7 @@ def _(
     op: agg_ops.DiffOp,
     column: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.Value:
     shifted = compile_unary_agg(agg_ops.ShiftOp(op.periods), column, window)
     if column.type().is_boolean():
@@ -524,7 +524,7 @@ def _(
     op: agg_ops.AllOp,
     column: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.BooleanValue:
     # BQ will return null for empty column, result would be true in pandas.
     result = _is_true(column).all()
@@ -539,7 +539,7 @@ def _(
     op: agg_ops.AnyOp,
     column: ibis_types.Column,
     window=None,
-    agg_order_by: typing.Sequence[ibis_types.Value] = [],
+    order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.BooleanValue:
     # BQ will return null for empty column, result would be false in pandas.
     result = _is_true(column).any()

bigframes/core/compile/compiled.py

@@ -608,10 +608,9 @@ def aggregate(
         """
         table = self._to_ibis_expr(ordering_mode="unordered", expose_hidden_cols=True)
         bindings = {col: table[col] for col in self.column_ids}
-        agg_order_by = [table[col] for col in self._hidden_ordering_column_names]
         stats = {
             col_out: agg_compiler.compile_aggregate(
-                aggregate, bindings, agg_order_by=agg_order_by
+                aggregate, bindings, order_by=self._ibis_order
             )
             for aggregate, col_out in aggregations
         }

tests/system/small/bigquery/test_array.py

@@ -83,6 +83,40 @@ def test_array_agg_w_dataframe():
         expected.to_pandas(),
     )
 
+
+@pytest.mark.parametrize(
+    ("ascending", "expected_b", "expected_c"),
+    [
+        pytest.param(
+            True, [["a", "b"], ["e", "d", "c"]], [[4, 5], [1, 2, 3]], id="asc"
+        ),
+        pytest.param(
+            False, [["b", "a"], ["c", "d", "e"]], [[5, 4], [3, 2, 1]], id="des"
+        ),
+    ],
+)
+def test_array_agg_reserve_order(ascending, expected_b, expected_c):
+    data = {
+        "a": [1, 1, 2, 2, 2],
+        "b": ["a", "b", "c", "d", "e"],
+        "c": [4, 5, 3, 2, 1],
+    }
+    df = bpd.DataFrame(data)
+
+    result = bbq.array_agg(df.sort_values("c", ascending=ascending).groupby(by=["a"]))
+    expected_data = {
+        "a": [1, 2],
+        "b": expected_b,
+        "c": expected_c,
+    }
+    expected = bpd.DataFrame(expected_data).set_index("a")
+
+    pd.testing.assert_frame_equal(
+        result.to_pandas(),
+        expected.to_pandas(),
+    )
+
+
 def assert_array_agg_matches_after_explode():
     data = {
         "index": np.arange(10),

third_party/bigframes_vendored/ibis/backends/bigquery/registry.py

@@ -42,11 +42,9 @@ def _array_aggregate(translator, op: vendored_ibis_ops.ArrayAggregate):
     arg = translator.translate(op.arg)
 
     order_by_sql = ""
-    if len(op.order_by_columns) > 0:
-        order_by_columns = ", ".join(
-            [translator.translate(column) for column in op.order_by_columns]
-        )
-        order_by_sql = f"ORDER BY {order_by_columns}"
+    if len(op.order_by) > 0:
+        order_by = ", ".join([translator.translate(column) for column in op.order_by])
+        order_by_sql = f"ORDER BY {order_by}"
 
     return f"ARRAY_AGG({arg} IGNORE NULLS {order_by_sql})"
 

third_party/bigframes_vendored/ibis/expr/operations/reductions.py

@@ -27,7 +27,7 @@ class ArrayAggregate(Filterable, Reduction):
     """
 
     arg: ibis_ops_core.Column
-    order_by_columns: VarTuple[ibis_ops_core.Value] = ()
+    order_by: VarTuple[ibis_ops_core.Value] = ()
 
     @ibis_annotations.attribute
     def dtype(self):