[data] Add Dataset.mix() public API and user guide for weighted dataset mixing #63168
Merged
Conversation
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
![gemini-code-assist[bot]](https://cdn.statically.io/img/avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces the Dataset.mix public API to Ray Data, enabling weighted interleaving of multiple datasets. The changes include the implementation of the mix method, the addition of a comprehensive user guide, and the promotion of MixStoppingCondition to the public API. Feedback focuses on aligning documentation terminology with enum names, implementing validation for input weights to prevent logical errors, and ensuring consistent context usage in the execution plan.
ray-gardener
Bot
added
docs
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
richardliaw
reviewed
May 8, 2026
richardliaw
reviewed
May 8, 2026
Comment on lines
+2959
to
+2963
| stats = DatasetStats( | ||
| metadata={"Mix": []}, | ||
| parent=[d._raw_stats() for d in datasets], | ||
| ) | ||
| stats.time_total_s = time.perf_counter() - start_time |
Contributor
There was a problem hiding this comment.
what's up with all of this stuff?
Contributor
Author
There was a problem hiding this comment.
This is also done in Union. let's just keep it consistent and consider removing later since it's just measuring a small amount of logical operator creation time.
Contributor
There was a problem hiding this comment.
ack, i was looking at join and didn't see the equivalent
richardliaw
reviewed
May 8, 2026
| ) | ||
|
|
||
| @classmethod | ||
| @PublicAPI(stability="alpha", api_group=SMJ_API_GROUP) |
Contributor
There was a problem hiding this comment.
feel like this should not be in SMJ group but rather like training ingest group or something
Contributor
Author
There was a problem hiding this comment.
I think it's ok for now since it falls in "merging datasets." Let's land this and reorganize APIs in a followup.
richardliaw
reviewed
May 8, 2026
richardliaw
reviewed
May 8, 2026
richardliaw
reviewed
May 8, 2026
…ncy with union/zip Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Justin Yu <justinvyu@anyscale.com>
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Justin Yu <justinvyu@anyscale.com>
![cursor[bot]](https://cdn.statically.io/img/avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
Reviewed by Cursor Bugbot for commit aac2eab. Configure here.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Justin Yu <justinvyu@anyscale.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Justin Yu <justinvyu@anyscale.com>
matthewdeng
approved these changes
May 8, 2026
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
Lucas61000
pushed a commit
to Lucas61000/ray
that referenced
this pull request
May 15, 2026
…aset mixing (ray-project#63168) * Adds Dataset.mix() classmethod (alpha) for streaming weighted interleaving of multiple datasets, building on the internal MixOperator from ray-project#62450. * Exports MixStoppingCondition from `ray.data`. * Adds a user guide under Ray Data docs covering per-block mixing, random mixing, stopping conditions, and limitations. Moved the other "scaling collation" user guide to Ray Data docs. * Updates test_mix.py to use the public API instead of the internal helper. --------- Signed-off-by: Justin Yu <justinvyu@anyscale.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
justinvyu
added a commit
that referenced
this pull request
May 15, 2026
Adds a release test benchmark for `Dataset.mix()` (introduced in #63168) that measures mixing throughput and ratio accuracy. Benchmark design: - Creates 8 datasets reading ImageNet parquet, each stamped with a ds_index column - Mixes with Dataset.mix(), repartitions to 4 * batch_size rows per block - Consumes via TorchTrainer to mimic the seen weighting ratio when ingesting multiple local batches which are split across workers. - Tracks per-batch mixing ratios per worker, aggregates mean/std across workers via all_reduce to get the mean and standard deviation across **global batches.** - Asserts ratio mean is within 0.05 of target and std < 0.1 - Tests with and without a shuffling step after mixing with `--num-workers=1` to showcase the effectiveness of shuffling removing the dependency of mixing quality on the number of workers. [See here for more details.](https://docs.ray.io/en/master/data/mixing-data.html#random-mixing) --------- Signed-off-by: Justin Yu <justinvyu@anyscale.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
TruongQuangPhat
pushed a commit
to cyhapun/ray-fix-issue
that referenced
this pull request
May 27, 2026
…3286) Adds a release test benchmark for `Dataset.mix()` (introduced in ray-project#63168) that measures mixing throughput and ratio accuracy. Benchmark design: - Creates 8 datasets reading ImageNet parquet, each stamped with a ds_index column - Mixes with Dataset.mix(), repartitions to 4 * batch_size rows per block - Consumes via TorchTrainer to mimic the seen weighting ratio when ingesting multiple local batches which are split across workers. - Tracks per-batch mixing ratios per worker, aggregates mean/std across workers via all_reduce to get the mean and standard deviation across **global batches.** - Asserts ratio mean is within 0.05 of target and std < 0.1 - Tests with and without a shuffling step after mixing with `--num-workers=1` to showcase the effectiveness of shuffling removing the dependency of mixing quality on the number of workers. [See here for more details.](https://docs.ray.io/en/master/data/mixing-data.html#random-mixing) --------- Signed-off-by: Justin Yu <justinvyu@anyscale.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: phattruong <23120318@student.hcmus.edu.vn>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
MixOperatorfor weighted dataset mixing #62450.ray.data.Testing
See additional testing and result here: https://gist.github.com/justinvyu/0b73d66397a3fa0d9286f88b5e3ec3c3
To be added as a release test.