Function test_random_api_reseed_behavior

python/ray/data/tests/test_random_api.py:115–151

Test reseed_after_execution behavior across multiple epochs. For shuffle/block_order we compare full row order (no sort); for random_sample we compare sorted items so we only assert same set of sampled rows.

(base_dataset)

def test_random_api_reseed_behavior(base_dataset):
    """Test reseed_after_execution behavior across multiple epochs.

    For shuffle/block_order we compare full row order (no sort); for random_sample
    we compare sorted items so we only assert same set of sampled rows.
    """
    seed_value = 42
    seed_false = RandomSeedConfig(seed=seed_value, reseed_after_execution=False)
    seed_true = RandomSeedConfig(seed=seed_value, reseed_after_execution=True)

    # Test random_shuffle with reseed_after_execution=False
    ds = base_dataset.random_shuffle(seed=seed_false)
    epochs = collect_batches_from_epochs(ds, num_epochs=3, batch_size=5)
    assert (
        epochs[0] == epochs[1] == epochs[2]
    ), "reseed_after_execution=False should produce same results"

    # Test random_shuffle with reseed_after_execution=True
    ds = base_dataset.random_shuffle(seed=seed_true)
    epochs = collect_batches_from_epochs(ds, num_epochs=3, batch_size=5)
    assert (epochs[0] != epochs[1]) or (
        epochs[1] != epochs[2]
    ), "reseed_after_execution=True should produce different results per epoch"

    # Test random_sample with reseed_after_execution=False
    ds = base_dataset.random_sample(fraction=0.5, seed=seed_false)
    epochs = collect_batches_from_epochs(ds, num_epochs=2, batch_size=5)
    assert sorted(epochs[0]) == sorted(
        epochs[1]
    ), "reseed_after_execution=False should produce same sampled items"

    # Test randomize_block_order with reseed_after_execution=False
    ds = base_dataset.randomize_block_order(seed=seed_false)
    epochs = collect_batches_from_epochs(ds, num_epochs=2, batch_size=5)
    assert (
        epochs[0] == epochs[1]
    ), "reseed_after_execution=False should produce same block order"


@pytest.mark.parametrize(