Preserve rows when concatenating pandas ArrayXD batches

src/datasets/features/features.py

@@ -982,14 +982,18 @@ def _from_sequence(
 
     @classmethod
     def _concat_same_type(cls, to_concat: Sequence_["PandasArrayExtensionArray"]) -> "PandasArrayExtensionArray":
-        if len(to_concat) > 1 and all(
-            va._data.shape == to_concat[0]._data.shape and va._data.dtype == to_concat[0]._data.dtype
+        if to_concat and all(
+            va._data.shape[1:] == to_concat[0]._data.shape[1:] and va._data.dtype == to_concat[0]._data.dtype
             for va in to_concat
         ):
-            data = np.vstack([va._data for va in to_concat])
+            data = np.concatenate([va._data for va in to_concat], axis=0)
         else:
-            data = np.empty(len(to_concat), dtype=object)
-            data[:] = [va._data for va in to_concat]
+            data = np.empty(sum(len(va) for va in to_concat), dtype=object)
+            offset = 0
+            for va in to_concat:
+                for row in va._data:
+                    data[offset] = row
+                    offset += 1
         return cls(data, copy=False)
 
     @property

tests/features/test_array_xd.py

@@ -15,6 +15,7 @@
 from datasets.features.features import (
     Array2DExtensionType,
     Array3DExtensionType,
+    PandasArrayExtensionArray,
     PandasArrayExtensionDtype,
     _ArrayXD,
 )
@@ -36,6 +37,57 @@
 )
 
 
+@pytest.mark.parametrize("shape", [(2, 2), (2, 1, 3)])
+@pytest.mark.parametrize("dtype", ["int32", "float32"])
+@pytest.mark.parametrize("chunk_lengths", [(2, 1), (3,), (2, 2), (0, 2), (0, 0)])
+def test_pandas_array_concat_dense(shape, dtype, chunk_lengths):
+    values = np.arange(sum(chunk_lengths) * np.prod(shape), dtype=dtype).reshape((-1, *shape))
+    chunks = []
+    offset = 0
+    for length in chunk_lengths:
+        chunks.append(PandasArrayExtensionArray(values[offset : offset + length]))
+        offset += length
+
+    result = PandasArrayExtensionArray._concat_same_type(chunks)
+
+    assert len(result) == sum(chunk_lengths)
+    assert result._data.shape == values.shape
+    assert result._data.dtype == values.dtype
+    np.testing.assert_array_equal(result._data, values)
+    extension_type = (Array2DExtensionType if len(shape) == 2 else Array3DExtensionType)(shape, dtype)
+    assert pa.array(result, type=extension_type).to_pylist() == values.tolist()
+
+
+@pytest.mark.parametrize("mixed_dense", [False, True])
+@pytest.mark.parametrize("first_length", [1, 2])
+def test_pandas_array_concat_ragged(mixed_dense, first_length):
+    values = [
+        np.arange(4, dtype="float32").reshape(2, 2),
+        np.ones((2, 2), dtype="float32"),
+        np.ones((1, 2), dtype="float32"),
+        np.nan,
+    ]
+    first = np.stack(values[:first_length]) if mixed_dense else np.empty(first_length, dtype=object)
+    if not mixed_dense:
+        first[:] = values[:first_length]
+    second = np.empty(len(values) - first_length, dtype=object)
+    second[:] = values[first_length:]
+
+    result = PandasArrayExtensionArray._concat_same_type(
+        [PandasArrayExtensionArray(first), PandasArrayExtensionArray(second)]
+    )
+
+    assert len(result) == len(values)
+    for actual, expected in zip(result._data, values):
+        np.testing.assert_array_equal(actual, expected)
+    assert pa.array(result, type=Array2DExtensionType((None, 2), "float32")).to_pylist() == [
+        values[0].tolist(),
+        values[1].tolist(),
+        values[2].tolist(),
+        None,
+    ]
+
+
 def generate_examples(features: dict, num_examples=100, seq_shapes=None):
     dummy_data = []
     seq_shapes = seq_shapes or {}