From 173ac1c189dad56d0c968b1e217359b8865bae7a Mon Sep 17 00:00:00 2001 From: Gijs Burghoorn Date: Thu, 30 May 2024 08:07:21 +0200 Subject: [PATCH] fix(rust): bail on bool `floordiv` (#16578) --- crates/polars-ops/src/series/ops/floor_divide.rs | 4 ++++ .../unit/operations/arithmetic/test_arithmetic.py | 10 ++++++++++ 2 files changed, 14 insertions(+) diff --git a/crates/polars-ops/src/series/ops/floor_divide.rs b/crates/polars-ops/src/series/ops/floor_divide.rs index 85d8750a5313..4a78cd179890 100644 --- a/crates/polars-ops/src/series/ops/floor_divide.rs +++ b/crates/polars-ops/src/series/ops/floor_divide.rs @@ -29,6 +29,10 @@ pub fn floor_div_series(a: &Series, b: &Series) -> PolarsResult { _ => {}, } + if !a.dtype().is_numeric() { + polars_bail!(op = "floor_div", a.dtype()); + } + let logical_type = a.dtype(); let a = a.to_physical_repr(); diff --git a/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py b/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py index 5d31bd1e2559..c0d6ce80e5d4 100644 --- a/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py +++ b/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py @@ -293,3 +293,13 @@ def test_null_column_arithmetic(op: Any) -> None: # test broadcast left output_df = df.select(op(pl.Series("a", [None]), pl.col("a"))) assert_frame_equal(expected_df, output_df) + + +def test_bool_floordiv() -> None: + df = pl.DataFrame({"x": [True]}) + + with pytest.raises( + pl.InvalidOperationError, + match="floor_div operation not supported for dtype `bool`", + ): + df.with_columns(pl.col("x").floordiv(2))