feat(pyspark): implement zeroifnull

ibis-project · May 23, 2022 · db13241 · db13241
1 parent 4735e9a
commit db13241
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/ibis/backends/pyspark/compiler.py b/ibis/backends/pyspark/compiler.py
@@ -2012,3 +2012,9 @@ def compile_degrees(t, expr, scope, timecontext, **kwargs):
 @compiles(ops.Radians)
 def compile_radians(t, expr, scope, timecontext, **kwargs):
     return F.radians(t.translate(expr.op().arg, scope, timecontext, **kwargs))
+
+
+@compiles(ops.ZeroIfNull)
+def compile_zero_if_null(t, expr, scope, timecontext, **kwargs):
+    col = t.translate(expr.op().arg, scope, timecontext, **kwargs)
+    return F.when(col.isNull() | F.isnan(col), F.lit(0)).otherwise(col)
diff --git a/ibis/backends/tests/test_generic.py b/ibis/backends/tests/test_generic.py
@@ -546,7 +546,7 @@ def test_logical_negation_column(backend, alltypes, df, op):
     backend.assert_series_equal(result, expected, check_names=False)
 
 
-@pytest.mark.notimpl(["dask", "datafusion", "pandas", "pyspark"])
+@pytest.mark.notimpl(["dask", "datafusion", "pandas"])
 @pytest.mark.parametrize(
     ("dtype", "zero", "expected"),
     [("int64", 0, 1), ("float64", 0.0, 1.0)],
@@ -559,7 +559,7 @@ def test_zeroifnull_literals(con, dtype, zero, expected):
     )
 
 
-@pytest.mark.notimpl(["dask", "datafusion", "pandas", "pyspark"])
+@pytest.mark.notimpl(["dask", "datafusion", "pandas"])
 def test_zeroifnull_column(backend, alltypes, df):
     expr = alltypes.int_col.nullif(1).zeroifnull()
     result = expr.execute()