[pavanmudumba]

PandasArrayExtensionArray._concat_same_type can turn whole input batches into individual rows. For example, concatenating dense batches with 2 and 1 rows returns 2 rows instead of 3, with an extra batch dimension inside each row. A single input chunk is also wrapped as one row. Equally sized object-backed chunks take the vstack branch and gain an extra dimension.

Reproduction:

import numpy as np
from datasets.features.features import PandasArrayExtensionArray

a = PandasArrayExtensionArray(np.arange(8, dtype="int32").reshape(2, 2, 2))
b = PandasArrayExtensionArray(np.arange(8, 12, dtype="int32").reshape(1, 2, 2))
result = PandasArrayExtensionArray._concat_same_type([a, b])
print(len(result), result._data.shape)
# Before: 2 (2,)
# After:  3 (3, 2, 2)

Compare each chunk's per-row shape rather than its full batch shape, and concatenate compatible backing arrays along axis 0. For mixed dense/object-backed chunks, allocate one object slot per input row and copy rows individually. This preserves ragged arrays and null rows without retaining the batch grouping.

Added 24 regression/control cases covering unequal and equal batch lengths, a single chunk, empty chunks, int32/float32, Array2D/Array3D, ragged arrays, mixed dense/object storage, null rows, and conversion back to Arrow. Tests exercise the extension-array concatenation hook directly; the separate pandas dtype-comparison issue is tracked in #8375 and #8464.

Validation on Windows / Python 3.12:

• With the original upstream concatenation method substituted in memory: 16 regression cases fail and 8 controls pass.
• python -m pytest tests/features/test_array_xd.py tests/test_formatting.py -q -rs: 155 passed, 29 skipped for optional dependencies or platform requirements. A short --basetemp path avoids Windows cache-lock filename limits.
• Ruff lint and format checks pass for both changed files.
• git diff --check passes.

[pavanmudumba]

Hi! Following up on this fix when you have a chance to review. It preserves row counts and shapes when concatenating pandas ArrayXD batches, with regression coverage for dense, ragged, empty, and null-containing inputs. The reproduction and local test results are in the description. Happy to adjust the implementation or add any further coverage you’d find useful. Thanks for your time!