Function test_random_shuffle

python/ray/data/tests/test_random_e2e.py:113–168

(
    shutdown_only, configure_shuffle_method, disable_fallback_to_object_extension
)

def test_random_shuffle(
    shutdown_only, configure_shuffle_method, disable_fallback_to_object_extension
):
    # Assert random 2 distinct random-shuffle pipelines yield different orders
    r1 = ray.data.range(100).random_shuffle().take(999)
    r2 = ray.data.range(100).random_shuffle().take(999)
    assert r1 != r2, (r1, r2)

    # Assert same random-shuffle pipeline yielding 2 different orders,
    # when executed
    ds = ray.data.range(100).random_shuffle()
    r1 = ds.take(999)
    r2 = ds.take(999)
    assert r1 != r2, (r1, r2)

    r1 = ray.data.range(100, override_num_blocks=1).random_shuffle().take(999)
    r2 = ray.data.range(100, override_num_blocks=1).random_shuffle().take(999)
    assert r1 != r2, (r1, r2)

    assert (
        ray.data.range(100)
        .random_shuffle()
        .repartition(1)
        ._logical_plan.initial_num_blocks()
        == 1
    )
    r1 = ray.data.range(100).random_shuffle().repartition(1).take(999)
    r2 = ray.data.range(100).random_shuffle().repartition(1).take(999)
    assert r1 != r2, (r1, r2)

    r0 = ray.data.range(100, override_num_blocks=5).take(999)
    r1 = ray.data.range(100, override_num_blocks=5).random_shuffle(seed=0).take(999)
    r2 = ray.data.range(100, override_num_blocks=5).random_shuffle(seed=0).take(999)
    r3 = ray.data.range(100, override_num_blocks=5).random_shuffle(seed=12345).take(999)
    assert r1 == r2, (r1, r2)
    assert r1 != r0, (r1, r0)
    assert r1 != r3, (r1, r3)

    r0 = ray.data.range(100, override_num_blocks=5).take(999)
    r1 = ray.data.range(100, override_num_blocks=5).random_shuffle(seed=0).take(999)
    r2 = ray.data.range(100, override_num_blocks=5).random_shuffle(seed=0).take(999)
    assert r1 == r2, (r1, r2)
    assert r1 != r0, (r1, r0)

    # Test move.
    ds = ray.data.range(100, override_num_blocks=2)
    r1 = ds.random_shuffle().take(999)
    ds = ds.map(lambda x: x).take(999)
    r2 = ray.data.range(100).random_shuffle().take(999)
    assert r1 != r2, (r1, r2)

    # Test empty dataset.
    ds = ray.data.from_items([])
    r1 = ds.random_shuffle()
    assert r1.count() == 0
    assert r1.take() == ds.take()