FIX-modin-project#4604: address review comments

Signed-off-by: Myachev <anatoly.myachev@intel.com>
anmyachev · Jul 4, 2022 · c0ee221 · c0ee221
1 parent bd26aa0
commit c0ee221
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 9 deletions.
diff --git a/modin/pandas/groupby.py b/modin/pandas/groupby.py
@@ -513,14 +513,17 @@ def try_get_str_func(fn):
             relabeling_required, func_dict, new_columns, order = reconstruct_func(
                 func, **kwargs
             )
-            func_dict = {col: try_get_str_func(fn) for col, fn in func_dict.items()}
+            _wrapper = lambda fn: fn  # noqa: E731
             if any(isinstance(fn, list) for fn in func_dict.values()):
                 # multicolumn case
                 # putting functions in a `list` allows to achieve multicolumn in each partition
-                func_dict = {
-                    col: fn if isinstance(fn, list) else [fn]
-                    for col, fn in func_dict.items()
-                }
+                _wrapper = lambda fn: [fn]  # noqa: E731
+            func_dict = {
+                col: try_get_str_func(fn)
+                if isinstance(fn, list)
+                else _wrapper(try_get_str_func(fn))
+                for col, fn in func_dict.items()
+            }
             if (
                 relabeling_required
                 and not self._as_index

diff --git a/modin/pandas/test/test_groupby.py b/modin/pandas/test/test_groupby.py
@@ -1511,12 +1511,11 @@ def test_agg_4604():
     def col3(x):
         return np.max(x)
 
+    by = ["col1"]
     agg_func = {"col2": ["sum", "min"], "col3": col3}
 
-    md_res = md_df.groupby(["col1"]).agg(agg_func)
-    pd_res = pd_df.groupby(["col1"]).agg(agg_func)
-
-    df_equals(md_res, pd_res)
+    md_groupby, pd_groupby = md_df.groupby(by), pd_df.groupby(by)
+    eval_agg(md_groupby, pd_groupby, agg_func)
 
 
 @pytest.mark.parametrize(