import numpy as np import pytest import ray from ray.data._internal.execution.operators.mix_operator import MixOperator from ray.data._internal.logical.operators.n_ary_operator import ( MixStoppingCondition, estimate_num_mix_outputs, ) def _make_ds(source_id, num_rows, rows_per_block): """Create a dataset where every row has {"source": source_id}.""" num_blocks = num_rows // rows_per_block return ray.data.from_items([{"source": source_id}] * num_rows).repartition( num_blocks ) @pytest.mark.parametrize("weights", [[1, 1], None]) def test_mix_equal_weights(ray_start_10_cpus_shared, weights): """Equal weights with uniform blocks should produce 50/50 every 2 blocks. If weights are not provided, they should default to 1.0. """ rows_per_block = 10 ds1 = _make_ds(source_id=0, num_rows=500, rows_per_block=rows_per_block) ds2 = _make_ds(source_id=1, num_rows=500, rows_per_block=rows_per_block) mixed = ds1.mix( ds2, weights=weights, stopping_condition=MixStoppingCondition.STOP_ON_LONGEST_DROP, ) # We should round robin between the two datasets. # The output should alternate 10 rows for each dataset. for batch in mixed.iter_batches(batch_size=2 * rows_per_block): ratio = np.sum(batch["source"] == 0) / len(batch["source"]) assert ratio == 0.5 def test_mix_uneven_weights(ray_start_10_cpus_shared): """75/25 weights with uniform blocks should produce the right ratio every 4 blocks.""" rows_per_block = 10 ds1 = _make_ds(source_id=0, num_rows=750, rows_per_block=rows_per_block) ds2 = _make_ds(source_id=1, num_rows=250, rows_per_block=rows_per_block) mixed = ds1.mix( ds2, weights=[0.75, 0.25], stopping_condition=MixStoppingCondition.STOP_ON_LONGEST_DROP, ) for batch in mixed.iter_batches(batch_size=4 * rows_per_block): ratio = np.sum(batch["source"] == 0) / len(batch["source"]) assert ratio == 0.75 def test_mix_single_dataset(ray_start_10_cpus_shared): ds1 = _make_ds(source_id=0, num_rows=100, rows_per_block=10) mixed = ds1.mix(weights=[1.0]) result = mixed.take_all() assert len(result) == 100 def test_mix_three_datasets(ray_start_10_cpus_shared): """Three datasets with 50/30/20 weights should produce the right ratio every 10 blocks.""" rows_per_block = 10 ds1 = _make_ds(source_id=0, num_rows=500, rows_per_block=rows_per_block) ds2 = _make_ds(source_id=1, num_rows=300, rows_per_block=rows_per_block) ds3 = _make_ds(source_id=2, num_rows=200, rows_per_block=rows_per_block) mixed = ds1.mix( ds2, ds3, weights=[0.5, 0.3, 0.2], stopping_condition=MixStoppingCondition.STOP_ON_LONGEST_DROP, ) for batch in mixed.iter_batches(batch_size=10 * rows_per_block): ratio_ds1 = np.sum(batch["source"] == 0) / len(batch["source"]) ratio_ds2 = np.sum(batch["source"] == 1) / len(batch["source"]) ratio_ds3 = np.sum(batch["source"] == 2) / len(batch["source"]) assert ratio_ds1 == 0.5 assert ratio_ds2 == 0.3 assert ratio_ds3 == 0.2 def test_mix_stop_on_shortest(ray_start_10_cpus_shared): """With STOP_ON_SHORTEST, the pipeline stops when the shorter dataset is exhausted. The ratio should hold up to the stop point.""" ds1 = _make_ds(source_id=0, num_rows=20, rows_per_block=10) ds2 = _make_ds(source_id=1, num_rows=50, rows_per_block=10) mixed = ds1.mix( ds2, weights=[0.5, 0.5], stopping_condition=MixStoppingCondition.STOP_ON_SHORTEST, ) result = mixed.take_all() count_0 = sum(1 for r in result if r["source"] == 0) count_1 = sum(1 for r in result if r["source"] == 1) # ds1 is shorter (20 rows). All of ds1 should be consumed. assert count_0 == 20 # ds2 should have contributed roughly the same amount. The tolerance # of 1 block accounts for the case where we let another block through # before the ds1 input is marked as exhausted. assert abs(count_1 - count_0) <= 10 def test_mix_stop_on_longest_drop(ray_start_10_cpus_shared): """With STOP_ON_LONGEST_DROP, shorter datasets drop out and longer datasets continue until all are exhausted.""" rows_per_block = 10 ds1 = _make_ds(source_id=0, num_rows=500, rows_per_block=rows_per_block) ds2 = _make_ds(source_id=1, num_rows=200, rows_per_block=rows_per_block) mixed = ds1.mix( ds2, weights=[0.5, 0.5], stopping_condition=MixStoppingCondition.STOP_ON_LONGEST_DROP, ) ds2_rows_seen = 0 for batch in mixed.iter_batches(batch_size=2 * rows_per_block): # We should round robin between the two datasets, until ds2 is exhausted. if ds2_rows_seen < 200: ratio = np.sum(batch["source"] == 0) / len(batch["source"]) assert ratio == 0.5 ds2_rows_seen += np.sum(batch["source"] == 1) # After that point, we should only see rows from ds1. else: ratio = np.sum(batch["source"] == 0) / len(batch["source"]) assert ratio == 1.0 def test_mix_non_uniform_block_sizes(ray_start_10_cpus_shared): """Non-uniform block sizes: the deficit algorithm should still track the target ratio, just with wider oscillations.""" ds1 = _make_ds(source_id=0, num_rows=480, rows_per_block=120) # 120-row blocks ds2 = _make_ds(source_id=1, num_rows=160, rows_per_block=10) # 10-row blocks mixed = ds1.mix( ds2, weights=[0.75, 0.25], stopping_condition=MixStoppingCondition.STOP_ON_LONGEST_DROP, ) # Mix ordering: [ds0: 120], [ds1: 10, 10, 10, 10], ... # Expect every window of 160 rows to have a ratio of 0.75:0.25 for batch in mixed.iter_batches(batch_size=160): ratio = np.sum(batch["source"] == 0) / len(batch["source"]) assert ratio == 0.75 class TestEstimateNumMixOutputs: def test_stop_on_longest_drop(self): # Sum of all inputs. assert ( estimate_num_mix_outputs( [100, 200], [0.5, 0.5], MixStoppingCondition.STOP_ON_LONGEST_DROP ) == 300 ) def test_stop_on_shortest_equal_weights(self): # Limited by the smaller input: 100 / 0.5 = 200. assert ( estimate_num_mix_outputs( [100, 200], [0.5, 0.5], MixStoppingCondition.STOP_ON_SHORTEST ) == 200 ) def test_stop_on_shortest_uneven_weights(self): # ds1: 300 / 0.75 = 400, ds2: 100 / 0.25 = 400. Both sustain 400. assert ( estimate_num_mix_outputs( [300, 100], [0.75, 0.25], MixStoppingCondition.STOP_ON_SHORTEST ) == 400 ) def test_stop_on_shortest_bottleneck(self): # ds1: 100 / 0.75 = 133, ds2: 100 / 0.25 = 400. ds1 is bottleneck. assert ( estimate_num_mix_outputs( [100, 100], [0.75, 0.25], MixStoppingCondition.STOP_ON_SHORTEST ) == 133 ) def test_none_input(self): assert ( estimate_num_mix_outputs( [100, None], [0.5, 0.5], MixStoppingCondition.STOP_ON_SHORTEST ) is None ) def test_three_datasets(self): # ds1: 500/0.5=1000, ds2: 300/0.3=1000, ds3: 200/0.2=1000. assert ( estimate_num_mix_outputs( [500, 300, 200], [0.5, 0.3, 0.2], MixStoppingCondition.STOP_ON_SHORTEST, ) == 1000 ) def test_single_dataset(self): assert ( estimate_num_mix_outputs( [100], [1.0], MixStoppingCondition.STOP_ON_SHORTEST ) == 100 ) def test_invalid_stopping_condition(self): with pytest.raises(ValueError, match="Unknown stopping condition"): estimate_num_mix_outputs( [100, 200], [0.5, 0.5], "invalid" # type: ignore[arg-type] ) class TestMixOperatorEstimates: """Test that MixOperator.num_outputs_total and num_output_rows_total query the correct methods on input operators (blocks vs rows).""" def _make_mix_op(self, num_blocks, num_rows, weights, stopping_condition): """Create a MixOperator stub with mock input dependencies. Bypasses the full constructor (which requires real PhysicalOperators) and only sets the fields needed by num_outputs_total / num_output_rows_total. """ from unittest.mock import MagicMock from ray.data._internal.execution.interfaces import PhysicalOperator mock_inputs = [] for blocks, rows in zip(num_blocks, num_rows): mock = MagicMock(spec=PhysicalOperator) mock._name = "MockInput" mock.num_outputs_total.return_value = blocks mock.num_output_rows_total.return_value = rows mock.num_output_splits.return_value = 1 mock._output_dependencies = [] mock_inputs.append(mock) op = MixOperator( ray.data.DataContext.get_current(), *mock_inputs, weights=weights, stopping_condition=stopping_condition, ) return op, mock_inputs def test_num_outputs_total_stop_on_longest_drop(self): op, _ = self._make_mix_op( num_blocks=[10, 20], num_rows=[1000, 500], weights=[0.5, 0.5], stopping_condition=MixStoppingCondition.STOP_ON_LONGEST_DROP, ) # Should sum block counts (10 + 20 = 30), not row counts. assert op.num_outputs_total() == 30 def test_num_outputs_total_stop_on_shortest(self): op, _ = self._make_mix_op( num_blocks=[10, 20], num_rows=[1000, 500], weights=[0.5, 0.5], stopping_condition=MixStoppingCondition.STOP_ON_SHORTEST, ) # Can't estimate block count for STOP_ON_SHORTEST. assert op.num_outputs_total() is None def test_num_output_rows_total_stop_on_longest_drop(self): op, mocks = self._make_mix_op( num_blocks=[10, 20], num_rows=[1000, 500], weights=[0.5, 0.5], stopping_condition=MixStoppingCondition.STOP_ON_LONGEST_DROP, ) # Should sum row counts (1000 + 500 = 1500), not block counts. assert op.num_output_rows_total() == 1500 def test_num_output_rows_total_stop_on_shortest(self): op, _ = self._make_mix_op( num_blocks=[10, 20], num_rows=[1000, 500], weights=[0.75, 0.25], stopping_condition=MixStoppingCondition.STOP_ON_SHORTEST, ) # ds1: 1000/0.75 = 1333, ds2: 500/0.25 = 2000. Min = 1333. assert op.num_output_rows_total() == 1333 if __name__ == "__main__": import sys sys.exit(pytest.main(["-v", __file__]))