huggingface / huggingface/datasets
batch(0) silently returns one batch with the whole dataset instead of raising
- Dominant language
- Python
- Stars
- 22k
- Forks
- 3.4k
- Avg merge
- 5d 7h
- Merged PRs (30d)
- 17
Description
### 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
```python
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:
- #8446 adds `_check_batch_size` but wires it only into the **export** methods (`to_csv`, `to_json`, ...), not `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:
```python
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
Contributor guide
Research direction
Start at Dataset.batch and IterableDataset.batch, then inspect the existing _check_batch_size helper from #8446. Reproduce the issue with batch sizes 0, -1, and 2; done means both paths reject non-positive sizes with the specified ValueError while the positive case remains unchanged.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- data
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100