huggingface / huggingface/datasets

map(batched=True, drop_last_batch=True) returns the input unchanged when the dataset is smaller than batch_size

Open
#8,386 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.