feat(pyspark): implement nullifzero

ibis-project · May 19, 2022 · 640234b · 640234b
1 parent addb2c1
commit 640234b
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/ibis/backends/pyspark/compiler.py b/ibis/backends/pyspark/compiler.py
@@ -1923,3 +1923,10 @@ def compile_unnest(t, expr, scope, timecontext, **kwargs):
     op = expr.op()
     column = t.translate(op.arg, scope, timecontext, **kwargs)
     return F.explode(column)
+
+
+@compiles(ops.NullIfZero)
+def compile_null_if_zero(t, expr, scope, timecontext, **kwargs):
+    op = expr.op()
+    arg = t.translate(op.arg, scope, timecontext, **kwargs)
+    return F.when(arg == 0, F.lit(None)).otherwise(arg)