fix aritmetic bug introduced in #3709 (#3741)

pola-rs · Jun 19, 2022 · 6cafa8b · 6cafa8b
1 parent 4c6f9cd
commit 6cafa8b
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
diff --git a/polars/polars-lazy/src/physical_plan/expressions/binary.rs b/polars/polars-lazy/src/physical_plan/expressions/binary.rs
@@ -313,7 +313,7 @@ impl PhysicalExpr for BinaryExpr {
             // flatten the Series and apply the operators
             (AggState::AggregatedList(_), AggState::AggregatedList(_), _) => {
                 let lhs = ac_l.flat_naive().as_ref().clone();
-                let rhs = ac_l.flat_naive().as_ref().clone();
+                let rhs = ac_r.flat_naive().as_ref().clone();
 
                 // drop lhs so that we might operate in place
                 {

diff --git a/py-polars/tests/test_queries.py b/py-polars/tests/test_queries.py
@@ -183,3 +183,28 @@ def test_groupby_agg_equals_zero_3535() -> None:
         "val1": [10, 0, -99],
         "val2": [None, 0.0, 10.5],
     }
+
+
+def test_arithmetic_in_aggregation_3739() -> None:
+    def demean_dot() -> pl.Expr:
+        x = pl.col("x")
+        y = pl.col("y")
+        x1 = x - x.mean()
+        y1 = y - y.mean()
+        return (x1 * y1).sum().alias("demean_dot")
+
+    assert (
+        pl.DataFrame(
+            {
+                "key": ["a", "a", "a", "a"],
+                "x": [4, 2, 2, 4],
+                "y": [2, 0, 2, 0],
+            }
+        )
+        .groupby("key")
+        .agg(
+            [
+                demean_dot(),
+            ]
+        )
+    ).to_dict(False) == {"key": ["a"], "demean_dot": [0.0]}