diff --git a/bigframes/core/compile/polars/compiler.py b/bigframes/core/compile/polars/compiler.py index d48ddba0cc..3211d6ebf7 100644 --- a/bigframes/core/compile/polars/compiler.py +++ b/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): diff --git a/bigframes/session/polars_executor.py b/bigframes/session/polars_executor.py index a1e1d436e1..a399b289bb 100644 --- a/bigframes/session/polars_executor.py +++ b/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, ) diff --git a/tests/system/small/engines/test_aggregation.py b/tests/system/small/engines/test_aggregation.py index d71013c648..3e6d4843de 100644 --- a/tests/system/small/engines/test_aggregation.py +++ b/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,