Merge remote-tracking branch 'origin/main' into refactor-isnull-op

Author: tswast

bigframes/session/polars_executor.py

@@ -47,7 +47,15 @@
     bigframes.operations.ge_op,
     bigframes.operations.le_op,
 )
-_COMPATIBLE_AGG_OPS = (agg_ops.SizeOp, agg_ops.SizeUnaryOp)
+_COMPATIBLE_AGG_OPS = (
+    agg_ops.SizeOp,
+    agg_ops.SizeUnaryOp,
+    agg_ops.MinOp,
+    agg_ops.MaxOp,
+    agg_ops.SumOp,
+    agg_ops.MeanOp,
+    agg_ops.CountOp,
+)
 
 
 def _get_expr_ops(expr: expression.Expression) -> set[bigframes.operations.ScalarOp]:

bigframes/testing/engine_utils.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pandas.testing
+
 from bigframes.core import nodes
 from bigframes.session import semi_executor
 
@@ -25,7 +27,8 @@ def assert_equivalence_execution(
     e2_result = engine2.execute(node, ordered=True)
     assert e1_result is not None
     assert e2_result is not None
-    # Schemas might have extra nullity markers, normalize to node expected schema, which should be looser
-    e1_table = e1_result.to_arrow_table().cast(node.schema.to_pyarrow())
-    e2_table = e2_result.to_arrow_table().cast(node.schema.to_pyarrow())
-    assert e1_table.equals(e2_table), f"{e1_table} is not equal to {e2_table}"
+    # Convert to pandas, as pandas has better comparison utils than arrow
+    assert e1_result.schema == e2_result.schema
+    e1_table = e1_result.to_pandas()
+    e2_table = e2_result.to_pandas()
+    pandas.testing.assert_frame_equal(e1_table, e2_table, rtol=1e-10)

tests/system/small/engines/test_aggregation.py

@@ -25,6 +25,28 @@
 REFERENCE_ENGINE = polars_executor.PolarsExecutor()
 
 
+def apply_agg_to_all_valid(
+    array: array_value.ArrayValue, op: agg_ops.UnaryAggregateOp, excluded_cols=[]
+) -> array_value.ArrayValue:
+    """
+    Apply the aggregation to every column in the array that has a compatible datatype.
+    """
+    exprs_by_name = []
+    for arg in array.column_ids:
+        if arg in excluded_cols:
+            continue
+        try:
+            _ = op.output_type(array.get_column_type(arg))
+            expr = expression.UnaryAggregation(op, expression.deref(arg))
+            name = f"{arg}-{op.name}"
+            exprs_by_name.append((expr, name))
+        except TypeError:
+            continue
+    assert len(exprs_by_name) > 0
+    new_arr = array.aggregate(exprs_by_name)
+    return new_arr
+
+
 @pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
 def test_engines_aggregate_size(
     scalars_array_value: array_value.ArrayValue,
@@ -48,6 +70,20 @@ def test_engines_aggregate_size(
     assert_equivalence_execution(node, REFERENCE_ENGINE, engine)
 
 
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+@pytest.mark.parametrize(
+    "op",
+    [agg_ops.min_op, agg_ops.max_op, agg_ops.mean_op, agg_ops.sum_op, agg_ops.count_op],
+)
+def test_engines_unary_aggregates(
+    scalars_array_value: array_value.ArrayValue,
+    engine,
+    op,
+):
+    node = apply_agg_to_all_valid(scalars_array_value, op).node
+    assert_equivalence_execution(node, REFERENCE_ENGINE, engine)
+
+
 @pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
 @pytest.mark.parametrize(
     "grouping_cols",