fix: Properly broadcast array arithmetic (#18851)

Author: ritchie46

.typos.toml

@@ -4,6 +4,7 @@ extend-exclude = [
   "*.csv",
   "*.gz",
   "dists.dss",
+  "**/images/*",
 ]
 ignore-hidden = false
 

crates/polars-core/src/series/arithmetic/borrowed.rs

@@ -130,20 +130,42 @@ fn array_shape(dt: &DataType, infer: bool) -> Vec<i64> {
     buf
 }
 
+#[cfg(feature = "dtype-array")]
+fn broadcast_array(lhs: &ArrayChunked, rhs: &Series) -> PolarsResult<(ArrayChunked, Series)> {
+    let out = match (lhs.len(), rhs.len()) {
+        (1, _) => (lhs.new_from_index(0, rhs.len()), rhs.clone()),
+        (_, 1) => {
+            // Numeric scalars will be broadcasted implicitly without intermediate allocation.
+            if rhs.dtype().is_numeric() {
+                (lhs.clone(), rhs.clone())
+            } else {
+                (lhs.clone(), rhs.new_from_index(0, lhs.len()))
+            }
+        },
+        (a, b) if a == b => (lhs.clone(), rhs.clone()),
+        _ => {
+            polars_bail!(InvalidOperation: "can only do arithmetic of array's of the same type and shape; got {} and {}", lhs.dtype(), rhs.dtype())
+        },
+    };
+    Ok(out)
+}
+
 #[cfg(feature = "dtype-array")]
 impl ArrayChunked {
     fn arithm_helper(
         &self,
         rhs: &Series,
         op: &dyn Fn(Series, Series) -> PolarsResult<Series>,
     ) -> PolarsResult<Series> {
-        let l_leaf_array = self.clone().into_series().get_leaf_array();
-        let shape = array_shape(self.dtype(), true);
+        let (lhs, rhs) = broadcast_array(self, rhs)?;
+
+        let l_leaf_array = lhs.clone().into_series().get_leaf_array();
+        let shape = array_shape(lhs.dtype(), true);
 
         let r_leaf_array = if rhs.dtype().is_numeric() && rhs.len() == 1 {
             rhs.clone()
         } else {
-            polars_ensure!(self.dtype() == rhs.dtype(), InvalidOperation: "can only do arithmetic of array's of the same type and shape; got {} and {}", self.dtype(), rhs.dtype());
+            polars_ensure!(lhs.dtype() == rhs.dtype(), InvalidOperation: "can only do arithmetic of array's of the same type and shape; got {} and {}", self.dtype(), rhs.dtype());
             rhs.get_leaf_array()
         };
 

py-polars/tests/unit/operations/arithmetic/test_array.py

@@ -0,0 +1,26 @@
+import polars as pl
+
+
+def test_literal_broadcast_array() -> None:
+    df = pl.DataFrame({"A": [[0.1, 0.2], [0.3, 0.4]]}).cast(pl.Array(float, 2))
+
+    lit = pl.lit([3, 5], pl.Array(float, 2))
+    assert df.select(
+        mul=pl.all() * lit,
+        div=pl.all() / lit,
+        add=pl.all() + lit,
+        sub=pl.all() - lit,
+        div_=lit / pl.all(),
+        add_=lit + pl.all(),
+        sub_=lit - pl.all(),
+        mul_=lit * pl.all(),
+    ).to_dict(as_series=False) == {
+        "mul": [[0.30000000000000004, 1.0], [0.8999999999999999, 2.0]],
+        "div": [[0.03333333333333333, 0.04], [0.09999999999999999, 0.08]],
+        "add": [[3.1, 5.2], [3.3, 5.4]],
+        "sub": [[-2.9, -4.8], [-2.7, -4.6]],
+        "div_": [[30.0, 25.0], [10.0, 12.5]],
+        "add_": [[3.1, 5.2], [3.3, 5.4]],
+        "sub_": [[2.9, 4.8], [2.7, 4.6]],
+        "mul_": [[0.30000000000000004, 1.0], [0.8999999999999999, 2.0]],
+    }