fix[rust]: allow explode aggregation context (#4664)

pola-rs · Sep 1, 2022 · 298b978 · 298b978
1 parent 5afb328
commit 298b978
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 2 deletions.
diff --git a/polars/polars-lazy/src/physical_plan/expressions/take.rs b/polars/polars-lazy/src/physical_plan/expressions/take.rs
@@ -165,7 +165,7 @@ impl PhysicalExpr for TakeExpr {
         let s = idx.cast(&DataType::List(Box::new(IDX_DTYPE)))?;
         let idx = s.list().unwrap();
 
-        let taken = ac
+        let mut taken = ac
             .aggregated()
             .list()
             .unwrap()
@@ -182,6 +182,8 @@ impl PhysicalExpr for TakeExpr {
             })
             .collect::<Result<ListChunked>>()?;
 
+        taken.rename(ac.series().name());
+
         ac.with_series(taken.into_series(), true);
         ac.with_update_groups(UpdateGroups::WithGroupsLen);
         Ok(ac)

diff --git a/polars/polars-lazy/src/physical_plan/planner/expr.rs b/polars/polars-lazy/src/physical_plan/planner/expr.rs
@@ -602,7 +602,7 @@ impl PhysicalPlanner {
                     vec![input],
                     function,
                     node_to_expr(expression, expr_arena),
-                    ApplyOptions::ApplyFlat,
+                    ApplyOptions::ApplyGroups,
                 )))
             }
             Wildcard => panic!("should be no wildcard at this point"),

diff --git a/py-polars/tests/test_explode.py b/py-polars/tests/test_explode.py
@@ -163,3 +163,23 @@ def test_utf8_explode() -> None:
         "a",
         "r",
     ]
+
+
+def test_explode_in_agg_context() -> None:
+    df = pl.DataFrame(
+        {"idxs": [[0], [1], [0, 2]], "array": [[0.0, 3.5], [4.6, 0.0], [0.0, 7.8, 0.0]]}
+    )
+
+    assert (
+        df.with_row_count("row_nr")
+        .explode("idxs")
+        .groupby("row_nr")
+        .agg(
+            [
+                pl.col("array").explode(),
+            ]
+        )
+    ).to_dict(False) == {
+        "row_nr": [0, 1, 2],
+        "array": [[0.0, 3.5], [4.6, 0.0], [0.0, 7.8, 0.0, 0.0, 7.8, 0.0]],
+    }