feat(sqlite): implement ops.Arbitrary

Author: cpcloud

ibis/backends/sqlite/registry.py

@@ -199,6 +199,15 @@ def _literal(t, op):
     return base_literal(t, op)
 
 
+def _arbitrary(t, op):
+    if (how := op.how) == "heavy":
+        raise com.OperationNotDefinedError(
+            "how='heavy' not implemented for the SQLite backend"
+        )
+
+    return reduction(getattr(sa.func, f"_ibis_sqlite_arbitrary_{how}"))(t, op)
+
+
 operation_registry.update(
     {
         # TODO(kszucs): don't dispatch on op.arg since that should be always an
@@ -322,5 +331,6 @@ def _literal(t, op):
         ops.RandomScalar: fixed_arity(
             lambda: 0.5 + sa.func.random() / sa.cast(-1 << 64, sa.REAL), 0
         ),
+        ops.Arbitrary: _arbitrary,
     }
 )

ibis/backends/sqlite/udf.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import abc
 import functools
 import inspect
 import math
@@ -350,6 +351,32 @@ def __init__(self):
         super().__init__(operator.xor)
 
 
+class _ibis_sqlite_arbitrary(abc.ABC):
+    def __init__(self) -> None:
+        self.value = None
+
+    @abc.abstractmethod
+    def step(self, value):
+        ...
+
+    def finalize(self) -> int | None:
+        return self.value
+
+
+@udaf
+class _ibis_sqlite_arbitrary_first(_ibis_sqlite_arbitrary):
+    def step(self, value):
+        if self.value is None:
+            self.value = value
+
+
+@udaf
+class _ibis_sqlite_arbitrary_last(_ibis_sqlite_arbitrary):
+    def step(self, value):
+        if value is not None:
+            self.value = value
+
+
 def _number_of_arguments(callable):
     signature = inspect.signature(callable)
     parameters = signature.parameters.values()

ibis/backends/tests/test_aggregation.py

@@ -554,7 +554,6 @@ def mean_and_std(v):
                 [
                     'impala',
                     'mysql',
-                    'sqlite',
                     'polars',
                     'datafusion',
                     "mssql",
@@ -571,7 +570,6 @@ def mean_and_std(v):
                 [
                     'impala',
                     'mysql',
-                    'sqlite',
                     'polars',
                     'datafusion',
                     "mssql",
@@ -589,7 +587,6 @@ def mean_and_std(v):
                     [
                         'impala',
                         'mysql',
-                        'sqlite',
                         'polars',
                         'datafusion',
                         "mssql",