Function test_async_flat_map

python/ray/data/tests/test_map.py:1113–1142

(
    shutdown_only, udf_kind, target_max_block_size_infinite_or_default
)

@pytest.mark.parametrize("udf_kind", ["coroutine", "async_gen"])
def test_async_flat_map(
    shutdown_only, udf_kind, target_max_block_size_infinite_or_default
):
    class AsyncActor:
        def __init__(self):
            pass

        if udf_kind == "async_gen":

            async def __call__(self, row):
                id = row["id"]
                yield {"id": id}
                await asyncio.sleep(random.randint(0, 5) / 100)
                yield {"id": id + 1}

        elif udf_kind == "coroutine":

            async def __call__(self, row):
                id = row["id"]
                await asyncio.sleep(random.randint(0, 5) / 100)
                return [{"id": id}, {"id": id + 1}]

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

    n = 10
    ds = ray.data.from_items([{"id": i} for i in range(0, n, 2)])
    ds = ds.flat_map(AsyncActor, concurrency=1, max_concurrency=2)
    output = ds.take_all()
    assert sorted(extract_values("id", output)) == list(range(n))


class TestGenerateTransformFnForAsyncMap: