Repository: pola-rs/polars (language: python)

Issue:
Array arithmetic is not elementwise for zero-width arrays
### Checks

- [X] I have checked that this issue has not already been reported.
- [X] I have confirmed this bug exists on the [latest version](https://pypi.org/project/polars/) of Polars.

### Reproducible example

```python
import polars as pl

a = pl.Series('a', [[]], pl.Array(pl.Int64, 0))
print(a + a)
assert len(a + a) == 1
```


### Log output

```shell
shape: (0,)
Series: 'a' [array[i64, 0]]
[
]
Traceback (most recent call last):
  File "/home/johndoe/Projects/polars/array_elementwise_bug.py", line 5, in <module>
    assert len(a + a) == 1
           ^^^^^^^^^^^^^^^
AssertionError
```


### Issue description

`[[]] + [[]]` produces `[]`

### Expected behavior

Arithmetic operations should be elementwise and thus produce an output that has the same length as the inputs

### Installed versions

<details>

```
Replace this line with the output of pl.show_versions(). Leave the backticks in place.
```

</details>



Discussion / hints:

Fix patch:
diff --git a/crates/polars-core/src/chunked_array/array/mod.rs b/crates/polars-core/src/chunked_array/array/mod.rs
index 3e0e47a7e86a..49f0bf7ce1bf 100644
--- a/crates/polars-core/src/chunked_array/array/mod.rs
+++ b/crates/polars-core/src/chunked_array/array/mod.rs
@@ -81,4 +81,13 @@ impl ArrayChunked {
 
         ArrayChunked::try_from_chunk_iter(self.name().clone(), chunks)
     }
+
+    /// Recurse nested types until we are at the leaf array.
+    pub fn get_leaf_array(&self) -> Series {
+        let mut current = self.get_inner();
+        while let Some(child_array) = current.try_array() {
+            current = child_array.get_inner();
+        }
+        current
+    }
 }
