Probabilistic interleave without replacement samples empty datasets
Describe the bug
With probabilities and stopping_strategy="all_exhausted_without_replacement", an empty map-style input can add a duplicate row from another source or cause an out-of-range index. This can happen when filtering one of the datasets before mixing them.
Steps to reproduce
from datasets import Dataset, interleave_datasets
left = Dataset.from_dict({"a": [0, 1]})
right = Dataset.from_dict({"a": [10]})
empty = left.select([])
mixed = interleave_datasets(
[empty, left, right],
probabilities=[1 / 3] * 3,
seed=42,
stopping_strategy="all_exhausted_without_replacement",
)
print(sorted(mixed["a"]))
Actual on current main:
[0, 0, 1, 10]
Expected:
[0, 1, 10]
Moving the empty source to the end, [left, right, empty], instead raises:
IndexError: Index 3 out of range for dataset of size 3.
The without-replacement strategy promises to include each input sample exactly once. An empty source has no sample to contribute.
Investigation
The probability-sampling loop initializes every entry of is_exhausted to False. When an empty source is drawn, its offset is appended before the length check marks it exhausted. That offset either points into the next source or lies beyond the concatenated dataset.
Initializing empty inputs as exhausted for this strategy lets the existing skip-exhausted-source check handle them.
Related work was checked: #8399 addresses all_exhausted without probabilities. #8318 vectorizes the probability-based first_exhausted / all_exhausted modes and keeps this without-replacement loop, including its all-False initialization. This report is scoped to the remaining all_exhausted_without_replacement path.
Environment
datasets 5.0.2.dev0, currentmain- Python 3.13.15, macOS
- NumPy 2.5.3, PyArrow 25.0.1
- All inputs are local; no Hub access or model downloads
A small fix and regression tests are ready: the nine empty-source cases fail before the fix; all fourteen focused cases pass afterward, including non-empty and all-empty controls.
Prepared with OpenAI Codex; reproduction and validation were run locally.
Source: huggingface/datasets