[Data] Support async callable classes in flat_map() (#51180)

<!-- Thank you for your contribution! Please review
https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before
opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR.
If you don't have the access to it, we will shortly find a reviewer and
assign them to your PR. -->

## Why are these changes needed?
This PR adds async generator support to `flat_map`.
The implementation is similar to how #46129 handled async callable
classes for map_batches(), changes include:
* Generalize the logic in `_generate_transform_fn_for_async_flat_map` so
it can process both batches and rows
* Add test case for async `flat_map`
<!-- Please give a short summary of the change and the problem this
solves. -->

## Related issue number
#50329 
<!-- For example: "Closes #1234" -->

## Checks

- [x] I've signed off every commit(by using the -s flag, i.e., `git
commit -s`) in this PR.
- [x] I've run `scripts/format.sh` to lint the changes in this PR.
- [ ] I've included any doc changes needed for
https://docs.ray.io/en/master/.
- [ ] I've added any new APIs to the API Reference. For example, if I
added a
method in Tune, I've added it in `doc/source/tune/api/` under the
           corresponding `.rst` file.
- [ ] I've made sure the tests are passing. Note that there might be a
few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [x] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(

---------

Signed-off-by: Drice1999 <chenxh267@gmail.com>

Author: Drice1999

python/ray/data/_internal/planner/plan_udf_map_op.py

@@ -376,7 +376,7 @@ def _generate_transform_fn_for_map_batches(
 ) -> MapTransformCallable[DataBatch, DataBatch]:
     if inspect.iscoroutinefunction(fn):
         # UDF is a callable class with async generator `__call__` method.
-        transform_fn = _generate_transform_fn_for_async_map_batches(fn)
+        transform_fn = _generate_transform_fn_for_async_map(fn, _validate_batch_output)
 
     else:
 
@@ -423,64 +423,66 @@ def transform_fn(
     return transform_fn
 
 
-def _generate_transform_fn_for_async_map_batches(
+def _generate_transform_fn_for_async_map(
     fn: UserDefinedFunction,
-) -> MapTransformCallable[DataBatch, DataBatch]:
-    def transform_fn(
-        input_iterable: Iterable[DataBatch], _: TaskContext
-    ) -> Iterable[DataBatch]:
+    validate_fn,
+) -> MapTransformCallable:
+    # Generates a transform function for asynchronous mapping of items (either batches or rows)
+    # using a user-defined function (UDF). This consolidated function handles both asynchronous
+    # batch processing and asynchronous flat mapping (e.g., rows) based on the provided UDF.
+    def transform_fn(input_iterable: Iterable, _: TaskContext) -> Iterable:
         # Use a queue to store outputs from async generator calls.
-        # We will put output batches into this queue from async
+        # We will put output items into this queue from async
         # generators, and in the main event loop, yield them from
         # the queue as they become available.
-        output_batch_queue = queue.Queue()
+        output_item_queue = queue.Queue()
         # Sentinel object to signal the end of the async generator.
         sentinel = object()
 
-        async def process_batch(batch: DataBatch):
+        async def process_item(item):
             try:
-                output_batch_iterator = await fn(batch)
+                output_item_iterator = await fn(item)
                 # As soon as results become available from the async generator,
                 # put them into the result queue so they can be yielded.
-                async for output_batch in output_batch_iterator:
-                    output_batch_queue.put(output_batch)
+                async for output_item in output_item_iterator:
+                    output_item_queue.put(output_item)
             except Exception as e:
-                output_batch_queue.put(
+                output_item_queue.put(
                     e
                 )  # Put the exception into the queue to signal an error
 
-        async def process_all_batches():
+        async def process_all_items():
             try:
                 loop = ray.data._map_actor_context.udf_map_asyncio_loop
-                tasks = [loop.create_task(process_batch(x)) for x in input_iterable]
+                tasks = [loop.create_task(process_item(x)) for x in input_iterable]
 
                 ctx = ray.data.DataContext.get_current()
                 if ctx.execution_options.preserve_order:
                     for task in tasks:
-                        await task()
+                        await task
                 else:
                     for task in asyncio.as_completed(tasks):
                         await task
             finally:
-                output_batch_queue.put(sentinel)
+                output_item_queue.put(sentinel)
 
-        # Use the existing event loop to create and run Tasks to process each batch
+        # Use the existing event loop to create and run Tasks to process each item
         loop = ray.data._map_actor_context.udf_map_asyncio_loop
-        asyncio.run_coroutine_threadsafe(process_all_batches(), loop)
+        asyncio.run_coroutine_threadsafe(process_all_items(), loop)
 
         # Yield results as they become available.
         while True:
-            # Here, `out_batch` is a one-row output batch
+            # Here, `out_item` is a one-row output item
             # from the async generator, corresponding to a
-            # single row from the input batch.
-            out_batch = output_batch_queue.get()
-            if out_batch is sentinel:
+            # single row from the input item.
+            out_item = output_item_queue.get()
+            if out_item is sentinel:
                 # Break out of the loop when the sentinel is received.
                 break
-            if isinstance(out_batch, Exception):
-                raise out_batch
-            _validate_batch_output(out_batch)
-            yield out_batch
+            if isinstance(out_item, Exception):
+                raise out_item
+            validate_fn(out_item)
+            yield out_item
 
     return transform_fn
 
@@ -511,11 +513,17 @@ def transform_fn(rows: Iterable[Row], _: TaskContext) -> Iterable[Row]:
 def _generate_transform_fn_for_flat_map(
     fn: UserDefinedFunction,
 ) -> MapTransformCallable[Row, Row]:
-    def transform_fn(rows: Iterable[Row], _: TaskContext) -> Iterable[Row]:
-        for row in rows:
-            for out_row in fn(row):
-                _validate_row_output(out_row)
-                yield out_row
+    if inspect.iscoroutinefunction(fn):
+        # UDF is a callable class with async generator `__call__` method.
+        transform_fn = _generate_transform_fn_for_async_map(fn, _validate_row_output)
+
+    else:
+
+        def transform_fn(rows: Iterable[Row], _: TaskContext) -> Iterable[Row]:
+            for row in rows:
+                for out_row in fn(row):
+                    _validate_row_output(out_row)
+                    yield out_row
 
     return transform_fn
 

python/ray/data/tests/test_map.py

@@ -1687,6 +1687,32 @@ async def __call__(self, batch):
     )
 
 
+def test_flat_map_async_generator(shutdown_only):
+    async def fetch_data(id):
+        return {"id": id}
+
+    class AsyncActor:
+        def __init__(self):
+            pass
+
+        async def __call__(self, row):
+            id = row["id"]
+            task1 = asyncio.create_task(fetch_data(id))
+            task2 = asyncio.create_task(fetch_data(id + 1))
+            print(f"yield task1: {id}")
+            yield await task1
+            print(f"sleep: {id}")
+            await asyncio.sleep(id % 5)
+            print(f"yield task2: {id}")
+            yield await task2
+
+    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(0, n)), output
+
+
 def test_map_batches_async_exception_propagation(shutdown_only):
     ray.shutdown()
     ray.init(num_cpus=2)