Function test_split_is_not_disruptive

python/ray/data/tests/test_split.py:863–885

(ray_start_cluster)

def test_split_is_not_disruptive(ray_start_cluster):
    ray.shutdown()
    ds = ray.data.range(100, override_num_blocks=10).map_batches(lambda x: x)

    def verify_integrity(splits):
        for dss in splits:
            for batch in dss.iter_batches():
                pass
        for batch in ds.iter_batches():
            pass

    # No block splitting invovled: split 10 even blocks into 2 groups.
    verify_integrity(ds.split(2, equal=True))
    # Block splitting invovled: split 10 even blocks into 3 groups.
    verify_integrity(ds.split(3, equal=True))

    # Same as above but having tranforms post converting to lazy.
    verify_integrity(ds.map_batches(lambda x: x).split(2, equal=True))
    verify_integrity(ds.map_batches(lambda x: x).split(3, equal=True))

    # Same as above but having in-place tranforms post converting to lazy.
    verify_integrity(ds.randomize_block_order().split(2, equal=True))
    verify_integrity(ds.randomize_block_order().split(3, equal=True))


def test_streaming_train_test_split_hash(ray_start_regular_shared_2_cpus):