huggingface / huggingface/datasets

Probabilistic interleave without replacement samples empty datasets

Open
#8,626 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

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

```python
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`:

```text
[0, 0, 1, 10]
```

Expected:

```text
[0, 1, 10]
```

Moving the empty source to the end, `[left, right, empty]`, instead raises:

```text
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`, current `main`
- 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.

Contributor guide

Open the contributing guide

Research direction

Start by running the provided interleave_datasets reproduction, then locate the probability-sampling loop for the all_exhausted_without_replacement strategy. Add regression coverage for the nine empty-source cases and verify the fourteen focused cases, including non-empty and all-empty controls, pass with the expected samples and no IndexError.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
74/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.