Skip to content

Commit

Permalink
feat(rust, python): raise error on invalid binary cmp (pola-rs#6564)
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 authored and cojmeister committed Jan 30, 2023
1 parent a7e037d commit ce70556
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 1 deletion.
11 changes: 10 additions & 1 deletion polars/polars-lazy/src/physical_plan/expressions/binary.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ use polars_core::series::unstable::UnstableSeries;
use polars_core::POOL;
use rayon::prelude::*;

use crate::physical_plan::errors::expression_err;
use crate::physical_plan::state::ExecutionState;
use crate::prelude::*;

Expand Down Expand Up @@ -108,7 +109,15 @@ impl PhysicalExpr for BinaryExpr {
|| self.right.evaluate(df, &state),
)
});
apply_operator_owned(lhs?, rhs?, self.op)
let lhs = lhs?;
let rhs = rhs?;
let lhs_len = lhs.len();
let rhs_len = rhs.len();
if lhs_len != rhs_len && !(lhs_len == 1 || rhs_len == 1) {
let msg = format!("Cannot evaluate two Series of different length. Got lhs of length: {lhs_len} and rhs of length: {rhs_len}.");
return Err(expression_err!(msg, self.expr, ComputeError));
}
apply_operator_owned(lhs, rhs, self.op)
}

#[allow(clippy::ptr_arg)]
Expand Down
14 changes: 14 additions & 0 deletions py-polars/tests/unit/test_errors.py
Original file line number Diff line number Diff line change
Expand Up @@ -418,3 +418,17 @@ def test_err_on_multiple_column_expansion() -> None:
"d": [4],
}
).select([pl.col(["a", "b"]) + pl.col(["c", "d"])])


def test_compare_different_len() -> None:
df = pl.DataFrame(
{
"idx": list(range(5)),
}
)

s = pl.Series([2, 5, 8])
with pytest.raises(
pl.ComputeError, match=r"annot evaluate two Series of different length"
):
df.filter(pl.col("idx") == s)

0 comments on commit ce70556

Please sign in to comment.