From 62bff919993a41d4debdce9ce012225644c45c6f Mon Sep 17 00:00:00 2001 From: Sam Foreman Date: Thu, 9 Jul 2026 23:41:36 -0500 Subject: [PATCH 1/2] Vectorize interleave_datasets index generation (probabilities + first/all_exhausted) `_interleave_map_style_datasets` builds the output index list in a pure-Python for-loop (one iteration per output row) when `probabilities` is given. For large interleaves this dominates runtime -- e.g. interleaving NVIDIA OpenMathInstruct-2 (~14M rows) with `all_exhausted` produces ~93M rows and takes ~90 min, almost all of it in that loop (the RNG is already batched; it is Python interpreter overhead, not compute). The sibling `probabilities is None` `all_exhausted` branch is already vectorized with numpy (modulo/offset). This brings the probabilities-given `first_exhausted` and `all_exhausted` branches to parity: replay the same 1000-sized `rng.choice(..., p=probabilities)` draw blocks, find the stop position from each source's length-th occurrence (min for first_exhausted, max for all_exhausted), and map each source's k-th appearance to `(k % length) + offset` with numpy. Output is bit-identical for a fixed `seed` (same RNG consumption + same rolling-window mapping): the existing hardcoded tests `test_interleave_datasets_probabilities` and `..._probabilities_oversampling_strategy` pass unchanged, and 80 randomized (lengths, probabilities, seed) cases across both strategies match the previous implementation exactly. `all_exhausted_without_replacement` keeps the explicit loop (its skip-on-exhaustion semantics make the output length data-dependent). Benchmark (3-source mix, ~93M output rows): ~90 min -> ~5 s. Adds a randomized determinism/balance test for the probabilities-given paths. --- src/datasets/arrow_dataset.py | 79 +++++++++++++++++++++++++++-------- tests/test_arrow_dataset.py | 27 ++++++++++++ 2 files changed, 88 insertions(+), 18 deletions(-) diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py index 59451a640e6..d3344bbf360 100644 --- a/src/datasets/arrow_dataset.py +++ b/src/datasets/arrow_dataset.py @@ -7210,16 +7210,64 @@ def _interleave_map_style_datasets( # We have to keep the indices to their respective dataset offsets and to flatten to effectively interleave the datasets indices = (indices + offsets).flatten().tolist() + elif stopping_strategy in ("first_exhausted", "all_exhausted"): + # Vectorized equivalent of the per-example Python loop below (kept for + # `all_exhausted_without_replacement`). For `first_exhausted` / + # `all_exhausted` the loop just appends, for each randomly drawn source, + # that source's next index with wrap-around on exhaustion (rolling + # window), and stops when the first (any) / every (all) source has been + # drawn at least `length` times. That is fully vectorizable, giving a + # large speedup for big outputs (e.g. ~90 min -> ~5 s for a ~93M-row + # interleave) while producing bit-identical output for a fixed `seed`, + # since we consume the RNG in the exact same 1000-sized `choice` blocks. + lengths_arr = np.asarray(lengths, dtype=np.int64) + n_datasets = len(lengths) + rng = np.random.default_rng(seed) + + # Draw source indices in 1000-sized blocks (matching the original + # iter_random_indices) until the stopping condition can be evaluated, + # i.e. until enough sources have reached their length. + blocks = [] + counts = np.zeros(n_datasets, dtype=np.int64) + reached = np.zeros(n_datasets, dtype=bool) + while not (reached.any() if not oversampling else reached.all()): + block = rng.choice(n_datasets, size=1000, p=probabilities) + blocks.append(block) + counts += np.bincount(block, minlength=n_datasets) + reached = counts >= lengths_arr + draws = np.concatenate(blocks) + + # For each source, the draw-position of its `length`-th occurrence (when + # it becomes exhausted). The original loop checks the stop condition + # BEFORE appending, so the last kept draw is at: + # - first_exhausted: the earliest such position (min over sources that + # reached length) minus... no: stop fires as soon as ANY source is + # exhausted, which happens right AFTER that source's length-th draw. + # So we keep draws up to and including that length-th occurrence. + # - all_exhausted: the latest such position (max), inclusive. + exhaust_pos = np.full(n_datasets, -1, dtype=np.int64) + for s in range(n_datasets): + occ = np.flatnonzero(draws == s) + if len(occ) >= lengths[s]: + exhaust_pos[s] = occ[lengths[s] - 1] + valid_pos = exhaust_pos[exhaust_pos >= 0] + stop_at = valid_pos.min() if not oversampling else valid_pos.max() + used = draws[: stop_at + 1] + + # Map each source's k-th appearance to its k-th index with wrap-around: + # concatenated row = (k % length) + offset. + indices_arr = np.empty(len(used), dtype=np.int64) + for s in range(n_datasets): + pos = np.flatnonzero(used == s) + indices_arr[pos] = (np.arange(len(pos)) % lengths[s]) + offsets[s] + indices = indices_arr.tolist() + else: - # boolean array indicating if at index i if the dataset_i has been fully exhausted + # all_exhausted_without_replacement: each sample appears exactly once, so + # an exhausted source is skipped rather than wrapped -- the output length + # is not simply the draw count, so keep the explicit per-example loop. is_exhausted = np.full(len(lengths), False) - # if undersampling ("first_exhausted"), we stop as soon as one dataset is exhausted - # if oversampling ("all_exhausted"), we stop as soons as every dataset is exhausted, i.e as soon as every samples of every dataset has been visited at least once - bool_strategy_func = ( - np.all if (oversampling or stopping_strategy == "all_exhausted_without_replacement") else np.any - ) - def iter_random_indices(): """Get an infinite iterator that randomly samples the index of the source to pick examples from.""" rng = np.random.default_rng(seed) @@ -7229,24 +7277,19 @@ def iter_random_indices(): current_index = [0] * len(datasets) indices = [] for source_idx in iter_random_indices(): - # If no oversampling, we stop as soon as a dataset has ran out of examples (np.any) - # Otherwise, we stop as soon as every dataset has ran out of examples (np.all) - if bool_strategy_func(is_exhausted): - # the stopping condition was reached, let's stop + # Stop as soon as every dataset has run out of examples (np.all). + if np.all(is_exhausted): break - # let's add the example at the current index of the `source_idx`-th dataset - # For without replacement sampling we additionally need to make sure the current source is not exhausted to not oversample. - if stopping_strategy != "all_exhausted_without_replacement" or not is_exhausted[source_idx]: + # Add the example at the current index of the source_idx-th dataset, + # unless that source is already exhausted (no oversampling). + if not is_exhausted[source_idx]: indices.append(current_index[source_idx] + offsets[source_idx]) current_index[source_idx] += 1 - # we've ran out of examples for the current dataset, let's update our boolean array and bring the current_index back to 0 + # Mark exhaustion; do not reset the index (no replacement). if current_index[source_idx] >= lengths[source_idx]: is_exhausted[source_idx] = True - # We don't want to reset the iterator when stopping strategy is without replacement. - if stopping_strategy != "all_exhausted_without_replacement": - current_index[source_idx] = 0 return concatenated_datasets.select(indices, **kwargs) diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py index 18a8a6038fe..498041f3b21 100644 --- a/tests/test_arrow_dataset.py +++ b/tests/test_arrow_dataset.py @@ -3809,6 +3809,33 @@ def test_interleave_datasets_probabilities_oversampling_strategy(): ) +@pytest.mark.parametrize("stopping_strategy", ["first_exhausted", "all_exhausted"]) +@pytest.mark.parametrize("seed", [0, 42, 1234]) +def test_interleave_datasets_probabilities_is_deterministic_and_balanced(stopping_strategy, seed): + # Regression guard for the vectorized index generation in + # _interleave_map_style_datasets (probabilities-given first/all_exhausted): + # it must stay deterministic for a fixed seed and respect the requested + # sampling proportions. Uses larger, uneven sources so the result is not + # trivially short. + probabilities = [0.6, 0.3, 0.1] + d1 = Dataset.from_dict({"a": list(range(0, 500))}) + d2 = Dataset.from_dict({"a": list(range(1000, 1200))}) + d3 = Dataset.from_dict({"a": list(range(2000, 2050))}) + kwargs = dict(probabilities=probabilities, seed=seed, stopping_strategy=stopping_strategy) + ds_a = interleave_datasets([d1, d2, d3], **kwargs) + ds_b = interleave_datasets([d1, d2, d3], **kwargs) + # deterministic: identical values and fingerprint across calls + assert ds_a["a"] == ds_b["a"] + assert ds_a._fingerprint == ds_b._fingerprint + # every yielded value comes from one of the sources + allowed = set(d1["a"]) | set(d2["a"]) | set(d3["a"]) + assert set(ds_a["a"]) <= allowed + # source 0 (prob 0.6) is drawn more than source 2 (prob 0.1) + from collections import Counter + src = Counter("d1" if v < 1000 else ("d2" if v < 2000 else "d3") for v in ds_a["a"]) + assert src["d1"] > src["d3"] + + @pytest.mark.parametrize("batch_size", [4, 5]) @pytest.mark.parametrize("drop_last_batch", [False, True]) def test_dataset_iter_batch(batch_size, drop_last_batch): From 71b03df887923aef30bfd4cd6727e786d16169ac Mon Sep 17 00:00:00 2001 From: Sam Foreman Date: Fri, 10 Jul 2026 00:04:19 -0500 Subject: [PATCH 2/2] Address review: empty-source handling + comment cleanup - Empty source (length 0): the previous vectorized code crashed on np.concatenate([]) (blocks never populated), and stock crashed with a cryptic `IndexError: Index N out of range`. Now raise a clear ValueError naming the empty dataset indices, for both first_exhausted and all_exhausted (an empty source is degenerate either way; silently dropping it would change results). Added a parametrized test. - Tightened the stop-position comment (removed the in-line "minus... no:" thought process) to a clear final statement per strategy. Re the suggestion to replace the per-source np.flatnonzero grouping with an argsort-based single pass: benchmarked both at 93M draws -- flatnonzero is actually faster (3 datasets: 1.5s vs 5.2s; 50 datasets: 7.6s vs 12.1s), since the O(n log n) sort dominates while the per-source vectorized compare stays cheap well past 50 datasets. Keeping flatnonzero; will note this on the thread. Equivalence unchanged: 80/80 randomized cases + the existing hardcoded tests still match the previous implementation bit-for-bit. --- src/datasets/arrow_dataset.py | 89 ++++++++++++++++++++--------------- tests/test_arrow_dataset.py | 14 ++++++ 2 files changed, 64 insertions(+), 39 deletions(-) diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py index d3344bbf360..a8bb5692f0e 100644 --- a/src/datasets/arrow_dataset.py +++ b/src/datasets/arrow_dataset.py @@ -7222,45 +7222,56 @@ def _interleave_map_style_datasets( # since we consume the RNG in the exact same 1000-sized `choice` blocks. lengths_arr = np.asarray(lengths, dtype=np.int64) n_datasets = len(lengths) - rng = np.random.default_rng(seed) - - # Draw source indices in 1000-sized blocks (matching the original - # iter_random_indices) until the stopping condition can be evaluated, - # i.e. until enough sources have reached their length. - blocks = [] - counts = np.zeros(n_datasets, dtype=np.int64) - reached = np.zeros(n_datasets, dtype=bool) - while not (reached.any() if not oversampling else reached.all()): - block = rng.choice(n_datasets, size=1000, p=probabilities) - blocks.append(block) - counts += np.bincount(block, minlength=n_datasets) - reached = counts >= lengths_arr - draws = np.concatenate(blocks) - - # For each source, the draw-position of its `length`-th occurrence (when - # it becomes exhausted). The original loop checks the stop condition - # BEFORE appending, so the last kept draw is at: - # - first_exhausted: the earliest such position (min over sources that - # reached length) minus... no: stop fires as soon as ANY source is - # exhausted, which happens right AFTER that source's length-th draw. - # So we keep draws up to and including that length-th occurrence. - # - all_exhausted: the latest such position (max), inclusive. - exhaust_pos = np.full(n_datasets, -1, dtype=np.int64) - for s in range(n_datasets): - occ = np.flatnonzero(draws == s) - if len(occ) >= lengths[s]: - exhaust_pos[s] = occ[lengths[s] - 1] - valid_pos = exhaust_pos[exhaust_pos >= 0] - stop_at = valid_pos.min() if not oversampling else valid_pos.max() - used = draws[: stop_at + 1] - - # Map each source's k-th appearance to its k-th index with wrap-around: - # concatenated row = (k % length) + offset. - indices_arr = np.empty(len(used), dtype=np.int64) - for s in range(n_datasets): - pos = np.flatnonzero(used == s) - indices_arr[pos] = (np.arange(len(pos)) % lengths[s]) + offsets[s] - indices = indices_arr.tolist() + + # Empty sources: a length-0 dataset can never be sampled to its length, + # so the stopping condition is ill-defined. The previous implementation + # crashed here with a cryptic `IndexError: Index N out of range` (it + # sampled the empty source and indexed into it). Raise a clear error + # instead -- for both strategies, since an empty source is degenerate + # either way and silently dropping it would change results. + if np.any(lengths_arr == 0): + empty = [i for i, ln in enumerate(lengths) if ln == 0] + raise ValueError( + "interleave_datasets with probabilities requires every dataset " + f"to be non-empty; datasets at indices {empty} are empty." + ) + else: + rng = np.random.default_rng(seed) + + # Draw source indices in 1000-sized blocks (matching the original + # iter_random_indices) until the stopping condition can be evaluated, + # i.e. until enough sources have reached their length. + blocks = [] + counts = np.zeros(n_datasets, dtype=np.int64) + reached = np.zeros(n_datasets, dtype=bool) + while not (reached.any() if not oversampling else reached.all()): + block = rng.choice(n_datasets, size=1000, p=probabilities) + blocks.append(block) + counts += np.bincount(block, minlength=n_datasets) + reached = counts >= lengths_arr + draws = np.concatenate(blocks) + + # A source becomes exhausted right AFTER its `length`-th draw, and + # the original loop checks the stop condition BEFORE appending. So + # the last draw we keep (inclusive) is at that length-th occurrence: + # - first_exhausted (any): the earliest such position over sources + # - all_exhausted (all): the latest such position over sources + exhaust_pos = np.full(n_datasets, -1, dtype=np.int64) + for s in range(n_datasets): + occ = np.flatnonzero(draws == s) + if len(occ) >= lengths[s]: + exhaust_pos[s] = occ[lengths[s] - 1] + valid_pos = exhaust_pos[exhaust_pos >= 0] + stop_at = valid_pos.min() if not oversampling else valid_pos.max() + used = draws[: stop_at + 1] + + # Map each source's k-th appearance to its k-th index with + # wrap-around: concatenated row = (k % length) + offset. + indices_arr = np.empty(len(used), dtype=np.int64) + for s in range(n_datasets): + pos = np.flatnonzero(used == s) + indices_arr[pos] = (np.arange(len(pos)) % lengths[s]) + offsets[s] + indices = indices_arr.tolist() else: # all_exhausted_without_replacement: each sample appears exactly once, so diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py index 498041f3b21..773e49b4be4 100644 --- a/tests/test_arrow_dataset.py +++ b/tests/test_arrow_dataset.py @@ -3836,6 +3836,20 @@ def test_interleave_datasets_probabilities_is_deterministic_and_balanced(stoppin assert src["d1"] > src["d3"] +@pytest.mark.parametrize("stopping_strategy", ["first_exhausted", "all_exhausted"]) +def test_interleave_datasets_probabilities_empty_source(stopping_strategy): + # A length-0 source can never be sampled to its length; the stop condition + # is ill-defined. Previously this crashed with a cryptic IndexError -- now + # it raises a clear ValueError for both strategies. + d_full = Dataset.from_dict({"a": [0, 1, 2]}) + d_empty = Dataset.from_dict({"a": []}) + with pytest.raises(ValueError): + interleave_datasets( + [d_full, d_empty], probabilities=[0.5, 0.5], seed=42, + stopping_strategy=stopping_strategy, + ) + + @pytest.mark.parametrize("batch_size", [4, 5]) @pytest.mark.parametrize("drop_last_batch", [False, True]) def test_dataset_iter_batch(batch_size, drop_last_batch):