Skip to content

Commit

Permalink
fix: binary operations in aggregation context on literals (pola-rs#12430
Browse files Browse the repository at this point in the history
)
  • Loading branch information
nameexhaustion authored Nov 14, 2023
1 parent 8af2048 commit dac29fb
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 2 deletions.
8 changes: 6 additions & 2 deletions crates/polars-lazy/src/physical_plan/expressions/binary.rs
Original file line number Diff line number Diff line change
Expand Up @@ -114,9 +114,13 @@ impl BinaryExpr {
let left_s = ac_l.series().rechunk();
let right_s = ac_r.series().rechunk();
let res_s = apply_operator(&left_s, &right_s, self.op)?;
let ca = ListChunked::full(&name, &res_s, ac_l.groups.len());
ac_l.with_update_groups(UpdateGroups::WithSeriesLen);
ac_l.with_series(ca.into_series(), true, Some(&self.expr))?;
let res_s = if res_s.len() == 1 {
res_s.new_from_index(0, ac_l.groups.len())
} else {
ListChunked::full(&name, &res_s, ac_l.groups.len()).into_series()
};
ac_l.with_series(res_s, true, Some(&self.expr))?;
Ok(ac_l)
}

Expand Down
13 changes: 13 additions & 0 deletions py-polars/tests/unit/operations/test_aggregations.py
Original file line number Diff line number Diff line change
Expand Up @@ -315,3 +315,16 @@ def test_first_last_unit_length_12363() -> None:
"a_last": [2],
"b_last": [None],
}


def test_binary_op_agg_context_no_simplify_expr_12423() -> None:
expect = pl.DataFrame({"x": [1], "y": [1]}, schema={"x": pl.Int64, "y": pl.Int32})

for simplify_expression in (True, False):
assert_frame_equal(
expect,
pl.LazyFrame({"x": [1]})
.group_by("x")
.agg(y=pl.lit(1) * pl.lit(1))
.collect(simplify_expression=simplify_expression),
)

0 comments on commit dac29fb

Please sign in to comment.