diff --git a/.typos.toml b/.typos.toml index ed5f7bde4b64..d4b3e9f418f3 100644 --- a/.typos.toml +++ b/.typos.toml @@ -4,6 +4,7 @@ extend-exclude = [ "*.csv", "*.gz", "dists.dss", + "**/images/*", ] ignore-hidden = false diff --git a/crates/polars-core/src/series/arithmetic/borrowed.rs b/crates/polars-core/src/series/arithmetic/borrowed.rs index 6003d0b05792..fba1b03a1332 100644 --- a/crates/polars-core/src/series/arithmetic/borrowed.rs +++ b/crates/polars-core/src/series/arithmetic/borrowed.rs @@ -130,6 +130,26 @@ fn array_shape(dt: &DataType, infer: bool) -> Vec { 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( @@ -137,13 +157,15 @@ impl ArrayChunked { rhs: &Series, op: &dyn Fn(Series, Series) -> PolarsResult, ) -> PolarsResult { - 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() }; diff --git a/py-polars/tests/unit/operations/arithmetic/test_array.py b/py-polars/tests/unit/operations/arithmetic/test_array.py new file mode 100644 index 000000000000..35142f139ba7 --- /dev/null +++ b/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]], + }