fix: Raise error when casting Array to different width (#14995)

pola-rs · Mar 15, 2024 · 4f069f5 · 4f069f5
1 parent 71a6563
commit 4f069f5
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 5 deletions.
diff --git a/crates/polars-core/src/chunked_array/cast.rs b/crates/polars-core/src/chunked_array/cast.rs
@@ -379,11 +379,12 @@ impl ChunkCast for ListChunked {
         match data_type {
             List(child_type) => {
                 match (self.inner_dtype(), &**child_type) {
+                    (old, new) if old == *new => Ok(self.clone().into_series()),
                     #[cfg(feature = "dtype-categorical")]
                     (dt, Categorical(None, _) | Enum(_, _))
                         if !matches!(dt, Categorical(_, _) | Enum(_, _) | String | Null) =>
                     {
-                        polars_bail!(ComputeError: "cannot cast List inner type: '{:?}' to Categorical", dt)
+                        polars_bail!(InvalidOperation: "cannot cast List inner type: '{:?}' to Categorical", dt)
                     },
                     _ => {
                         // ensure the inner logical type bubbles up
@@ -417,7 +418,7 @@ impl ChunkCast for ListChunked {
             },
             _ => {
                 polars_bail!(
-                    ComputeError: "cannot cast List type (inner: '{:?}', to: '{:?}')",
+                    InvalidOperation: "cannot cast List type (inner: '{:?}', to: '{:?}')",
                     self.inner_dtype(),
                     data_type,
                 )
@@ -442,10 +443,16 @@ impl ChunkCast for ArrayChunked {
         use DataType::*;
         match data_type {
             Array(child_type, width) => {
+                polars_ensure!(
+                    *width == self.width(),
+                    InvalidOperation: "cannot cast Array to a different width"
+                );
+
                 match (self.inner_dtype(), &**child_type) {
+                    (old, new) if old == *new => Ok(self.clone().into_series()),
                     #[cfg(feature = "dtype-categorical")]
                     (dt, Categorical(None, _) | Enum(_, _)) if !matches!(dt, String) => {
-                        polars_bail!(ComputeError: "cannot cast fixed-size-list inner type: '{:?}' to dtype: {:?}", dt, child_type)
+                        polars_bail!(InvalidOperation: "cannot cast Array inner type: '{:?}' to dtype: {:?}", dt, child_type)
                     },
                     _ => {
                         // ensure the inner logical type bubbles up
@@ -476,7 +483,13 @@ impl ChunkCast for ArrayChunked {
                     ))
                 }
             },
-            _ => polars_bail!(ComputeError: "cannot cast list type"),
+            _ => {
+                polars_bail!(
+                    InvalidOperation: "cannot cast Array type (inner: '{:?}', to: '{:?}')",
+                    self.inner_dtype(),
+                    data_type,
+                )
+            },
         }
     }
 

diff --git a/py-polars/tests/unit/operations/test_cast.py b/py-polars/tests/unit/operations/test_cast.py
@@ -632,3 +632,11 @@ def test_cast_decimal() -> None:
         df.select(pl.col("s").cast(pl.Boolean)),
         pl.DataFrame({"s": [False, True, True]}),
     )
+
+
+def test_cast_array_to_different_width() -> None:
+    s = pl.Series([[1, 2], [3, 4]], dtype=pl.Array(pl.Int8, 2))
+    with pytest.raises(
+        pl.InvalidOperationError, match="cannot cast Array to a different width"
+    ):
+        s.cast(pl.Array(pl.Int16, 3))
diff --git a/py-polars/tests/unit/test_errors.py b/py-polars/tests/unit/test_errors.py
@@ -512,7 +512,8 @@ def test_err_on_invalid_time_zone_cast() -> None:
 def test_invalid_inner_type_cast_list() -> None:
     s = pl.Series([[-1, 1]])
     with pytest.raises(
-        pl.ComputeError, match=r"cannot cast List inner type: 'Int64' to Categorical"
+        pl.InvalidOperationError,
+        match=r"cannot cast List inner type: 'Int64' to Categorical",
     ):
         s.cast(pl.List(pl.Categorical))