Repository: ray-project/ray
Issue category: Bug Report

Issue:
[Data] `GroupedData.map_groups()` doesn't allow partial callables
### What happened + What you expected to happen

In [this PR](https://github.com/ray-project/ray/pull/45310), a new requirement was imposed on the `fn` callable given as input to [`GroupedData.map_groups()`](https://docs.ray.io/en/latest/data/api/doc/ray.data.grouped_data.GroupedData.map_groups.html#ray.data.grouped_data.GroupedData.map_groups): that it have a `__name__` attribute. Unfortunately, callables partially parametrized using [`functools.partial()`](https://docs.python.org/3/library/functools.html#functools.partial) have no such attribute, so passing them into `.map_groups()` raises an error: `AttributeError: 'functools.partial' object has no attribute '__name__'`. This did not happen prior to the linked PR.

It's not a huge deal, but it did cause code to break unexpectedly, and I guess technically is in conflict with the type annotations on this method.

```
In [1]: import functools

In [2]: callable(functools.partial(lambda x: x))
Out[2]: True

In [3]: functools.partial(lambda x: x).__name__
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[3], line 1
----> 1 functools.partial(lambda x: x).__name__

AttributeError: 'functools.partial' object has no attribute '__name__'
```

### Versions / Dependencies

`ray >= 2.21`
PY3.10
macOS 14.4

### Reproduction script

```python
>>> import functools
>>> import ray
>>> ds = ray.data.range(10)
>>> ds.groupby("id").map_groups(lambda x: x)  # this is fine
MapBatches(<lambda>)
+- Sort
   +- Dataset(num_rows=10, schema={id: int64})
>>> ds.groupby("id").map_groups(functools.partial(lambda x: x))  # this errors
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[8], line 1
----> 1 ds.groupby("id").map_groups(functools.partial(lambda x: x))

File ~/.pyenv/versions/3.10.13/envs/ev-detection-py310/lib/python3.10/site-packages/ray/data/grouped_data.py:253, in GroupedData.map_groups(self, fn, compute, batch_format, fn_args, fn_kwargs, fn_constructor_args, fn_constructor_kwargs, num_cpus, num_gpus, concurrency, **ray_remote_args)
    249         yield from apply_udf_to_groups(fn, batch, *args, **kwargs)
    251 # Change the name of the wrapped function so that users see the name of their
    252 # function rather than `wrapped_fn` in the progress bar.
--> 253 wrapped_fn.__name__ = fn.__name__
    255 # Note we set batch_size=None here, so it will use the entire block as a batch,
    256 # which ensures that each group will be contained within a batch in entirety.
    257 return sorted_ds._map_batches_without_batch_size_validation(
    258     wrapped_fn,
    259     batch_size=None,
   (...)
    271     **ray_remote_args,
    272 )

AttributeError: 'functools.partial' object has no attribute '__name__'
```

### Issue Severity

Low: It annoys or frustrates me.


Hints:
Hi, can I take on this issue?

Patch:
diff --git a/python/ray/data/grouped_data.py b/python/ray/data/grouped_data.py
index 8f7b7dde118d..427ea18b7bbf 100644
--- a/python/ray/data/grouped_data.py
+++ b/python/ray/data/grouped_data.py
@@ -1,3 +1,4 @@
+from functools import partial
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 from ray.data._internal.aggregate import Count, Max, Mean, Min, Std, Sum
@@ -261,7 +262,10 @@ def wrapped_fn(batch, *args, **kwargs):
 
         # Change the name of the wrapped function so that users see the name of their
         # function rather than `wrapped_fn` in the progress bar.
-        wrapped_fn.__name__ = fn.__name__
+        if isinstance(fn, partial):
+            wrapped_fn.__name__ = fn.func.__name__
+        else:
+            wrapped_fn.__name__ = fn.__name__
 
         # Note we set batch_size=None here, so it will use the entire block as a batch,
         # which ensures that each group will be contained within a batch in entirety.


Test patch:
diff --git a/python/ray/data/tests/test_all_to_all.py b/python/ray/data/tests/test_all_to_all.py
index cf0cb8b2b2e7..a6b173383145 100644
--- a/python/ray/data/tests/test_all_to_all.py
+++ b/python/ray/data/tests/test_all_to_all.py
@@ -1167,7 +1167,6 @@ def test_groupby_map_groups_multicolumn(
     ray_start_regular_shared, ds_format, num_parts, use_push_based_shuffle
 ):
     # Test built-in count aggregation
-    print(f"Seeding RNG for test_groupby_arrow_count with: {RANDOM_SEED}")
     random.seed(RANDOM_SEED)
     xs = list(range(100))
     random.shuffle(xs)
@@ -1190,6 +1189,33 @@ def test_groupby_map_groups_multicolumn(
     ]
 
 
+def test_groupby_map_groups_with_partial():
+    """
+    The partial function name should show up as
+    +- Sort
+       +- MapBatches(func)
+    """
+    from functools import partial
+
+    def func(x, y):
+        return {f"x_add_{y}": [len(x["id"]) + y]}
+
+    df = pd.DataFrame({"id": list(range(100))})
+    df["key"] = df["id"] % 5
+
+    ds = ray.data.from_pandas(df).groupby("key").map_groups(partial(func, y=5))
+    result = ds.take_all()
+
+    assert result == [
+        {"x_add_5": 25},
+        {"x_add_5": 25},
+        {"x_add_5": 25},
+        {"x_add_5": 25},
+        {"x_add_5": 25},
+    ]
+    assert "MapBatches(func)" in ds.__repr__()
+
+
 def test_random_block_order_schema(ray_start_regular_shared):
     df = pd.DataFrame({"a": np.random.rand(10), "b": np.random.rand(10)})
     ds = ray.data.from_pandas(df).randomize_block_order()