[hello-peter-tang]

Why are these changes needed?

Dataset.mix only rejected weights with weight <= 0. That lets NaN and +inf through, and they are then divided by the total weight, producing non-finite normalized weights that corrupt the deficit-based block selection. -inf was rejected, but with the misleading "Weights must be positive" message.

What is the change?

• Added validate_mix_weights(), shared by the logical Mix operator and the physical MixOperator.
• Non-finite weights now raise ValueError: Weights must be finite. Got weights: [...]; zero and negative weights keep the existing positive-weights error.
• Documented the validation on Dataset.mix.

Related issue number

Fixes #67387

Tests

• Added parameterized tests for nan, inf, -inf, 0.0, and -1.0 at both the Dataset.mix level and the MixOperator level.
• python/ray/data/tests/test_mix.py: 30 passed (20 existing + 10 new).
• black and ruff clean on the touched files.

[gemini-code-assist]

Code Review

This pull request introduces a centralized validate_mix_weights function to ensure that mix weights are positive and finite, preventing issues with NaN or infinite values. It updates the dataset documentation and adds unit tests to verify this validation. A review comment identifies a missing import of MixOperator in the new test cases, which would lead to a NameError during execution.

[gemini-code-assist]

The test uses MixOperator on line 56, but it is not imported within this test function. If MixOperator is not imported at the top of test_mix.py, this will raise a NameError during test execution. To ensure the test is self-contained and robust, please import MixOperator locally.

Suggested change

	from unittest.mock import MagicMock
	from ray.data._internal.execution.interfaces import PhysicalOperator
	from unittest.mock import MagicMock
	from ray.data._internal.execution.interfaces import PhysicalOperator
	from ray.data._internal.execution.operators.mix_operator import MixOperator

[Jadiouo]

Verified against base 8bfa990 and head 702036c (Python 3.13, CPU only, Ray nightly with ray/data swapped).

Through the public API (ds1.mix(ds2, weights=...), then take_all()): on base, weights=[nan, 1] and [inf, 1] are accepted by mix() and execution then hangs (I stopped it after 30s). With the PR, NaN, +inf and -inf raise ValueError: Weights must be finite. Got weights: [...] at the mix() call. Zero, negative and all-zero weights still raise the "positive" error, and finite weights produce the same block order as base. test_mix.py: 30 passed; the nan/inf/-inf cases fail on base.

Two small things: the gemini comment about a missing MixOperator import doesn't apply, since it's imported at the top of the file. The two parametrize lists are identical and could share one constant.

One question: richardliaw suggested on #67387 holding off until more people report this. Has that been cleared?

[hello-peter-tang]

@richardliaw — there's now an independent reproduction on the linked issue: @Jadiouo verified against base 8bfa990 vs this PR's head (702036c) that on base weights=[nan, 1] and [inf, 1] are accepted by mix() and execution hangs, while this PR raises ValueError immediately. That's the "someone actually ran into it" signal you asked for on #67387. Given the failure mode is a hang rather than a clear error, would you be open to lifting the hold-off and having this land as a cheap guard? Happy to close it if you'd still rather wait.