fix: Raise on invalid shape dataframe arithmetic (#17322)

pola-rs · Jul 1, 2024 · 120ec7f · 120ec7f
1 parent ccdfbda
commit 120ec7f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/crates/polars-core/src/frame/arithmetic.rs b/crates/polars-core/src/frame/arithmetic.rs
@@ -145,7 +145,7 @@ impl DataFrame {
             let df = if col_len < self.width() { self } else { other };
 
             for i in col_len..max_len {
-                let s = &df.get_columns()[i];
+                let s = &df.get_columns().get(i).ok_or_else(|| polars_err!(InvalidOperation: "cannot do arithmetic on DataFrames with shapes: {:?} and {:?}", self.shape(), other.shape()))?;
                 let name = s.name();
                 let dtype = s.dtype();
 

diff --git a/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py b/py-polars/tests/unit/operations/arithmetic/test_arithmetic.py
@@ -760,3 +760,8 @@ def test_date_datetime_sub() -> None:
         "foo": [timedelta(days=-4)],
         "bar": [timedelta(days=4)],
     }
+
+
+def test_raise_invalid_shape() -> None:
+    with pytest.raises(pl.exceptions.InvalidOperationError):
+        pl.DataFrame([[1, 2], [3, 4]]) * pl.DataFrame([1, 2, 3])