batch(0) silently returns one batch with the whole dataset instead of raising

[shashvat-singham]

Describe the bug

Dataset.batch() and IterableDataset.batch() accept batch_size=0 and negative sizes, and quietly return one batch holding the whole dataset. Asking for a zero-sized batch produces the largest possible batch, which is as close to the opposite of the request as the API can get.

Steps or code to reproduce the bug

from datasets import Dataset

ds = Dataset.from_dict({"a": list(range(5))})
it = ds.to_iterable_dataset()

[len(b["a"]) for b in it.batch(0)]    # [5]
[len(b["a"]) for b in it.batch(-1)]   # [5]
[len(b["a"]) for b in it.batch(2)]    # [2, 2, 1]   <- sane case

len(ds.batch(0))    # 1
len(ds.batch(-1))   # 1
len(ds.batch(2))    # 3

Both the eager and the streaming path behave the same way, so this is consistent — just consistently wrong.

Expected behavior

batch_size should be required to be a positive integer, raising something like

ValueError: batch_size must be a positive integer, but got 0.

Silently substituting "everything in one batch" is the worst of the options: it looks like it worked, and a batch_size computed at runtime that comes out as 0 (an empty shard, an integer division, a config default that did not get filled in) turns into a single unbounded batch. On a large dataset that is a memory blow-up rather than an error.

Note on overlap with open PRs

I checked the open PRs before filing, and this specific path is not covered by either of the two nearby ones:

• Reject a non-positive batch_size in the Dataset export methods #8446 adds _check_batch_size but wires it only into the export methods (to_csv, to_json, ...), not batch().
• Validate by_column and fix the error message in Dataset.batch #8443 validates by_column and the "missing batch_size or by_column" message in Dataset.batch, but does not check the value of batch_size.

So the helper that batch() wants already exists in #8446:

def _check_batch_size(batch_size: Optional[int]):
    if batch_size is not None and batch_size <= 0:
        raise ValueError(f"batch_size must be a positive integer, but got {batch_size}.")

and the fix is largely calling it from Dataset.batch and IterableDataset.batch too. That does mean this is best done after #8446 lands, to avoid two copies of the same helper.

Related, lower confidence

Same family, listed separately because the right answer is less obvious and I am not proposing a change: IterableDataset.take(-1) yields 0 rows and skip(-1) yields all rows, i.e. negatives are silently clamped rather than rejected. That is defensible as "clamp like a slice", unlike batch(0), where there is no reading under which one giant batch is the answer.

Environment info

• datasets version: 5.0.2.dev0 (main @ 48b7ee7)
• Python version: 3.11.9
• Platform: Windows 11
• PyArrow version: 25.0.1

[Anakintano]

Hi @shashvat-singham i did also try and reproduced it and got the same results described :

IterableDataset.batch(0/-1) → [5]
Dataset.batch(0/-1)         → 1 batch
batch(2)                    → [2, 2, 1] / 3 batches

I’d like to give this a try

Thanks

#self-assign