Function test_train_test_split

python/ray/data/tests/test_split.py:744–776

(ray_start_regular_shared_2_cpus)

def test_train_test_split(ray_start_regular_shared_2_cpus):
    ds = ray.data.range(8)

    # float
    train, test = ds.train_test_split(test_size=0.25)
    assert extract_values("id", train.take()) == [0, 1, 2, 3, 4, 5]
    assert extract_values("id", test.take()) == [6, 7]

    # int
    train, test = ds.train_test_split(test_size=2)
    assert extract_values("id", train.take()) == [0, 1, 2, 3, 4, 5]
    assert extract_values("id", test.take()) == [6, 7]

    # shuffle
    train, test = ds.train_test_split(test_size=0.25, shuffle=True, seed=1)
    assert extract_values("id", train.take()) == [7, 4, 6, 0, 5, 2]
    assert extract_values("id", test.take()) == [1, 3]

    # error handling
    with pytest.raises(TypeError):
        ds.train_test_split(test_size=[1])

    with pytest.raises(ValueError):
        ds.train_test_split(test_size=-1)

    with pytest.raises(ValueError):
        ds.train_test_split(test_size=0)

    with pytest.raises(ValueError):
        ds.train_test_split(test_size=1.1)

    with pytest.raises(ValueError):
        ds.train_test_split(test_size=9)


def test_train_test_split_stratified(ray_start_regular_shared_2_cpus):