`.batch()` error on formatted datasets

[mariosasko]

The .batch() method currently assumes the input (batch) is always a dictionary, which causes errors when it isn't. This can happen with formatted datasets, since formats like "pyarrow", "pandas" (only affects IterableDataset), and "polars" return tables/dataframes instead of dictionaries.

For example:

from datasets import IterableDataset, Dataset
list(IterableDataset.from_dict({"a": [1, 2, 3, 4]}).with_format("pyarrow").batch(2))
# AttributeError: 'pyarrow.lib.Table' object has no attribute 'items'

Ideally, the result should be the same whether the format is applied before or after batching, i.e., the following should hold for all the format types:

assert list(IterableDataset.from_dict({"a": [1, 2, 3, 4]}).with_format(format_type).batch(2)) == list(IterableDataset.from_dict({"a": [1, 2, 3, 4]}).batch(2).with_format(format_type))
assert list(Dataset.from_dict({"a": [1, 2, 3, 4]}).with_format(format_type).batch(2)) == list(Dataset.from_dict({"a": [1, 2, 3, 4]}).batch(2).with_format(format_type))

[lhoestq]

Hi ! Good catch :) Since table-formatted iterable datasets have an Arrow path (i.e. they have .iter_arrow()) I guess a simple fix would be to make .batch() make the dataset arrow-formatted and use a batch_arrow_fn in that case (instead of the current _batch_fn that expects dictionaries). It would also be better performance-wise. How does that sound ?

[mariosasko]

Yes, and a similar approach can be used for the Dataset case: temporarily switching to the "arrow" format, running the batched .map() (which should also be much faster since it avoids Python ↔ Arrow conversions), and then restoring the original format.