Cast multi-dimensional numpy arrays when inferring IterableDataset features

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))

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