Fix Dataset.map(batched=True, drop_last_batch=True) returning unchanged input when dataset is smaller than batch_size

[mayuriphad]

Fixes #8386

Bug

When drop_last_batch=True and the dataset has fewer rows than batch_size, Dataset.map(batched=True, ...) was returning the input dataset unchanged instead of an empty dataset. Every batch is incomplete in this situation, so the mapped function is never invoked and update_data stays falsy; _map_single then fell through to yield rank, True, shard, silently passing the original, unprocessed rows straight through.

This disagreed with IterableDataset.batch() and Dataset.iter(), which both correctly yield zero rows when every batch is incomplete, and it also affected Dataset.batch() since it is implemented on top of map(batched=True).

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))  # 3 before the fix, expected 0

print(len(ds.batch(5, drop_last_batch=True)))  # 3 before the fix, expected 0

Fix

In Dataset._map_single, when batched=True, drop_last_batch=True, and the shard has fewer rows than batch_size (so the function was never called and update_data is still falsy), return an empty dataset (shard.select([])) carrying the input schema instead of the untouched shard. This matches option (1) discussed in the issue, and aligns the eager map/batch behavior with the streaming IterableDataset.batch/Dataset.iter behavior.

Test

Added test_map_batched_drop_last_batch_smaller_than_batch_size in tests/test_arrow_dataset.py, which fails on main (AssertionError: 30 != 0) and passes with this fix.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.