From 3592ee10e4a001df7233d300305f6fbd604eeb06 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Fri, 10 Jul 2026 15:35:17 -0700 Subject: [PATCH] Cast multi-dimensional numpy arrays when inferring IterableDataset features `_infer_features_from_batch` built the Arrow table with `pa.Table.from_pydict(batch)`, which only accepts 1-dimensional array values. When a map function yields multi-dimensional numpy arrays (e.g. an image or any 2D+ array), `IterableDataset._resolve_features()` therefore crashed with `pyarrow.lib.ArrowInvalid: Can only convert 1-dimensional array values`. Every other Arrow-table builder in this file first routes the data through `cast_to_python_objects(..., only_1d_for_numpy=True)` (see lines 173 and 184), which recursively turns multi-dimensional numpy arrays into nested lists of 1-D arrays that PyArrow accepts; the non-streaming `Dataset`/`ArrowWriter` path does the same. Route the batch through the same helper (already imported in this module) so feature inference matches the rest of the codebase. Fixes #7100. Adds a regression test asserting features resolve from a map that yields multi-dimensional numpy arrays. --- src/datasets/iterable_dataset.py | 2 +- tests/test_iterable_dataset.py | 16 ++++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/src/datasets/iterable_dataset.py b/src/datasets/iterable_dataset.py index 17b9a2020fc..c5fdcad4371 100644 --- a/src/datasets/iterable_dataset.py +++ b/src/datasets/iterable_dataset.py @@ -127,7 +127,7 @@ def add_column_fn(example: dict, idx: int, name: str, column: list[dict]): def _infer_features_from_batch(batch: dict[str, list], try_features: Optional[Features] = None) -> Features: - pa_table = pa.Table.from_pydict(batch) + pa_table = pa.Table.from_pydict(cast_to_python_objects(batch, only_1d_for_numpy=True)) if try_features is not None: try: pa_table = table_cast(pa_table, pa.schema(try_features.type)) diff --git a/tests/test_iterable_dataset.py b/tests/test_iterable_dataset.py index 3c9a0ee32a6..2aef77c3cbf 100644 --- a/tests/test_iterable_dataset.py +++ b/tests/test_iterable_dataset.py @@ -3182,3 +3182,19 @@ def gen(): texts = ds["text"] assert isinstance(texts, IterableColumn) assert list(texts) == [["Good", "Bad"], ["Good again", "Bad again"]] + + +def test_iterable_dataset_resolve_features_from_multidim_numpy(): + # Regression for https://github.com/huggingface/datasets/issues/7100: resolving + # the features of an IterableDataset whose map yields multi-dimensional numpy + # arrays must not raise "Can only convert 1-dimensional array values". The batch + # is routed through cast_to_python_objects(..., only_1d_for_numpy=True) like the + # other Arrow-table builders in iterable_dataset.py. + ds = ( + Dataset.from_dict({"a": [[[1, 2, 3], [1, 2, 3]]]}) + .to_iterable_dataset() + .map(lambda x: {"a": [np.array(x["a"])]}) + ) + ds = ds._resolve_features() + assert ds.features is not None + assert "a" in ds.features