improve expression engine aggstates (#2282)

pola-rs · Jan 6, 2022 · 653585d · 653585d
1 parent e137974
commit 653585d
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 1 deletion.
diff --git a/polars/polars-lazy/src/physical_plan/expressions/mod.rs b/polars/polars-lazy/src/physical_plan/expressions/mod.rs
@@ -228,7 +228,15 @@ impl<'a> AggregationContext<'a> {
                 assert_eq!(series.len(), self.groups.len());
                 AggState::AggregatedList(series)
             }
-            _ => AggState::NotAggregated(series),
+            _ => {
+                // already aggregated to sum, min even this series was flattened it never could
+                // retrieve the length before grouping, so it stays  in this state.
+                if let AggState::AggregatedFlat(_) = self.series {
+                    AggState::AggregatedFlat(series)
+                } else {
+                    AggState::NotAggregated(series)
+                }
+            }
         };
         self
     }

diff --git a/polars/polars-lazy/src/tests/queries.rs b/polars/polars-lazy/src/tests/queries.rs
@@ -1948,6 +1948,47 @@ fn test_round_after_agg() -> Result<()> {
         .collect()?;
 
     assert!(out.column("foo")?.f32().is_ok());
+
+    let df = df![
+        "groups" => ["pigeon",
+                 "rabbit",
+                 "rabbit",
+                 "Chris",
+                 "pigeon",
+                 "fast",
+                 "fast",
+                 "pigeon",
+                 "rabbit",
+                 "Chris"],
+        "b" => [5409, 4848, 4864, 3540, 8103, 3083, 8575, 9963, 8809, 5425],
+        "c" => [0.4517241160719615,
+                  0.2551467646274673,
+                  0.8682045191407308,
+                  0.9925316385786037,
+                  0.5392027792928116,
+                  0.7633847828107002,
+                  0.7967295231651537,
+                  0.01444779067224733,
+                  0.23807484087472652,
+                  0.10985868798350984]
+    ]?;
+
+    let out = df
+        .lazy()
+        .groupby_stable([col("groups")])
+        .agg([((col("b") * col("c")).sum() / col("b").sum())
+            .round(2)
+            .alias("foo")])
+        .collect()?;
+
+    let out = out.column("foo")?;
+    let out = out.f64()?;
+
+    assert_eq!(
+        Vec::from(out),
+        &[Some(0.3), Some(0.41), Some(0.46), Some(0.79)]
+    );
+
     Ok(())
 }