FEAT-#2375: addressing comments

Signed-off-by: Dmitry Chigarev <dmitry.chigarev@intel.com>
modin-project · Dec 8, 2020 · 6558f86 · 6558f86
1 parent 7474170
commit 6558f86
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 2 deletions.
diff --git a/modin/backends/pandas/query_compiler.py b/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")

diff --git a/modin/engines/base/frame/data.py b/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

diff --git a/modin/engines/base/frame/partition_manager.py b/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.
 

diff --git a/modin/pandas/test/test_groupby.py b/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)