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