From c0ee221c2d6073442b5c69df2bd3412674e6ced6 Mon Sep 17 00:00:00 2001 From: Myachev Date: Mon, 4 Jul 2022 15:28:22 +0200 Subject: [PATCH] FIX-#4604: address review comments Signed-off-by: Myachev --- modin/pandas/groupby.py | 13 ++++++++----- modin/pandas/test/test_groupby.py | 7 +++---- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/modin/pandas/groupby.py b/modin/pandas/groupby.py index e4fe5908fdf..e81a0a171f0 100644 --- 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 index 8939d24699c..763634a2e07 100644 --- 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(