feat(datafusion): add functions sign, power, nullifzero, log

ibis-project · Jul 25, 2023 · ef72e40 · ef72e40
1 parent d4ec5c2
commit ef72e40
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 43 deletions.
diff --git a/ibis/backends/datafusion/compiler.py b/ibis/backends/datafusion/compiler.py
@@ -505,6 +505,40 @@ def degrees(op, **kw):
     return translate(op.arg, **kw) * df.lit(180) / df.lit(math.pi)
 
 
+@translate.register(ops.Power)
+def power(op, **kw):
+    base = translate(op.left, **kw)
+    exponent = translate(op.right, **kw)
+    return df.functions.power(base, exponent)
+
+
+@translate.register(ops.Sign)
+def sign(op, **kw):
+    arg = translate(op.arg, **kw)
+
+    arrow_sign = df.udf(
+        pc.sign,
+        input_types=[PyArrowType.from_ibis(op.arg.output_dtype)],
+        return_type=PyArrowType.from_ibis(op.output_dtype),
+        volatility="immutable",
+    )
+
+    return arrow_sign(arg)
+
+
+@translate.register(ops.NullIfZero)
+def null_if_zero(op, **kw):
+    arg = translate(op.arg, **kw)
+    return df.functions.nullif(arg, df.literal(0))
+
+
+@translate.register(ops.Log)
+def log(op, **kw):
+    arg = translate(op.arg, **kw)
+    base = translate(op.base, **kw)
+    return df.functions.log(base, arg)
+
+
 @translate.register(ops.RandomScalar)
 def random_scalar(_, **__):
     return df.functions.random()

diff --git a/ibis/backends/tests/test_numeric.py b/ibis/backends/tests/test_numeric.py
@@ -801,40 +801,23 @@ def test_isnan_isinf(
             L(5.556).sign(),
             1,
             id='sign-pos',
-            marks=[
-                pytest.mark.notimpl(
-                    ["datafusion"], raises=com.OperationNotDefinedError
-                ),
-            ],
         ),
         param(
             L(-5.556).sign(),
             -1,
             id='sign-neg',
-            marks=[
-                pytest.mark.notimpl(
-                    ["datafusion"], raises=com.OperationNotDefinedError
-                ),
-            ],
         ),
         param(
             L(0).sign(),
             0,
             id='sign-zero',
-            marks=[
-                pytest.mark.notimpl(
-                    ["datafusion"], raises=com.OperationNotDefinedError
-                ),
-            ],
         ),
         param(L(5.556).sqrt(), math.sqrt(5.556), id='sqrt'),
         param(
             L(5.556).log(2),
             math.log(5.556, 2),
             id='log-base',
-            marks=pytest.mark.notimpl(
-                ["datafusion", "druid"], raises=com.OperationNotDefinedError
-            ),
+            marks=pytest.mark.notimpl(["druid"], raises=com.OperationNotDefinedError),
         ),
         param(
             L(5.556).ln(),
@@ -911,14 +894,6 @@ def test_trig_functions_literals(con, expr, expected):
         param(_.dc.tan(), np.tan, id="tan"),
     ],
 )
-@pytest.mark.notyet(
-    ["datafusion"],
-    reason=(
-        "datafusion implements trig functions but can't easily test them due"
-        " to missing NullIfZero"
-    ),
-    raises=com.OperationNotDefinedError,
-)
 def test_trig_functions_columns(backend, expr, alltypes, df, expected_fn):
     dc_max = df.double_col.max()
     expr = alltypes.mutate(dc=(_.double_col / dc_max).nullifzero()).select(tmp=expr)
@@ -955,9 +930,6 @@ def test_trig_functions_columns(backend, expr, alltypes, df, expected_fn):
             lambda t: np.sign(t.double_col),
             id='sign',
             marks=[
-                pytest.mark.notimpl(
-                    ["datafusion"], raises=com.OperationNotDefinedError
-                ),
                 pytest.mark.notyet(
                     ["oracle"],
                     raises=AssertionError,
@@ -970,9 +942,6 @@ def test_trig_functions_columns(backend, expr, alltypes, df, expected_fn):
             lambda t: np.sign(-t.double_col),
             id='sign-negative',
             marks=[
-                pytest.mark.notimpl(
-                    ["datafusion"], raises=com.OperationNotDefinedError
-                ),
                 pytest.mark.notyet(
                     ["oracle"],
                     raises=AssertionError,
@@ -1013,9 +982,6 @@ def test_simple_math_functions_columns(
             lambda t: t.double_col.add(1).log(2),
             lambda t: np.log2(t.double_col + 1),
             id='log2',
-            marks=pytest.mark.notimpl(
-                ["datafusion"], raises=com.OperationNotDefinedError
-            ),
         ),
         param(
             lambda t: t.double_col.add(1).ln(),
@@ -1151,14 +1117,7 @@ def test_backend_specific_numerics(backend, con, df, alltypes, expr_fn, expected
         operator.mul,
         operator.truediv,
         operator.floordiv,
-        param(
-            operator.pow,
-            marks=[
-                pytest.mark.notimpl(
-                    ["datafusion"], raises=com.OperationNotDefinedError
-                ),
-            ],
-        ),
+        operator.pow,
     ],
     ids=lambda op: op.__name__,
 )