fix bug in flattened apply (groupby)

pola-rs · Dec 15, 2021 · 0d7555a · 0d7555a
1 parent 4923eb5
commit 0d7555a
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 1 deletion.
diff --git a/polars/polars-lazy/Cargo.toml b/polars/polars-lazy/Cargo.toml
@@ -69,6 +69,7 @@ test = [
   "regex",
   "polars-core/plain_fmt",
   "diff",
+  "abs",
 ]
 
 [dependencies]

diff --git a/polars/polars-lazy/src/physical_plan/expressions/apply.rs b/polars/polars-lazy/src/physical_plan/expressions/apply.rs
@@ -79,7 +79,9 @@ impl PhysicalExpr for ApplyExpr {
                 }
                 ApplyOptions::ApplyFlat => {
                     let s = self.function.call_udf(&mut [ac.flat().into_owned()])?;
-                    ac.with_update_groups(UpdateGroups::WithGroupsLen);
+                    if ac.is_aggregated() {
+                        ac.with_update_groups(UpdateGroups::WithGroupsLen);
+                    }
                     ac.with_series(s, false);
                     Ok(ac)
                 }

diff --git a/polars/polars-lazy/src/tests/queries.rs b/polars/polars-lazy/src/tests/queries.rs
@@ -2418,3 +2418,23 @@ fn empty_df() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_apply_flatten() -> Result<()> {
+    let df = df![
+         "A"=> [1.1435, 2.223456, 3.44732, -1.5234, -2.1238, -3.2923],
+        "B"=> ["a", "b", "a", "b", "a", "b"]
+    ]?;
+
+    let out = df
+        .lazy()
+        .stable_groupby([col("B")])
+        .agg([col("A").abs().sum()])
+        .collect()?;
+
+    let out = out.column("A_sum")?;
+    assert_eq!(out.get(0), AnyValue::Float64(6.71462));
+    assert_eq!(out.get(1), AnyValue::Float64(7.039156));
+
+    Ok(())
+}