fix groups state in complex aggregation (#3656)

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

@@ -319,6 +319,8 @@ impl<'a> AggregationContext<'a> {
     /// Calling this function will ensure both are sortened. This will be a no-op
     /// if already aggregated.
     pub(crate) fn sort_by_groups(&mut self) {
+        // make sure that the groups are updated before we use them to sort.
+        self.groups();
         match &self.state {
             AggState::NotAggregated(s) => {
                 // We should not aggregate literals!!

polars/tests/it/lazy/groupby.rs

@@ -83,3 +83,29 @@ fn test_list_arithmetic_in_groupby() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_filter_diff_arithmetic() -> Result<()> {
+    let df = df![
+        "user" => [1, 1, 1, 1, 2],
+        "group" => [1, 2, 1, 1, 2],
+        "value" => [1, 5, 14, 17, 20]
+    ]?;
+
+    let out = df
+        .lazy()
+        .groupby([col("user")])
+        .agg([(col("value")
+            .filter(col("group").eq(lit(1)))
+            .diff(1, Default::default())
+            * lit(2))
+        .alias("diff")])
+        .sort("user", Default::default())
+        .explode([col("diff")])
+        .collect()?;
+
+    let out = out.column("diff")?;
+    assert_eq!(out, &Series::new("diff", &[None, Some(26), Some(6), None]));
+
+    Ok(())
+}