FEAT-#2303: fix OmniSci aggregates and add mean

modin/experimental/backends/omnisci/query_compiler.py

@@ -289,22 +289,29 @@ def groupby_dict_agg(self, by, func_dict, groupby_args, agg_args, drop=False):
         )
         return self.__constructor__(new_frame)
 
-    def count(self, axis=0, level=None, **kwargs):
-        return self._agg("count")
+    def count(self, **kwargs):
+        return self._agg("count", **kwargs)
 
-    def max(self, axis=0, level=None, **kwargs):
-        return self._agg("max")
+    def max(self, **kwargs):
+        return self._agg("max", **kwargs)
 
-    def min(self, axis=0, level=None, **kwargs):
-        return self._agg("min")
+    def min(self, **kwargs):
+        return self._agg("min", **kwargs)
 
-    def sum(self, axis=0, level=None, **kwargs):
-        return self._agg("sum")
+    def sum(self, **kwargs):
+        return self._agg("sum", **kwargs)
+
+    def mean(self, **kwargs):
+        return self._agg("mean", **kwargs)
 
     def _agg(self, agg, axis=0, level=None, **kwargs):
         if level is not None or axis != 0:
             return getattr(super(), agg)(axis=axis, level=level, **kwargs)
 
+        skipna = kwargs.get("skipna", True)
+        if not skipna:
+            return getattr(super(), agg)(axis=axis, level=level, **kwargs)
+
         new_frame = self._modin_frame.agg(agg)
         new_frame = new_frame._set_index(
             pandas.Index.__new__(pandas.Index, data=["__reduced__"], dtype="O")

modin/experimental/engines/omnisci_on_ray/test/test_dataframe.py

@@ -852,12 +852,19 @@ class TestAgg:
         "d": [11, 22, 33, 22, 33, 22],
     }
 
-    @pytest.mark.parametrize("agg", ["count", "max", "min", "sum"])
-    def test_simple_agg(self, agg):
-        def agg(df, agg, **kwargs):
-            return getattr(df, agg)()
+    @pytest.mark.parametrize("agg", ["max", "min", "sum", "mean"])
+    @pytest.mark.parametrize("skipna", bool_arg_values)
+    def test_simple_agg(self, agg, skipna):
+        def apply(df, agg, skipna, **kwargs):
+            return getattr(df, agg)(skipna=skipna)
 
-        run_and_compare(agg, data=self.data, agg=agg, force_lazy=False)
+        run_and_compare(apply, data=self.data, agg=agg, skipna=skipna, force_lazy=False)
+
+    def test_count_agg(self):
+        def apply(df, **kwargs):
+            return df.count()
+
+        run_and_compare(apply, data=self.data, force_lazy=False)
 
     @pytest.mark.parametrize("cols", ["a", "d"])
     @pytest.mark.parametrize("dropna", [True, False])