Vectorize interleave_datasets index generation (probabilities + first/all_exhausted)

[saforem2]

What

_interleave_map_style_datasets (used by interleave_datasets for map-style Datasets) builds the output index list in a pure-Python for-loop, one iteration per output row, when probabilities is given:

for source_idx in iter_random_indices():
    if bool_strategy_func(is_exhausted):
        break
    indices.append(current_index[source_idx] + offsets[source_idx])
    current_index[source_idx] += 1
    if current_index[source_idx] >= lengths[source_idx]:
        is_exhausted[source_idx] = True
        current_index[source_idx] = 0

For large interleaves this loop dominates runtime. Interleaving e.g. nvidia/OpenMathInstruct-2 (~14M rows) with two smaller datasets under stopping_strategy="all_exhausted" produces ~93M output rows and the index build alone takes ~90 min — almost entirely Python interpreter overhead, since the RNG is already drawn in batches (rng.choice(..., size=1000)); it is not compute-bound.

Notably, the sibling probabilities is None all_exhausted branch in the same function is already vectorized with numpy (np.mod/offset). This PR brings the probabilities-given first_exhausted and all_exhausted branches to parity.

How

Reproduce the exact same algorithm with numpy:

1. Replay the identical draw sequence — np.random.default_rng(seed) drawn in the same 1000-sized rng.choice(n, p=probabilities) blocks — until the stop condition can be evaluated.
2. Find the stop position from each source's length-th occurrence: min over sources for first_exhausted (stop when any source is exhausted), max for all_exhausted (stop when every source is).
3. Map each source's k-th appearance to (k % length) + offset (the same rolling-window-on-exhaustion behavior), vectorized per source.

all_exhausted_without_replacement keeps the explicit loop — its skip-on-exhaustion semantics make the output length data-dependent, so it isn't a simple function of the draw sequence.

Bit-identical output

Because the RNG is consumed in exactly the same way and the index mapping is the same, output is bit-identical for a fixed seed:

• The existing hardcoded tests test_interleave_datasets_probabilities ([10, 11, 20, 12, 0, 21, 13]) and test_interleave_datasets_probabilities_oversampling_strategy (16-element sequence) pass unchanged.
• 80 randomized (num_datasets, lengths, probabilities, seed) cases across both first_exhausted and all_exhausted match the previous implementation exactly.
• Added test_interleave_datasets_probabilities_is_deterministic_and_balanced (parametrized over strategy + seed) as a regression guard.

Benchmark

3-source mix, ~93M output rows (all_exhausted, probabilities given): ~90 min → ~5 s.

Notes

• No public API or semantics change; same results, same _fingerprint.
• first_exhausted / all_exhausted only; all_exhausted_without_replacement and the probabilities is None paths are untouched.

[Copilot]

Pull request overview

This PR accelerates interleave_datasets for map-style Datasets when probabilities is provided by replacing the per-row Python loop in _interleave_map_style_datasets with a NumPy-based, blockwise RNG replay + vectorized index mapping for stopping_strategy="first_exhausted" and "all_exhausted".

Changes:

• Vectorizes index generation for probabilities-driven interleaving under first_exhausted and all_exhausted, keeping all_exhausted_without_replacement on the explicit loop.
• Adds a regression test to guard determinism and basic proportionality properties under both stopping strategies across multiple seeds.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
src/datasets/arrow_dataset.py	Introduces vectorized RNG replay and index mapping for probabilities-based first_exhausted/all_exhausted interleaving.
tests/test_arrow_dataset.py	Adds a regression test ensuring deterministic outputs/fingerprints and basic probability ordering across strategies/seeds.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[saforem2]

I benchmarked the argsort-based single-pass approach against the current per-source np.flatnonzero, and the per-source version is actually faster in both the few-dataset and many-dataset cases, so I've kept it.

At 93M draws (dummy lengths, index-mapping section only):

case	per-source flatnonzero (current)	argsort + bincount
3 datasets	1.47s	5.22s
50 datasets	7.58s	12.13s

The O(n_datasets * n_draws) scaling is real in complexity terms, but the constant factors favor flatnonzero: it's a cheap fully-vectorized compare-and-pack, whereas the argsort does an O(n_draws log n_draws) sort of a 93M-element array that dominates even at 50 datasets. flatnonzero's boolean temporaries are n_draws bools each (transient, one source at a time), so peak memory stays modest.

flatnonzero would only lose for a very large number of datasets (roughly where n_datasets approaches log(n_draws)); interleave mixes are typically a handful of sources, so this stays on the fast side in practice. Happy to switch to the argsort form if you'd prefer the better asymptotic behavior for pathological n_datasets regardless -- but for the common case the current code is the faster of the two.

[saforem2]

Hi @lhoestq — gentle ping on this one when you have a moment.

It vectorizes the probabilities-given index generation in _interleave_map_style_datasets (the first_exhausted / all_exhausted branches), which currently builds the index list in a pure-Python per-row loop. This brings it to parity with the already-vectorized probabilities=None branch in the same function.

• Bit-identical output for a fixed seed (added/extended tests assert new == old across seeds, lengths, probabilities, and n-datasets).
• ~90 min → ~5 s to build a real 93M-row interleaved mix (tulu-3 + OpenMathInstruct-2 + ultrachat at given weights) — the per-row loop was pure interpreter overhead, not compute.

Looks like CI is awaiting maintainer approval to run (external contributor). Happy to rebase or adjust anything — thanks!