address comments offline

Author: chelsea-lin

bigframes/core/compile/aggregate_compiler.py

@@ -38,7 +38,9 @@ def compile_aggregate(
 ) -> ibis_types.Value:
     if isinstance(aggregate, ex.UnaryAggregation):
         input = scalar_compiler.compile_expression(aggregate.arg, bindings=bindings)
-        return compile_unary_agg(aggregate.op, input, order_by=order_by)
+        return compile_unary_agg(
+            aggregate.op, input, order_by=order_by if aggregate.op.can_order_by else []
+        )
     elif isinstance(aggregate, ex.BinaryAggregation):
         left = scalar_compiler.compile_expression(aggregate.left, bindings=bindings)
         right = scalar_compiler.compile_expression(aggregate.right, bindings=bindings)

bigframes/core/compile/compiled.py

@@ -305,15 +305,17 @@ def aggregate(
         aggregations: typing.Sequence[typing.Tuple[ex.Aggregation, str]],
         by_column_ids: typing.Sequence[str] = (),
         dropna: bool = True,
-    ) -> UnorderedIR:
+    ) -> OrderedIR:
         """
         Apply aggregations to the expression.
         Arguments:
             aggregations: input_column_id, operation, output_column_id tuples
-            by_column_id: column id of the aggregation key, this is preserved through the transform
+            by_column_id: column id of the aggregation key, this is preserved through
+              the transform
             dropna: whether null keys should be dropped
         Returns:
-            UnorderedIR
+            OrderedIR: the grouping key is a unique-valued column and has ordering
+              information.
         """
         table = self._to_ibis_expr()
         bindings = {col: table[col] for col in self.column_ids}
@@ -323,18 +325,32 @@ def aggregate(
         }
         if by_column_ids:
             result = table.group_by(by_column_ids).aggregate(**stats)
+            # Must have deterministic ordering, so order by the unique "by" column
+            ordering = ExpressionOrdering(
+                tuple([ascending_over(column_id) for column_id in by_column_ids]),
+                total_ordering_columns=frozenset(by_column_ids),
+            )
             columns = tuple(result[key] for key in result.columns)
-            expr = UnorderedIR(result, columns=columns)
+            expr = OrderedIR(result, columns=columns, ordering=ordering)
             if dropna:
                 for column_id in by_column_ids:
                     expr = expr._filter(expr._get_ibis_column(column_id).notnull())
             return expr
         else:
             aggregates = {**stats, ORDER_ID_COLUMN: ibis_types.literal(0)}
             result = table.aggregate(**aggregates)
+            # Ordering is irrelevant for single-row output, but set ordering id regardless
+            # as other ops(join etc.) expect it.
+            # TODO: Maybe can make completely empty
+            ordering = ExpressionOrdering(
+                ordering_value_columns=tuple([]),
+                total_ordering_columns=frozenset([]),
+            )
             return UnorderedIR(
                 result,
                 columns=[result[col_id] for col_id in [*stats.keys()]],
+                hidden_ordering_columns=[result[ORDER_ID_COLUMN]],
+                ordering=ordering,
             )
 
     def _uniform_sampling(self, fraction: float) -> UnorderedIR:
@@ -523,7 +539,8 @@ def from_pandas(
         """
         Builds an in-memory only (SQL only) expr from a pandas dataframe.
 
-        Assumed that the dataframe has unique string column names and bigframes-suppported dtypes.
+        Assumed that the dataframe has unique string column names and bigframes-suppported
+        dtypes.
         """
 
         # ibis memtable cannot handle NA, must convert to None
@@ -560,7 +577,8 @@ def _hidden_column_ids(self) -> typing.Sequence[str]:
 
     @property
     def _ibis_order(self) -> Sequence[ibis_types.Value]:
-        """Returns a sequence of ibis values which can be directly used to order a table expression. Has direction modifiers applied."""
+        """Returns a sequence of ibis values which can be directly used to order a
+        table expression. Has direction modifiers applied."""
         return _convert_ordering_to_table_values(
             {**self._column_names, **self._hidden_ordering_column_names},
             self._ordering.all_ordering_columns,
@@ -602,7 +620,8 @@ def aggregate(
         Apply aggregations to the expression.
         Arguments:
             aggregations: input_column_id, operation, output_column_id tuples
-            by_column_id: column id of the aggregation key, this is preserved through the transform
+            by_column_id: column id of the aggregation key, this is preserved through
+              the transform
             dropna: whether null keys should be dropped
         Returns:
             OrderedIR

bigframes/operations/aggregations.py

@@ -52,6 +52,10 @@ def arguments(self) -> int:
     def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
         return input_types[0]
 
+    @property
+    def can_order_by(self):
+        return False
+
 
 @dataclasses.dataclass(frozen=True)
 class AggregateOp(WindowOp):
@@ -226,6 +230,10 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
 class ArrayAggOp(UnaryAggregateOp):
     name: ClassVar[str] = "arrayagg"
 
+    @property
+    def can_order_by(self):
+        return True
+
     @property
     def skips_nulls(self):
         return True