Repository: modin-project/modin

Issue:
Add 'dropna' support for OmniSci backend
This is required by #1294


Discussion:

Resolving patch:
diff --git a/modin/experimental/backends/omnisci/query_compiler.py b/modin/experimental/backends/omnisci/query_compiler.py
--- a/modin/experimental/backends/omnisci/query_compiler.py
+++ b/modin/experimental/backends/omnisci/query_compiler.py
@@ -413,6 +413,22 @@ def drop(self, index=None, columns=None):
             )
         )
 
+    def dropna(self, axis=0, how="any", thresh=None, subset=None):
+        """Returns a new QueryCompiler with null values dropped along given axis.
+
+        Return:
+            a new QueryCompiler
+        """
+        if thresh is not None or axis != 0:
+            return super().dropna(axis=axis, how=how, thresh=thresh, subset=subset)
+
+        if subset is None:
+            subset = self.columns
+        return self.__constructor__(
+            self._modin_frame.dropna(subset=subset, how=how),
+            shape_hint=self._shape_hint,
+        )
+
     def dt_year(self):
         """Extract year from Datetime info
 
diff --git a/modin/experimental/engines/omnisci_on_ray/frame/calcite_builder.py b/modin/experimental/engines/omnisci_on_ray/frame/calcite_builder.py
--- a/modin/experimental/engines/omnisci_on_ray/frame/calcite_builder.py
+++ b/modin/experimental/engines/omnisci_on_ray/frame/calcite_builder.py
@@ -312,6 +312,17 @@ def build(self, op):
         self._to_calcite(op)
         return self.res
 
+    def _add_projection(self, frame):
+        """
+        Applies projection of the frame columns for table scan and discards
+        virtual "rowid" column.
+        """
+        proj = CalciteProjectionNode(
+            frame._table_cols, [self._ref(frame, col) for col in frame._table_cols]
+        )
+        self._push(proj)
+        return proj
+
     def _input_ctx(self):
         return self._input_ctx_stack[-1]
 
@@ -392,9 +403,7 @@ def _process_mask(self, op):
 
         # mask is currently always applied over scan, it means
         # we need additional projection to remove rowid column
-        fields = frame._table_cols
-        exprs = [self._ref(frame, col) for col in frame._table_cols]
-        self._push(CalciteProjectionNode(fields, exprs))
+        self._add_projection(frame)
 
     def _process_groupby(self, op):
         frame = op.input[0]
@@ -521,13 +530,8 @@ def _process_union(self, op):
 
     def _process_sort(self, op):
         frame = op.input[0]
-
-        # Sort should be applied to projections.
         if not isinstance(self._input_node(0), CalciteProjectionNode):
-            proj = CalciteProjectionNode(
-                frame._table_cols, [self._ref(frame, col) for col in frame._table_cols]
-            )
-            self._push(proj)
+            proj = self._add_projection(frame)
             self._input_ctx().replace_input_node(frame, proj, frame._table_cols)
 
         nulls = op.na_position.upper()
@@ -542,3 +546,8 @@ def _process_sort(self, op):
     def _process_filter(self, op):
         condition = self._translate(op.condition)
         self._push(CalciteFilterNode(condition))
+
+        if isinstance(self._input_node(0), CalciteScanNode):
+            # if filter was applied over scan, then we need additional
+            # projection to remove rowid column
+            self._add_projection(op.input[0])
diff --git a/modin/experimental/engines/omnisci_on_ray/frame/data.py b/modin/experimental/engines/omnisci_on_ray/frame/data.py
--- a/modin/experimental/engines/omnisci_on_ray/frame/data.py
+++ b/modin/experimental/engines/omnisci_on_ray/frame/data.py
@@ -282,6 +282,9 @@ def groupby_agg(self, by, axis, agg, groupby_args, **kwargs):
         new_columns = []
         index_cols = None
 
+        if groupby_args["dropna"]:
+            base = base.dropna(subset=groupby_cols, how="any")
+
         if groupby_args["as_index"]:
             index_cols = groupby_cols.copy()
         else:
@@ -451,6 +454,32 @@ def fillna(
 
         return new_frame
 
+    def dropna(self, subset, how="any"):
+        how_to_merge = {
+            "any": "AND",
+            "all": "OR",
+        }
+
+        # If index columns are not presented in the frame, then we have to create them
+        # based on "rowid". This is needed because 'dropna' preserves index.
+        if self._index_cols is None:
+            base = self._materialize_rowid()
+
+        checks = [base.ref(col).is_not_null() for col in subset]
+        condition = (
+            checks[0]
+            if len(checks) == 1
+            else OpExpr(how_to_merge[how], checks, np.dtype("bool"))
+        )
+        result = base.__constructor__(
+            columns=base.columns,
+            dtypes=base._dtypes,
+            op=FilterNode(base, condition),
+            index_cols=base._index_cols,
+            force_execution_mode=base._force_execution_mode,
+        )
+        return result
+
     def dt_extract(self, obj):
         exprs = self._index_exprs()
         for col in self.columns:


Test patch:
diff --git a/modin/experimental/engines/omnisci_on_ray/test/test_dataframe.py b/modin/experimental/engines/omnisci_on_ray/test/test_dataframe.py
--- a/modin/experimental/engines/omnisci_on_ray/test/test_dataframe.py
+++ b/modin/experimental/engines/omnisci_on_ray/test/test_dataframe.py
@@ -1663,5 +1663,33 @@ def applier(df1, df2, *args, **kwargs):
         )
 
 
+class TestDropna:
+    data = {
+        "col1": [1, 2, None, 2, 1],
+        "col2": [None, 3, None, 2, 1],
+        "col3": [2, 3, 4, None, 5],
+        "col4": [1, 2, 3, 4, 5],
+    }
+
+    @pytest.mark.parametrize("subset", [None, ["col1", "col2"]])
+    @pytest.mark.parametrize("how", ["all", "any"])
+    def test_dropna(self, subset, how):
+        def applier(df, *args, **kwargs):
+            return df.dropna(subset=subset, how=how)
+
+        run_and_compare(applier, data=self.data)
+
+    @pytest.mark.parametrize("by", ["col1", ["col1", "col2"], ["col1", "col4"]])
+    @pytest.mark.parametrize("dropna", [True, False])
+    def test_dropna_groupby(self, by, dropna):
+        def applier(df, *args, **kwargs):
+            # OmniSci backend preserves NaNs at the result of groupby,
+            # so replacing NaNs with '0' to match with Pandas.
+            # https://github.com/modin-project/modin/issues/2878
+            return df.groupby(by=by, dropna=dropna).sum().fillna(0)
+
+        run_and_compare(applier, data=self.data)
+
+
 if __name__ == "__main__":
     pytest.main(["-v", __file__])