## 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 ray-project#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.

---------

Signed-off-by: my-vegetable-has-exploded <wy1109468038@gmail.com>
Signed-off-by: Aydin Abiar <aydin@anyscale.com>

## 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 ray-project#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.

---------

Signed-off-by: my-vegetable-has-exploded <wy1109468038@gmail.com>
Signed-off-by: YK <1811651+ykdojo@users.noreply.github.com>

## 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 ray-project#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.

---------

Signed-off-by: my-vegetable-has-exploded <wy1109468038@gmail.com>

## 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 ray-project#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.

---------

Signed-off-by: my-vegetable-has-exploded <wy1109468038@gmail.com>
Signed-off-by: Future-Outlier <eric901201@gmail.com>

## 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 ray-project#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.

---------

Signed-off-by: my-vegetable-has-exploded <wy1109468038@gmail.com>
Signed-off-by: peterxcli <peterxcli@gmail.com>