Function test_flat_map

python/ray/data/tests/test_map.py:348–374

(
    ray_start_regular_shared, udf_kind, target_max_block_size_infinite_or_default
)

@pytest.mark.parametrize("udf_kind", ["gen", "func"])
def test_flat_map(
    ray_start_regular_shared, udf_kind, target_max_block_size_infinite_or_default
):
    ds = ray.data.range(3)

    if udf_kind == "gen":

        def _udf(item: dict) -> Iterator[int]:
            for _ in range(2):
                yield {"id": item["id"] + 1}

    elif udf_kind == "func":

        def _udf(item: dict) -> dict:
            return [{"id": item["id"] + 1} for _ in range(2)]

    else:
        pytest.fail(f"Invalid udf_kind: {udf_kind}")

    assert sorted(extract_values("id", ds.flat_map(_udf).take())) == [
        1,
        1,
        2,
        2,
        3,
        3,
    ]


# Helper function to process timestamp data in nanoseconds