feat: Polars engine supports std, var

Author: TrevorBergeron

bigframes/core/compile/polars/compiler.py

@@ -535,9 +535,11 @@ def compile_agg_op(
             if isinstance(op, agg_ops.StdOp):
                 return pl.std(inputs[0])
             if isinstance(op, agg_ops.VarOp):
-                return pl.var(inputs[0])
+                # polars var doesnt' support decimal, so use std instead
+                return pl.std(inputs[0]).pow(2)
             if isinstance(op, agg_ops.PopVarOp):
-                return pl.var(inputs[0], ddof=0)
+                # polars var doesnt' support decimal, so use std instead
+                return pl.std(inputs[0], ddof=0).pow(2)
             if isinstance(op, agg_ops.FirstNonNullOp):
                 return pl.col(*inputs).drop_nulls().first()
             if isinstance(op, agg_ops.LastNonNullOp):

bigframes/session/polars_executor.py

@@ -106,6 +106,9 @@
     agg_ops.SumOp,
     agg_ops.MeanOp,
     agg_ops.CountOp,
+    agg_ops.VarOp,
+    agg_ops.PopVarOp,
+    agg_ops.StdOp,
 )
 
 

tests/system/small/engines/test_aggregation.py

@@ -111,6 +111,20 @@ def test_engines_unary_aggregates(
     assert_equivalence_execution(node, REFERENCE_ENGINE, engine)
 
 
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+@pytest.mark.parametrize(
+    "op",
+    [agg_ops.std_op, agg_ops.var_op, agg_ops.PopVarOp()],
+)
+def test_engines_unary_variance_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)
+
+
 def test_sql_engines_median_op_aggregates(
     scalars_array_value: array_value.ArrayValue,
     bigquery_client: bigquery.Client,