feat: implement dropna for SQL backends

Author: jcrist

ibis/backends/base/sql/compiler/select_builder.py

@@ -1,9 +1,12 @@
 from __future__ import annotations
 
+import functools
+import operator
 from typing import NamedTuple
 
 import toolz
 
+import ibis
 import ibis.common.exceptions as com
 import ibis.expr.analysis as L
 import ibis.expr.operations as ops
@@ -442,6 +445,28 @@ def _collect_Distinct(self, expr, toplevel=False):
 
         self._collect(expr.op().table, toplevel=toplevel)
 
+    def _collect_DropNa(self, expr, toplevel=False):
+        if toplevel:
+            op = expr.op()
+            if op.subset is None:
+                columns = [op.table[c] for c in op.table.columns]
+            else:
+                columns = op.subset
+            if columns:
+                filters = [
+                    functools.reduce(
+                        operator.and_ if op.how == "any" else operator.or_,
+                        [c.notnull() for c in columns],
+                    )
+                ]
+            elif op.how == "all":
+                filters = [ibis.literal(False)]
+            else:
+                filters = []
+            self.table_set = op.table
+            self.select_set = [op.table]
+            self.filters = filters
+
     def _collect_Limit(self, expr, toplevel=False):
         if not toplevel:
             return

ibis/backends/dask/execution/generic.py

@@ -117,7 +117,10 @@
             execute_difference_dataframe_dataframe,
         )
     ],
-    ops.DropNa: [((dd.DataFrame, tuple), execute_node_dropna_dataframe)],
+    ops.DropNa: [
+        ((dd.DataFrame, tuple), execute_node_dropna_dataframe),
+        ((dd.DataFrame, type(None)), execute_node_dropna_dataframe),
+    ],
     ops.FillNa: [
         ((dd.DataFrame, simple_types), execute_node_fillna_dataframe_scalar),
         ((dd.DataFrame,), execute_node_fillna_dataframe_dict),

ibis/backends/pandas/execution/generic.py

@@ -1144,9 +1144,11 @@ def execute_node_log_number_number(op, value, base, **kwargs):
     return math.log(value, base)
 
 
+@execute_node.register(ops.DropNa, pd.DataFrame, type(None))
 @execute_node.register(ops.DropNa, pd.DataFrame, tuple)
 def execute_node_dropna_dataframe(op, df, subset, **kwargs):
-    subset = [col.get_name() for col in subset] if subset else None
+    if subset is not None:
+        subset = [col.get_name() for col in subset]
     return df.dropna(how=op.how, subset=subset)
 
 

ibis/backends/pyspark/compiler.py

@@ -1986,7 +1986,9 @@ def compile_not_null(t, expr, scope, timecontext, **kwargs):
 def compile_dropna_table(t, expr, scope, timecontext, **kwargs):
     op = expr.op()
     table = t.translate(op.table, scope, timecontext, **kwargs)
-    subset = [col.get_name() for col in op.subset] if op.subset else None
+    subset = op.subset
+    if subset is not None:
+        subset = [col.get_name() for col in subset]
     return table.dropna(how=op.how, subset=subset)
 
 

ibis/backends/tests/test_generic.py

@@ -11,7 +11,6 @@
 
 import ibis
 import ibis.common.exceptions as com
-import ibis.util as util
 from ibis import _
 from ibis import literal as L
 
@@ -385,45 +384,29 @@ def test_mutate_rename(alltypes):
     assert list(result.columns) == ["bool_col", "string_col", "dupe_col"]
 
 
+@pytest.mark.parametrize('how', ['any', 'all'])
 @pytest.mark.parametrize(
-    ('how', 'subset'),
-    [
-        ('any', None),
-        ('any', []),
-        ('any', ['int_col', 'na_col']),
-        ('all', None),
-        ('all', ['int_col', 'na_col']),
-        ('all', 'none_col'),
-    ],
-)
-@pytest.mark.notimpl(
-    [
-        "clickhouse",
-        "datafusion",
-        "impala",
-        "mysql",
-        "postgres",
-        "sqlite",
-    ]
+    'subset', [None, [], 'col_1', ['col_1', 'col_2'], ['col_1', 'col_3']]
 )
-@pytest.mark.notyet(["duckdb"], reason="non-finite value support")
+@pytest.mark.notimpl(["datafusion"])
 def test_dropna_table(backend, alltypes, how, subset):
-    table = alltypes.mutate(na_col=np.nan)
-    table = table.mutate(none_col=None)
-    table = table.mutate(none_col=table['none_col'].cast('float64'))
+    is_two = alltypes.int_col == 2
+    is_four = alltypes.int_col == 4
+
+    table = alltypes.mutate(
+        col_1=is_two.ifelse(ibis.NA, alltypes.float_col),
+        col_2=is_four.ifelse(ibis.NA, alltypes.float_col),
+        col_3=(is_two | is_four).ifelse(ibis.NA, alltypes.float_col),
+    ).select("col_1", "col_2", "col_3")
+
     table_pandas = table.execute()
 
     result = table.dropna(subset, how).execute().reset_index(drop=True)
-    subset = util.promote_list(subset) if subset else table_pandas.columns
     expected = table_pandas.dropna(how=how, subset=subset).reset_index(
         drop=True
     )
 
-    # check_dtype is False here because there are dtype diffs between
-    # Pyspark and Pandas on Java 8 - the 'bool_col' of an empty DataFrame
-    # is type object in Pyspark, and type bool in Pandas. This diff does
-    # not exist in Java 11.
-    backend.assert_frame_equal(result, expected, check_dtype=False)
+    backend.assert_frame_equal(result, expected)
 
 
 def test_select_sort_sort(alltypes):

ibis/expr/operations/relations.py

@@ -794,7 +794,7 @@ class DropNa(TableNode, sch.HasSchema):
 
     table = rlz.table
     how = rlz.isin({'any', 'all'})
-    subset = rlz.optional(rlz.tuple_of(rlz.column_from("table")), default=())
+    subset = rlz.optional(rlz.tuple_of(rlz.column_from("table")), default=None)
 
     @property
     def schema(self):

ibis/expr/types/relations.py

@@ -874,9 +874,8 @@ def dropna(
         """
         from ibis.expr import operations as ops
 
-        if subset is None:
-            subset = []
-        subset = util.promote_list(subset)
+        if subset is not None:
+            subset = util.promote_list(subset)
         return ops.DropNa(self, how, subset).to_expr()
 
     def fillna(