From db132418171f66e2902ac921e78c4bc1dfa99668 Mon Sep 17 00:00:00 2001
From: Phillip Cloud <417981+cpcloud@users.noreply.github.com>
Date: Sat, 21 May 2022 12:58:31 -0400
Subject: [PATCH] feat(pyspark): implement zeroifnull

---
 ibis/backends/pyspark/compiler.py   | 6 ++++++
 ibis/backends/tests/test_generic.py | 4 ++--
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/ibis/backends/pyspark/compiler.py b/ibis/backends/pyspark/compiler.py
index ee62601a91dc..8b7e6684d83c 100644
--- 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
index 1093f25dff1d..fd0905be091b 100644
--- 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()