Skip to content

Commit ef5e83a

Browse files
feat: Polars engine supports std, var (#2215)
1 parent 5e006e4 commit ef5e83a

File tree

3 files changed

+21
-2
lines changed

3 files changed

+21
-2
lines changed

bigframes/core/compile/polars/compiler.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -535,9 +535,11 @@ def compile_agg_op(
535535
if isinstance(op, agg_ops.StdOp):
536536
return pl.std(inputs[0])
537537
if isinstance(op, agg_ops.VarOp):
538-
return pl.var(inputs[0])
538+
# polars var doesnt' support decimal, so use std instead
539+
return pl.std(inputs[0]).pow(2)
539540
if isinstance(op, agg_ops.PopVarOp):
540-
return pl.var(inputs[0], ddof=0)
541+
# polars var doesnt' support decimal, so use std instead
542+
return pl.std(inputs[0], ddof=0).pow(2)
541543
if isinstance(op, agg_ops.FirstNonNullOp):
542544
return pl.col(*inputs).drop_nulls().first()
543545
if isinstance(op, agg_ops.LastNonNullOp):

bigframes/session/polars_executor.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,9 @@
103103
agg_ops.SumOp,
104104
agg_ops.MeanOp,
105105
agg_ops.CountOp,
106+
agg_ops.VarOp,
107+
agg_ops.PopVarOp,
108+
agg_ops.StdOp,
106109
)
107110

108111

tests/system/small/engines/test_aggregation.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,20 @@ def test_engines_unary_aggregates(
111111
assert_equivalence_execution(node, REFERENCE_ENGINE, engine)
112112

113113

114+
@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
115+
@pytest.mark.parametrize(
116+
"op",
117+
[agg_ops.std_op, agg_ops.var_op, agg_ops.PopVarOp()],
118+
)
119+
def test_engines_unary_variance_aggregates(
120+
scalars_array_value: array_value.ArrayValue,
121+
engine,
122+
op,
123+
):
124+
node = apply_agg_to_all_valid(scalars_array_value, op).node
125+
assert_equivalence_execution(node, REFERENCE_ENGINE, engine)
126+
127+
114128
def test_sql_engines_median_op_aggregates(
115129
scalars_array_value: array_value.ArrayValue,
116130
bigquery_client: bigquery.Client,

0 commit comments

Comments
 (0)