fix groups update to match exploded offsets (#3010)

pola-rs · Mar 30, 2022 · 874913d · 874913d
1 parent 7adf768
commit 874913d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/polars/polars-lazy/src/physical_plan/expressions/mod.rs b/polars/polars-lazy/src/physical_plan/expressions/mod.rs
@@ -140,7 +140,10 @@ impl<'a> AggregationContext<'a> {
                             .iter()
                             .map(|&o| {
                                 let len = (o - previous) as IdxSize;
-                                let new_offset = offset + len;
+                                // explode will fill empty rows with null, so we must increment the group
+                                // offset accordingly
+                                let new_offset = offset + len + (len == 0) as IdxSize;
+
                                 previous = o;
                                 let out = [offset, len];
                                 offset = new_offset;

diff --git a/polars/tests/it/lazy/groupby.rs b/polars/tests/it/lazy/groupby.rs
@@ -14,12 +14,12 @@ fn test_filter_sort_diff_2984() -> Result<()> {
         .groupby([col("group")])
         .agg([col("id")
             .filter(col("id").lt(lit(3)))
-            .sort(true)
+            .sort(false)
             .diff(1, Default::default())
             .sum()])
         .sort("group", Default::default())
         .collect()?;
 
-    assert_eq!(Vec::from(out.column("id")?.i32()?), &[Some(-1), None]);
+    assert_eq!(Vec::from(out.column("id")?.i32()?), &[Some(1), None]);
     Ok(())
 }