Function test_zip_e2e

python/ray/data/tests/test_execution_optimizer_advanced.py:404–419

(
    ray_start_regular_shared_2_cpus, num_blocks1, num_blocks2, num_blocks3
)

    list(itertools.combinations_with_replacement(range(1, 4), 3)),
)
def test_zip_e2e(
    ray_start_regular_shared_2_cpus, num_blocks1, num_blocks2, num_blocks3
):
    n = 4
    ds1 = ray.data.range(n, override_num_blocks=num_blocks1)
    ds2 = ray.data.range(n, override_num_blocks=num_blocks2).map(
        column_udf("id", lambda x: x + 1)
    )
    ds3 = ray.data.range(n, override_num_blocks=num_blocks3).map(
        column_udf("id", lambda x: x + 2)
    )
    ds = ds1.zip(ds2, ds3)
    assert ds.take() == named_values(
        ["id", "id_1", "id_2"], zip(range(n), range(1, n + 1), range(2, n + 2))
    )
    _check_usage_record(["ReadRange", "Zip"])


def test_execute_to_legacy_block_list(