diff --git a/crates/polars-core/src/series/arithmetic/borrowed.rs b/crates/polars-core/src/series/arithmetic/borrowed.rs
index f9e5ff42139b..01065c66c1d1 100644
--- a/crates/polars-core/src/series/arithmetic/borrowed.rs
+++ b/crates/polars-core/src/series/arithmetic/borrowed.rs
@@ -51,6 +51,18 @@ where
     ChunkedArray<T>: IntoSeries,
 {
     fn subtract(lhs: &ChunkedArray<T>, rhs: &Series) -> PolarsResult<Series> {
+        #[cfg(feature = "dtype-array")]
+        if let Some(rhs) = rhs.try_array() {
+            return rhs.arithm_helper_scalar_lhs(lhs.clone().into_series(), &|l, r| l.subtract(&r));
+        }
+
+        polars_ensure!(
+            lhs.dtype() == rhs.dtype(),
+            opq = add,
+            rhs.dtype(),
+            rhs.dtype()
+        );
+
         // SAFETY:
         // There will be UB if a ChunkedArray is alive with the wrong datatype.
         // we now only create the potentially wrong dtype for a short time.
@@ -61,6 +73,18 @@ where
         Ok(out.into_series())
     }
     fn add_to(lhs: &ChunkedArray<T>, rhs: &Series) -> PolarsResult<Series> {
+        #[cfg(feature = "dtype-array")]
+        if let Some(rhs) = rhs.try_array() {
+            return rhs.arithm_helper_scalar_lhs(lhs.clone().into_series(), &|l, r| l.add_to(&r));
+        }
+
+        polars_ensure!(
+            lhs.dtype() == rhs.dtype(),
+            opq = add,
+            rhs.dtype(),
+            rhs.dtype()
+        );
+
         // SAFETY:
         // see subtract
         let rhs = unsafe { lhs.unpack_series_matching_physical_type(rhs) };
@@ -68,6 +92,18 @@ where
         Ok(out.into_series())
     }
     fn multiply(lhs: &ChunkedArray<T>, rhs: &Series) -> PolarsResult<Series> {
+        #[cfg(feature = "dtype-array")]
+        if let Some(rhs) = rhs.try_array() {
+            return rhs.arithm_helper_scalar_lhs(lhs.clone().into_series(), &|l, r| l.multiply(&r));
+        }
+
+        polars_ensure!(
+            lhs.dtype() == rhs.dtype(),
+            opq = add,
+            rhs.dtype(),
+            rhs.dtype()
+        );
+
         // SAFETY:
         // see subtract
         let rhs = unsafe { lhs.unpack_series_matching_physical_type(rhs) };
@@ -75,6 +111,18 @@ where
         Ok(out.into_series())
     }
     fn divide(lhs: &ChunkedArray<T>, rhs: &Series) -> PolarsResult<Series> {
+        #[cfg(feature = "dtype-array")]
+        if let Some(rhs) = rhs.try_array() {
+            return rhs.arithm_helper_scalar_lhs(lhs.clone().into_series(), &|l, r| l.divide(&r));
+        }
+
+        polars_ensure!(
+            lhs.dtype() == rhs.dtype(),
+            opq = add,
+            rhs.dtype(),
+            rhs.dtype()
+        );
+
         // SAFETY:
         // see subtract
         let rhs = unsafe { lhs.unpack_series_matching_physical_type(rhs) };
@@ -82,6 +130,19 @@ where
         Ok(out.into_series())
     }
     fn remainder(lhs: &ChunkedArray<T>, rhs: &Series) -> PolarsResult<Series> {
+        #[cfg(feature = "dtype-array")]
+        if let Some(rhs) = rhs.try_array() {
+            return rhs
+                .arithm_helper_scalar_lhs(lhs.clone().into_series(), &|l, r| l.remainder(&r));
+        }
+
+        polars_ensure!(
+            lhs.dtype() == rhs.dtype(),
+            opq = add,
+            rhs.dtype(),
+            rhs.dtype()
+        );
+
         // SAFETY:
         // see subtract
         let rhs = unsafe { lhs.unpack_series_matching_physical_type(rhs) };
@@ -114,24 +175,6 @@ impl NumOpsDispatchInner for BooleanType {
     }
 }
 
-#[cfg(feature = "dtype-array")]
-fn array_shape(dt: &DataType, infer: bool) -> Vec<ReshapeDimension> {
-    fn inner(dt: &DataType, buf: &mut Vec<ReshapeDimension>) {
-        if let DataType::Array(_, size) = dt {
-            buf.push(ReshapeDimension::Specified(
-                Dimension::try_from(*size as i64).unwrap(),
-            ))
-        }
-    }
-
-    let mut buf = vec![];
-    if infer {
-        buf.push(ReshapeDimension::Infer)
-    }
-    inner(dt, &mut buf);
-    buf
-}
-
 #[cfg(feature = "dtype-array")]
 fn broadcast_array(lhs: &ArrayChunked, rhs: &Series) -> PolarsResult<(ArrayChunked, Series)> {
     let out = match (lhs.len(), rhs.len()) {
@@ -165,19 +208,56 @@ impl ArrayChunked {
     ) -> PolarsResult<Series> {
         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);
+        polars_ensure!(
+            lhs.dtype() == rhs.dtype()
 
-        let r_leaf_array = if rhs.dtype().is_numeric() && rhs.len() == 1 {
-            rhs.clone()
-        } else {
-            polars_ensure!(lhs.dtype() == rhs.dtype(), InvalidOperation: "can only do arithmetic of arrays of the same type and shape; got {} and {}", self.dtype(), rhs.dtype());
-            rhs.get_leaf_array()
-        };
+            // @NOTE: we allow the arithmetic operations with a scalar of the leaf array
+            || rhs.dtype().is_numeric() && rhs.len() == 1,
+            InvalidOperation: "can only do arithmetic of arrays of the same type and shape; got {} and {}",
+            lhs.dtype(), rhs.dtype()
+        );
+
+        let l_leaf_array = lhs.get_leaf_array();
+        let r_leaf_array = rhs.get_leaf_array();
+
+        let mut dt = lhs.dtype();
+        let mut shape = vec![ReshapeDimension::Specified(
+            Dimension::new(lhs.len() as u64),
+        )];
+        while let DataType::Array(child, size) = dt {
+            shape.push(ReshapeDimension::Specified(Dimension::new(*size as u64)));
+            dt = child;
+        }
 
         let out = op(l_leaf_array, r_leaf_array)?;
         out.reshape_array(&shape)
     }
+
+    fn arithm_helper_scalar_lhs(
+        &self,
+        lhs: Series,
+        op: &dyn Fn(Series, Series) -> PolarsResult<Series>,
+    ) -> PolarsResult<Series> {
+        polars_ensure!(
+            lhs.len() == 1,
+            InvalidOperation: "can only do arithmetic of between arrays and a scalar the leaf type; got {} and {}",
+            lhs.dtype(), self.dtype()
+        );
+
+        let r_leaf_array = self.get_leaf_array();
+        let out = op(lhs, r_leaf_array)?;
+
+        let mut dt = self.dtype();
+        let mut shape = vec![ReshapeDimension::Specified(Dimension::new(
+            self.len() as u64
+        ))];
+        while let DataType::Array(child, size) = dt {
+            shape.push(ReshapeDimension::Specified(Dimension::new(*size as u64)));
+            dt = child;
+        }
+
+        out.reshape_array(&shape)
+    }
 }
 
 #[cfg(feature = "dtype-array")]
diff --git a/crates/polars-core/src/series/implementations/floats.rs b/crates/polars-core/src/series/implementations/floats.rs
index 846e326d35b2..85c0d87cf0f1 100644
--- a/crates/polars-core/src/series/implementations/floats.rs
+++ b/crates/polars-core/src/series/implementations/floats.rs
@@ -110,48 +110,18 @@ macro_rules! impl_dyn_series {
             }
 
             fn subtract(&self, rhs: &Series) -> PolarsResult<Series> {
-                polars_ensure!(
-                    self.dtype() == rhs.dtype(),
-                    opq = sub,
-                    self.dtype(),
-                    rhs.dtype()
-                );
                 NumOpsDispatch::subtract(&self.0, rhs)
             }
             fn add_to(&self, rhs: &Series) -> PolarsResult<Series> {
-                polars_ensure!(
-                    self.dtype() == rhs.dtype(),
-                    opq = add,
-                    self.dtype(),
-                    rhs.dtype()
-                );
                 NumOpsDispatch::add_to(&self.0, rhs)
             }
             fn multiply(&self, rhs: &Series) -> PolarsResult<Series> {
-                polars_ensure!(
-                    self.dtype() == rhs.dtype(),
-                    opq = mul,
-                    self.dtype(),
-                    rhs.dtype()
-                );
                 NumOpsDispatch::multiply(&self.0, rhs)
             }
             fn divide(&self, rhs: &Series) -> PolarsResult<Series> {
-                polars_ensure!(
-                    self.dtype() == rhs.dtype(),
-                    opq = div,
-                    self.dtype(),
-                    rhs.dtype()
-                );
                 NumOpsDispatch::divide(&self.0, rhs)
             }
             fn remainder(&self, rhs: &Series) -> PolarsResult<Series> {
-                polars_ensure!(
-                    self.dtype() == rhs.dtype(),
-                    opq = rem,
-                    self.dtype(),
-                    rhs.dtype()
-                );
                 NumOpsDispatch::remainder(&self.0, rhs)
             }
             #[cfg(feature = "algorithm_group_by")]
diff --git a/crates/polars-core/src/series/implementations/mod.rs b/crates/polars-core/src/series/implementations/mod.rs
index b2cb97e39b69..9447dc077e15 100644
--- a/crates/polars-core/src/series/implementations/mod.rs
+++ b/crates/polars-core/src/series/implementations/mod.rs
@@ -183,48 +183,18 @@ macro_rules! impl_dyn_series {
             }
 
             fn subtract(&self, rhs: &Series) -> PolarsResult<Series> {
-                polars_ensure!(
-                    self.dtype() == rhs.dtype(),
-                    opq = sub,
-                    self.dtype(),
-                    rhs.dtype()
-                );
                 NumOpsDispatch::subtract(&self.0, rhs)
             }
             fn add_to(&self, rhs: &Series) -> PolarsResult<Series> {
-                polars_ensure!(
-                    self.dtype() == rhs.dtype(),
-                    opq = add,
-                    self.dtype(),
-                    rhs.dtype()
-                );
                 NumOpsDispatch::add_to(&self.0, rhs)
             }
             fn multiply(&self, rhs: &Series) -> PolarsResult<Series> {
-                polars_ensure!(
-                    self.dtype() == rhs.dtype(),
-                    opq = mul,
-                    self.dtype(),
-                    rhs.dtype()
-                );
                 NumOpsDispatch::multiply(&self.0, rhs)
             }
             fn divide(&self, rhs: &Series) -> PolarsResult<Series> {
-                polars_ensure!(
-                    self.dtype() == rhs.dtype(),
-                    opq = div,
-                    self.dtype(),
-                    rhs.dtype()
-                );
                 NumOpsDispatch::divide(&self.0, rhs)
             }
             fn remainder(&self, rhs: &Series) -> PolarsResult<Series> {
-                polars_ensure!(
-                    self.dtype() == rhs.dtype(),
-                    opq = rem,
-                    self.dtype(),
-                    rhs.dtype()
-                );
                 NumOpsDispatch::remainder(&self.0, rhs)
             }
             #[cfg(feature = "algorithm_group_by")]


Test patch:
diff --git a/py-polars/tests/unit/datatypes/test_array.py b/py-polars/tests/unit/datatypes/test_array.py
index b578266b0c6f..df7a5c5cc4ec 100644
--- a/py-polars/tests/unit/datatypes/test_array.py
+++ b/py-polars/tests/unit/datatypes/test_array.py
@@ -383,3 +383,16 @@ def test_zero_width_array(fn: str) -> None:
 
                 df = pl.concat([a.to_frame(), b.to_frame()], how="horizontal")
                 df.select(c=expr_f(pl.col.a, pl.col.b))
+
+
+def test_elementwise_arithmetic_19682() -> None:
+    dt = pl.Array(pl.Int64, (2, 3))
+
+    a = pl.Series("a", [[[1, 2, 3], [4, 5, 6]]], dt)
+    sc = pl.Series("a", [1])
+    zfa = pl.Series("a", [[]], pl.Array(pl.Int64, 0))
+
+    assert_series_equal(a + a, pl.Series("a", [[[2, 4, 6], [8, 10, 12]]], dt))
+    assert_series_equal(a + sc, pl.Series("a", [[[2, 3, 4], [5, 6, 7]]], dt))
+    assert_series_equal(sc + a, pl.Series("a", [[[2, 3, 4], [5, 6, 7]]], dt))
+    assert_series_equal(zfa + zfa, pl.Series("a", [[]], pl.Array(pl.Int64, 0)))