Lazy: fix inconsistency in apply_flat

polars/polars-lazy/src/physical_plan/expressions/apply.rs

@@ -75,7 +75,7 @@ impl PhysicalExpr for ApplyExpr {
                 Ok(ac)
             }
             ApplyOptions::ApplyFlat => {
-                let s = self.function.call_udf(&mut [ac.take()])?;
+                let s = self.function.call_udf(&mut [ac.flat().into_owned()])?;
                 ac.with_series(s);
                 Ok(ac)
             }
@@ -143,7 +143,21 @@ impl PhysicalAggregation for ApplyExpr {
                 }
                 Ok(Some(ca.into_series()))
             }
-            ApplyOptions::ApplyFlat => self.function.call_udf(&mut [ac.take()]).map(Some),
+            ApplyOptions::ApplyFlat => {
+                // the function needs to be called on a flat series
+                // but the series may be flat or aggregated
+                // if its flat, we just apply and return
+                // if not flat, the flattening sorts by group, so we must create new group tuples
+                // and again aggregate.
+                let out = self.function.call_udf(&mut [ac.flat().into_owned()]);
+                if ac.is_flat() {
+                    out.map(Some)
+                } else {
+                    // TODO! maybe just apply over list?
+                    ac.with_update_groups(UpdateGroups::WithGroupsLen);
+                    Ok(out?.agg_list(ac.groups()))
+                }
+            }
             ApplyOptions::ApplyList => self
                 .function
                 .call_udf(&mut [ac.aggregated().into_owned()])

polars/polars-lazy/src/physical_plan/expressions/mod.rs

@@ -147,6 +147,10 @@ impl<'a> AggregationContext<'a> {
         }
     }
 
+    pub(crate) fn is_flat(&self) -> bool {
+        matches!(&self.series, AggState::Flat(_))
+    }
+
     pub(crate) fn combine_groups(&mut self, other: AggregationContext) -> &mut Self {
         if let (Cow::Borrowed(_), Cow::Owned(a)) = (&self.groups, other.groups) {
             self.groups = Cow::Owned(a);

polars/polars-lazy/src/test.rs

@@ -1848,3 +1848,25 @@ fn test_round_after_agg() -> Result<()> {
     assert!(out.column("foo")?.f64().is_ok());
     Ok(())
 }
+
+#[test]
+fn test_power_in_agg_list() -> Result<()> {
+    let df = fruits_cars();
+
+    let out = df
+        .lazy()
+        .groupby(vec![col("fruits")])
+        .agg(vec![col("A")
+            .rolling_min(1, None, false, 0)
+            .pow(2.0)
+            .alias("foo")])
+        .sort("fruits", true)
+        .collect()?;
+
+    let agg = out.column("foo")?.list()?;
+    let first = agg.get(0).unwrap();
+    let vals = first.f64()?;
+    assert_eq!(Vec::from(vals), &[Some(1.0), Some(4.0), Some(25.0)]);
+
+    Ok(())
+}