huggingface / huggingface/datasets
map(batched=True, drop_last_batch=True) returns the input unchanged when the dataset is smaller than batch_size
- Dominant language
- Python
- Stars
- 22k
- Forks
- 3.4k
- Avg merge
- 5d 7h
- Merged PRs (30d)
- 17
Description
### Describe the bug
When `drop_last_batch=True` and the dataset is smaller than `batch_size`, `Dataset.map(batched=True, ...)` returns the **input unchanged** instead of an empty dataset. The mapped function is never called, but the original rows are passed through with the original (unbatched) schema.
The iterable path and `Dataset.iter` both correctly yield nothing, so the map-style and streaming paths disagree.
### Steps to reproduce the bug
```python
from datasets import Dataset
ds = Dataset.from_dict({"a": [1, 2, 3]})
out = ds.map(lambda b: b, batched=True, batch_size=5, drop_last_batch=True)
print(len(out), out.features)
# 3 {'a': Value('int64')} <- expected 0 rows
print(len(ds.batch(5, drop_last_batch=True)))
# 3 <- expected 0 rows
print(sum(1 for _ in ds.to_iterable_dataset().batch(5, drop_last_batch=True)))
# 0 <- correct
print(sum(1 for _ in ds.iter(5, drop_last_batch=True)))
# 0 <- correct
```
Note the progress bar reports `Map: 0 examples` while the result still has 3 rows, which confirms the function was never invoked yet the data was passed through.
`Dataset.batch()` is implemented on top of `map(batched=True)`, so it inherits the same behavior — and its result is doubly wrong, since the rows come back *unbatched* (`Value('int64')` rather than `List(Value('int64'))`).
### Expected behavior
`drop_last_batch=True` means an incomplete final batch is dropped. When the whole dataset is smaller than one batch, every batch is incomplete, so the result should be empty — matching `IterableDataset.batch` and `Dataset.iter`.
### Root cause
In `_map_single`, `update_data` stays `None` because the function is never invoked for an incomplete batch. The writer is therefore never initialized and the code falls through to `yield rank, True, shard`, returning the original shard.
### Open question for maintainers
The right output *schema* is genuinely ambiguous when the function never runs, which is why I'm filing this rather than sending a patch directly. The options I see:
1. Return an empty dataset carrying the **input** features (simple, but wrong for `batch()`, where the output schema should be the batched `List(...)` form).
2. Initialize the writer eagerly from the function's declared/inferred output features.
3. Raise a clear error when `drop_last_batch=True` would drop everything.
Happy to open a PR once there's a preference — I have a fix for (1) working locally, and (2) is doable but a larger change.
### Environment info
- `datasets` 5.0.2.dev0 (`main` @ b7cb10b0e)
- Python 3.12.13
- pyarrow 25.0.0
Contributor guide
Research direction
Start at the _map_single implementation used by Dataset.map and reproduce the smaller-than-batch_size example from the issue. Compare its behavior with Dataset.batch(), to_iterable_dataset().batch(), and Dataset.iter() for drop_last_batch=True. Done means the chosen schema behavior is agreed with maintainers and the map-style result matches the expected empty-batch semantics.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- data
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100