FIX-modin-project#2269: Added required arguments for groupby_agg

Moved wrap_udf_function into backend because omnisci doesn't support executing lambdas. Signed-off-by: Gregory Shimansky <[email protected]>
YarShev · Oct 30, 2020 · a0e1b74 · a0e1b74
1 parent 657916f
commit a0e1b74
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 5 deletions.
diff --git a/modin/backends/base/query_compiler.py b/modin/backends/base/query_compiler.py
@@ -1396,14 +1396,30 @@ def groupby_size(
             drop=drop,
         )
 
-    def groupby_agg(self, by, axis, agg_func, groupby_args, agg_args, drop=False):
+    def groupby_agg(
+        self,
+        by,
+        is_multi_by,
+        idx_name,
+        axis,
+        agg_func,
+        agg_args,
+        agg_kwargs,
+        groupby_kwargs,
+        drop_,
+        drop=False,
+    ):
         return GroupByDefault.register(pandas.core.groupby.DataFrameGroupBy.aggregate)(
             self,
             by=by,
+            is_multi_by=is_multi_by,
+            idx_name=idx_name,
             axis=axis,
             agg_func=agg_func,
-            groupby_args=groupby_args,
             agg_args=agg_args,
+            agg_kwargs=agg_kwargs,
+            groupby_kwargs=groupby_kwargs,
+            drop_=drop_,
             drop=drop,
         )
 

diff --git a/modin/backends/pandas/query_compiler.py b/modin/backends/pandas/query_compiler.py
@@ -2596,6 +2596,8 @@ def groupby_agg(
         drop_,
         drop=False,
     ):
+        agg_func = wrap_udf_function(agg_func)
+
         if is_multi_by:
             return self.default_to_pandas(agg_func, *agg_args, **agg_kwargs)
 

diff --git a/modin/experimental/backends/omnisci/query_compiler.py b/modin/experimental/backends/omnisci/query_compiler.py
@@ -262,6 +262,24 @@ def groupby_count(self, by, axis, groupby_args, map_args, **kwargs):
         )
         return self.__constructor__(new_frame)
 
+    def groupby_agg(self,
+                    by,
+                    is_multi_by,
+                    idx_name,
+                    axis,
+                    agg_func,
+                    agg_args,
+                    agg_kwargs,
+                    groupby_kwargs,
+                    drop_,
+                    drop=False,
+    ):
+        # TODO: handle is_multi_by, idx_name and drop args
+        new_frame = self._modin_frame.groupby_agg(
+            by, axis, agg_func, groupby_kwargs, *agg_args, **agg_kwargs
+        )
+        return self.__constructor__(new_frame)
+
     def groupby_dict_agg(self, by, func_dict, groupby_args, agg_args, drop=False):
         """Apply aggregation functions to a grouped dataframe per-column.
 

diff --git a/modin/experimental/engines/omnisci_on_ray/test/test_dataframe.py b/modin/experimental/engines/omnisci_on_ray/test/test_dataframe.py
@@ -547,6 +547,7 @@ def groupby_count(df, cols, as_index, **kwargs):
 
         run_and_compare(groupby_count, data=self.data, cols=cols, as_index=as_index)
 
+    @pytest.mask.xfail(reason="Currently mean() passes a lambda into backend which cannot be executed on omnisci backend")
     @pytest.mark.parametrize("cols", cols_value)
     @pytest.mark.parametrize("as_index", bool_arg_values)
     def test_groupby_mean(self, cols, as_index):
@@ -577,6 +578,15 @@ def groupby(df, **kwargs):
 
         run_and_compare(groupby, data=self.data)
 
+    @pytest.mask.xfail(reason="Function specified as a string should be passed into backend API, but currently it is transformed into a lambda")
+    @pytest.mark.parametrize("cols", cols_value)
+    @pytest.mark.parametrize("as_index", bool_arg_values)
+    def test_groupby_agg_mean(self, cols, as_index):
+        def groupby_mean(df, cols, as_index, **kwargs):
+            return df.groupby(cols, as_index=as_index).agg("mean")
+
+        run_and_compare(groupby_mean, data=self.data, cols=cols, as_index=as_index)
+
     taxi_data = {
         "a": [1, 1, 2, 2],
         "b": [11, 21, 12, 11],

diff --git a/modin/pandas/groupby.py b/modin/pandas/groupby.py
@@ -29,7 +29,7 @@
 import pandas.core.common as com
 
 from modin.error_message import ErrorMessage
-from modin.utils import _inherit_docstrings, wrap_udf_function, try_cast_to_pandas
+from modin.utils import _inherit_docstrings, try_cast_to_pandas
 from modin.config import IsExperimental
 from .series import Series
 
@@ -834,8 +834,6 @@ def _apply_agg_function(self, f, drop=True, *args, **kwargs):
         """
         assert callable(f), "'{0}' object is not callable".format(type(f))
 
-        f = wrap_udf_function(f)
-
         new_manager = self._query_compiler.groupby_agg(
             by=self._by,
             is_multi_by=self._is_multi_by,