fix: Allow structured conversion to/from numpy with Array types, preserving shape (#23438)

Author: orlp

py-polars/polars/_utils/construction/dataframe.py

@@ -1213,9 +1213,6 @@ def numpy_to_pydf(
         n_columns = len(record_names)
         for nm in record_names:
             shape = data[nm].shape
-            if len(data[nm].shape) > 2:
-                msg = f"cannot create DataFrame from structured array with elements > 2D; shape[{nm!r}] = {shape}"
-                raise ValueError(msg)
         if not schema:
             schema = record_names
     else:

py-polars/polars/dataframe/frame.py

@@ -2010,7 +2010,7 @@ def to_numpy(
                 if s.dtype == String and not s.has_nulls():
                     arr = arr.astype(str, copy=False)
                 arrays.append(arr)
-                struct_dtype.append((s.name, arr.dtype))
+                struct_dtype.append((s.name, arr.dtype, arr.shape[1:]))
 
             out = np.empty(self.height, dtype=struct_dtype)
             for idx, c in enumerate(self.columns):

py-polars/tests/unit/interop/numpy/test_to_numpy_df.py

@@ -287,3 +287,20 @@ def test_to_numpy_c_order_1700() -> None:
         df_chunked,
         pl.from_numpy(df_chunked.to_numpy(order="c"), schema=df_chunked.schema),
     )
+
+
+def test_to_numpy_array_shape_23426() -> None:
+    df = pl.DataFrame(
+        {
+            "x": [1, 2],
+            "y": [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
+            "z": [[[-1, -1, -2], [4, 5, 6]], [[-3, -5, -8], [10, 20, 30]]],
+        },
+        schema={
+            "x": pl.UInt8,
+            "y": pl.Array(pl.Float32, 3),
+            "z": pl.Array(pl.Int16, (2, 3)),
+        },
+    )
+
+    assert_frame_equal(df, pl.from_numpy(df.to_numpy(structured=True)))