[data] support generator udf for map_groups

[my-vegetable-has-exploded]

Description

This pr support return a generator object from map_groups UDF. if UDF have a large output , we return iterator to reduce memory cost.

Related issues

Close #57935

Additional information

This change centers on the _apply_udf_to_groups helper function within the file ray/data/grouped_data.py.

map_groups internally calls map_batches, providing a wrapper function (wrapped_fn) that in turn calls _apply_udf_to_groups to apply the user's UDF to each group.

The key modification is that instead of directly yielding the UDF's return value, the logic now inspects the result first. If the result is an Iterator, it is consumed with yield from to produce each data batch individually. If it is not an iterator, the single data batch is yielded directly, preserving the original behavior.

[gemini-code-assist]

Code Review

This pull request introduces a valuable enhancement by adding support for generator UDFs in map_groups. This change can significantly reduce memory consumption for UDFs that produce large outputs. The implementation in _apply_udf_to_groups is clean and correctly handles both single DataBatch and Iterator[DataBatch] return types. The new test case test_map_groups_generator_udf is comprehensive and effectively validates the new functionality. I have one minor suggestion to improve code clarity by updating a type hint to align with this new capability.

[cursor]

Bug: UDF Type Checking Fails on Older Python Versions

The isinstance check for UDF results uses typing.Iterator, which is a generic alias not suitable for runtime type checking. This causes a TypeError on Python versions prior to 3.10, preventing the correct detection of iterator or generator objects.

 

[cursor]

Bug: Runtime Error with Typing Generics

The isinstance(result, Iterator) check uses typing.Iterator, which is a generic type alias. On Python 3.7-3.9, this raises a TypeError because typing generics aren't designed for runtime checks. This prevents UDFs returning generators from being correctly unpacked, causing map_groups to fail or produce incorrect results.