Skip to content

Cast multi-dimensional numpy arrays when inferring IterableDataset features#8319

Open
vineethsaivs wants to merge 1 commit into
huggingface:mainfrom
vineethsaivs:fix/iterable-features-multidim-numpy
Open

Cast multi-dimensional numpy arrays when inferring IterableDataset features#8319
vineethsaivs wants to merge 1 commit into
huggingface:mainfrom
vineethsaivs:fix/iterable-features-multidim-numpy

Conversation

@vineethsaivs

Copy link
Copy Markdown

Summary

IterableDataset._resolve_features() crashes when a map function yields multi-dimensional numpy arrays:

from datasets import Dataset
import numpy as np
ds = Dataset.from_dict({"a": [[[1, 2, 3], [1, 2, 3]]]}).to_iterable_dataset().map(lambda x: {"a": [np.array(x["a"])]})
ds._resolve_features()
# pyarrow.lib.ArrowInvalid: Can only convert 1-dimensional array values

Feature inference goes _resolve_features() -> _infer_features_from_batch(...) -> pa.Table.from_pydict(batch), and PyArrow's from_pydict only accepts 1-dimensional array values, so any 2D+ numpy array in a column crashes it.

Every other place in iterable_dataset.py that builds an Arrow table from raw examples first routes the data through cast_to_python_objects(..., only_1d_for_numpy=True) (lines 173 and 184), which recursively turns multi-dimensional numpy/tensor values into nested lists of 1-D arrays that PyArrow accepts. The non-streaming Dataset/ArrowWriter path does the same. _infer_features_from_batch was simply missing that cast.

Fix

Route the batch through cast_to_python_objects(batch, only_1d_for_numpy=True) before from_pydict. The helper is already imported in this module, and it is a no-op for ordinary Python data, so plain batches infer identically.

Fixes #7100.

Test

tests/test_iterable_dataset.py::test_iterable_dataset_resolve_features_from_multidim_numpy resolves the features of an IterableDataset whose map yields multi-dimensional numpy arrays. It raises ArrowInvalid before the fix and passes after (resolving to a nested List(...)).

…atures

`_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 huggingface#7100. Adds a regression test asserting features resolve from a map that
yields multi-dimensional numpy arrays.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

IterableDataset: cannot resolve features from list of numpy arrays

1 participant