Function test_mix_equal_weights

python/ray/data/tests/test_mix.py:21–37

Equal weights with uniform blocks should produce 50/50 every 2 blocks. If weights are not provided, they should default to 1.0.

(ray_start_10_cpus_shared, weights)

@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):