Skip to content

Commit

Permalink
fix: Fix projection pushdown with literal joins (#16981)
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 authored Jun 16, 2024
1 parent 24dafd9 commit 1effc24
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 7 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,15 @@ fn add_keys_to_accumulated_state(
// the JOIN executor
if add_local {
// take the left most name as output name
let name = aexpr_to_leaf_name(expr, expr_arena);
let node = expr_arena.add(AExpr::Column(name.clone()));
local_projection.push(ColumnNode(node));
Some(name)
let mut iter = aexpr_to_leaf_names_iter(expr, expr_arena);
if let Some(name) = iter.next() {
drop(iter);
let node = expr_arena.add(AExpr::Column(name.clone()));
local_projection.push(ColumnNode(node));
Some(name)
} else {
None
}
} else {
None
}
Expand Down Expand Up @@ -250,15 +255,14 @@ pub(super) fn process_join(
continue;
}

add_keys_to_accumulated_state(
let _ = add_keys_to_accumulated_state(
e.node(),
&mut pushdown_left,
&mut local_projection,
&mut names_left,
expr_arena,
true,
)
.unwrap();
);
}

// For left and innner joins we can set `coalesce` to `true` if the rhs key columns are not projected.
Expand Down
9 changes: 9 additions & 0 deletions py-polars/tests/unit/operations/test_join.py
Original file line number Diff line number Diff line change
Expand Up @@ -961,3 +961,12 @@ def test_join_raise_on_repeated_expression_key_names(coalesce: bool) -> None:
left.join(
right, on=[pl.col("a"), pl.col("a") % 2], how="full", coalesce=coalesce
)


def test_join_lit_panic_11410() -> None:
df = pl.LazyFrame({"date": [1, 2, 3], "symbol": [4, 5, 6]})
dates = df.select("date").unique(maintain_order=True)
symbols = df.select("symbol").unique(maintain_order=True)
assert symbols.join(dates, left_on=pl.lit(1), right_on=pl.lit(1)).drop(
"literal"
).collect().to_dict(as_series=False) == {"symbol": [4], "date": [1]}

0 comments on commit 1effc24

Please sign in to comment.