fix(rust, python): unnest only pushdown column if there are projectio…

…ns (#5360)
pola-rs · Oct 27, 2022 · d3e6e03 · d3e6e03
1 parent e97088a
commit d3e6e03
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 1 deletion.
diff --git a/polars/polars-lazy/polars-plan/src/logical_plan/optimizer/projection_pushdown.rs b/polars/polars-lazy/polars-plan/src/logical_plan/optimizer/projection_pushdown.rs
@@ -1149,7 +1149,7 @@ impl ProjectionPushDown {
                     input,
                     function: function.clone(),
                 };
-                if function.allow_projection_pd() {
+                if function.allow_projection_pd() && !acc_projections.is_empty() {
                     for name in function.additional_projection_pd_columns() {
                         let node = expr_arena.add(AExpr::Column(name.clone()));
                         add_expr_to_accumulated(

diff --git a/py-polars/tests/unit/test_projections.py b/py-polars/tests/unit/test_projections.py
@@ -62,3 +62,35 @@ def test_unnest_projection_pushdown() -> None:
         "col": ["z", "z", "c", "c"],
         "value": [1, 2, 2, 3],
     }
+
+
+def test_unnest_columns_available() -> None:
+    df = pl.DataFrame(
+        {
+            "title": ["Avatar", "spectre", "King Kong"],
+            "content_rating": ["PG-13"] * 3,
+            "genres": [
+                "Action|Adventure|Fantasy|Sci-Fi",
+                "Action|Adventure|Thriller",
+                "Action|Adventure|Drama|Romance",
+            ],
+        }
+    ).lazy()
+
+    q = df.with_column(
+        pl.col("genres")
+        .str.split("|")
+        .arr.to_struct(
+            n_field_strategy="max_width", name_generator=lambda i: f"genre{i+1}"
+        )
+    ).unnest("genres")
+
+    out = q.collect()
+    assert out.to_dict(False) == {
+        "title": ["Avatar", "spectre", "King Kong"],
+        "content_rating": ["PG-13", "PG-13", "PG-13"],
+        "genre1": ["Action", "Action", "Action"],
+        "genre2": ["Adventure", "Adventure", "Adventure"],
+        "genre3": ["Fantasy", "Thriller", "Drama"],
+        "genre4": ["Sci-Fi", None, "Romance"],
+    }