Skip to content

Commit d3e6e03

Browse files
authored
fix(rust, python): unnest only pushdown column if there are projections (#5360)
1 parent e97088a commit d3e6e03

File tree

2 files changed

+33
-1
lines changed

2 files changed

+33
-1
lines changed

polars/polars-lazy/polars-plan/src/logical_plan/optimizer/projection_pushdown.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1149,7 +1149,7 @@ impl ProjectionPushDown {
11491149
input,
11501150
function: function.clone(),
11511151
};
1152-
if function.allow_projection_pd() {
1152+
if function.allow_projection_pd() && !acc_projections.is_empty() {
11531153
for name in function.additional_projection_pd_columns() {
11541154
let node = expr_arena.add(AExpr::Column(name.clone()));
11551155
add_expr_to_accumulated(

py-polars/tests/unit/test_projections.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,3 +62,35 @@ def test_unnest_projection_pushdown() -> None:
6262
"col": ["z", "z", "c", "c"],
6363
"value": [1, 2, 2, 3],
6464
}
65+
66+
67+
def test_unnest_columns_available() -> None:
68+
df = pl.DataFrame(
69+
{
70+
"title": ["Avatar", "spectre", "King Kong"],
71+
"content_rating": ["PG-13"] * 3,
72+
"genres": [
73+
"Action|Adventure|Fantasy|Sci-Fi",
74+
"Action|Adventure|Thriller",
75+
"Action|Adventure|Drama|Romance",
76+
],
77+
}
78+
).lazy()
79+
80+
q = df.with_column(
81+
pl.col("genres")
82+
.str.split("|")
83+
.arr.to_struct(
84+
n_field_strategy="max_width", name_generator=lambda i: f"genre{i+1}"
85+
)
86+
).unnest("genres")
87+
88+
out = q.collect()
89+
assert out.to_dict(False) == {
90+
"title": ["Avatar", "spectre", "King Kong"],
91+
"content_rating": ["PG-13", "PG-13", "PG-13"],
92+
"genre1": ["Action", "Action", "Action"],
93+
"genre2": ["Adventure", "Adventure", "Adventure"],
94+
"genre3": ["Fantasy", "Thriller", "Drama"],
95+
"genre4": ["Sci-Fi", None, "Romance"],
96+
}

0 commit comments

Comments
 (0)