feat(datafusion): implement arbitrary

ibis-project · Mar 29, 2024 · 43a8f50 · 43a8f50
1 parent b514c23
commit 43a8f50
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/ibis/backends/datafusion/compiler.py b/ibis/backends/datafusion/compiler.py
@@ -64,7 +64,6 @@ class DataFusionCompiler(SQLGlotCompiler):
     )
 
     SIMPLE_OPS = {
-        ops.Arbitrary: "first_value",
         ops.ApproxMedian: "approx_median",
         ops.ArrayRemove: "array_remove_all",
         ops.BitAnd: "bit_and",
@@ -152,6 +151,12 @@ def visit_Substring(self, op, *, arg, start, length):
             return self.f.substr(arg, start, length)
         return self.f.substr(arg, start)
 
+    def visit_Arbitrary(self, op, *, arg, where):
+        cond = ~arg.is_(None)
+        if where is not None:
+            cond &= where
+        return self.agg.first_value(arg, where=cond)
+
     def visit_Variance(self, op, *, arg, how, where):
         if how == "sample":
             return self.agg.var_samp(arg, where=where)

diff --git a/ibis/backends/tests/test_aggregation.py b/ibis/backends/tests/test_aggregation.py
@@ -666,7 +666,6 @@ def test_reduction_ops(
         "impala",
         "mysql",
         "polars",
-        "datafusion",
         "mssql",
         "druid",
         "oracle",