FEAT-#2375: addressing comments

Signed-off-by: Dmitry Chigarev <[email protected]>

Author: dchigarev

modin/backends/pandas/query_compiler.py

@@ -2610,7 +2610,9 @@ def compute_groupby(df, drop=False, partition_idx=0):
                 except (DataError, TypeError):
                     result = pandas.DataFrame(index=grouped_df.size().index)
                 if isinstance(result, pandas.Series):
-                    result = result.to_frame("__reduced__")
+                    result = result.to_frame(
+                        result.name if result.name is not None else "__reduced__"
+                    )
 
                 result_cols = result.columns
                 result.drop(columns=missmatched_cols, inplace=True, errors="ignore")

modin/engines/base/frame/data.py

@@ -1699,6 +1699,9 @@ def broadcast_apply_full_axis(
                 advance, and if not provided it must be computed.
             apply_indices : list-like (optional),
                 Indices of `axis ^ 1` to apply function over.
+            enumerate_partitions : bool (optional, default False),
+                Whether or not to pass partition index into applied `func`.
+                Note that `func` must be able to obtain `partition_idx` kwarg.
             dtypes : list-like (optional)
                 The data types of the result. This is an optimization
                 because there are functions that always result in a particular data

modin/engines/base/frame/partition_manager.py

@@ -248,6 +248,9 @@ def broadcast_axis_partitions(
             The flag to keep partitions for Modin Frame.
         apply_indices : list of ints (optional),
             Indices of `axis ^ 1` to apply function over.
+        enumerate_partitions : bool (optional, default False),
+            Whether or not to pass partition index into `apply_func`.
+            Note that `apply_func` must be able to obtain `partition_idx` kwarg.
         lengths : list(int), default None
             The list of lengths to shuffle the object.
 

modin/pandas/test/test_groupby.py

@@ -349,7 +349,11 @@ def maybe_get_columns(df, by):
 
     # Workaround for Pandas bug #34656. Recreate groupby object for Pandas
     pandas_groupby = pandas_df.groupby(by=pandas_by, as_index=as_index)
-    apply_functions = [lambda df: df.sum(), min]
+    apply_functions = [
+        lambda df: df.sum(),
+        lambda df: pandas.Series([1, 2, 3, 4], name="result"),
+        min,
+    ]
     for func in apply_functions:
         eval_apply(modin_groupby, pandas_groupby, func)