feat(api): add approx_quantiles for computing approximate quantiles

Author: jcrist

ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_approx_quantiles/array/out.sql

@@ -0,0 +1,7 @@
+SELECT
+  [
+    approx_quantiles(`t0`.`double_col`, 4 IGNORE NULLS)[1],
+    approx_quantiles(`t0`.`double_col`, 4 IGNORE NULLS)[2],
+    approx_quantiles(`t0`.`double_col`, 4 IGNORE NULLS)[3]
+  ] AS `qs`
+FROM `functional_alltypes` AS `t0`

ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_approx_quantiles/complete-array/out.sql

@@ -0,0 +1,3 @@
+SELECT
+  approx_quantiles(`t0`.`double_col`, 4 IGNORE NULLS) AS `qs`
+FROM `functional_alltypes` AS `t0`

ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_approx_quantiles/scalar/out.sql

@@ -0,0 +1,3 @@
+SELECT
+  approx_quantiles(`t0`.`double_col`, 2 IGNORE NULLS)[1] AS `qs`
+FROM `functional_alltypes` AS `t0`

ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_approx_quantiles/shuffled-array/out.sql

@@ -0,0 +1,7 @@
+SELECT
+  [
+    approx_quantiles(`t0`.`double_col`, 4 IGNORE NULLS)[2],
+    approx_quantiles(`t0`.`double_col`, 4 IGNORE NULLS)[1],
+    approx_quantiles(`t0`.`double_col`, 4 IGNORE NULLS)[3]
+  ] AS `qs`
+FROM `functional_alltypes` AS `t0`

ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_approx_quantiles/tricky-scalar/out.sql

@@ -0,0 +1,3 @@
+SELECT
+  approx_quantiles(`t0`.`double_col`, 100000 IGNORE NULLS)[33333] AS `qs`
+FROM `functional_alltypes` AS `t0`

ibis/backends/bigquery/tests/unit/test_compiler.py

@@ -677,3 +677,19 @@ def test_time_from_hms_with_micros(snapshot):
     literal = ibis.literal(datetime.time(12, 34, 56))
     result = ibis.to_sql(literal, dialect="bigquery")
     snapshot.assert_match(result, "no_micros.sql")
+
+
+@pytest.mark.parametrize(
+    "quantiles",
+    [
+        param(0.5, id="scalar"),
+        param(1 / 3, id="tricky-scalar"),
+        param([0.25, 0.5, 0.75], id="array"),
+        param([0.5, 0.25, 0.75], id="shuffled-array"),
+        param([0, 0.25, 0.5, 0.75, 1], id="complete-array"),
+    ],
+)
+def test_approx_quantiles(alltypes, quantiles, snapshot):
+    query = alltypes.double_col.approx_quantile(quantiles).name("qs")
+    result = ibis.to_sql(query, dialect="bigquery")
+    snapshot.assert_match(result, "out.sql")

ibis/backends/polars/compiler.py

@@ -785,6 +785,7 @@ def execute_mode(op, **kw):
 
 
 @translate.register(ops.Quantile)
+@translate.register(ops.ApproxQuantile)
 def execute_quantile(op, **kw):
     arg = translate(op.arg, **kw)
     quantile = translate(op.quantile, **kw)

ibis/backends/sql/compilers/bigquery/__init__.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import decimal
+import math
 import re
 from typing import TYPE_CHECKING, Any
 
@@ -392,6 +394,41 @@ def visit_GroupConcat(self, op, *, arg, sep, where, order_by):
 
         return sge.GroupConcat(this=arg, separator=sep)
 
+    def visit_ApproxQuantile(self, op, *, arg, quantile, where):
+        if not isinstance(op.quantile, ops.Literal):
+            raise com.UnsupportedOperationError(
+                "quantile must be a literal in BigQuery"
+            )
+
+        # BigQuery syntax is `APPROX_QUANTILES(col, resolution)` to return
+        # `resolution + 1` quantiles array. To handle this, we compute the
+        # resolution ourselves then restructure the output array as needed.
+        # To avoid excessive resolution we arbitrarily cap it at 100,000 -
+        # since these are approximate quantiles anyway this seems fine.
+        quantiles = util.promote_list(op.quantile.value)
+        fracs = [decimal.Decimal(str(q)).as_integer_ratio() for q in quantiles]
+        resolution = min(math.lcm(*(den for _, den in fracs)), 100_000)
+        indices = [(num * resolution) // den for num, den in fracs]
+
+        if where is not None:
+            arg = self.if_(where, arg, NULL)
+
+        if not op.arg.dtype.is_floating():
+            arg = self.cast(arg, dt.float64)
+
+        array = self.f.approx_quantiles(
+            arg, sge.IgnoreNulls(this=sge.convert(resolution))
+        )
+        if isinstance(op, ops.ApproxQuantile):
+            return array[indices[0]]
+
+        if indices == list(range(resolution + 1)):
+            return array
+        else:
+            return sge.Array(expressions=[array[i] for i in indices])
+
+    visit_ApproxMultiQuantile = visit_ApproxQuantile
+
     def visit_FloorDivide(self, op, *, left, right):
         return self.cast(self.f.floor(self.f.ieee_divide(left, right)), op.dtype)
 

ibis/backends/sql/compilers/clickhouse.py

@@ -188,18 +188,24 @@ def visit_CountStar(self, op, *, where, arg):
             return self.f.countIf(where)
         return sge.Count(this=STAR)
 
-    def visit_Quantile(self, op, *, arg, quantile, where):
-        if where is None:
-            return self.agg.quantile(arg, quantile, where=where)
-
-        func = "quantile" + "s" * isinstance(op, ops.MultiQuantile)
+    def _visit_quantile(self, func, arg, quantile, where):
         return sge.ParameterizedAgg(
-            this=f"{func}If",
+            this=f"{func}If" if where is not None else func,
             expressions=util.promote_list(quantile),
-            params=[arg, where],
+            params=[arg, where] if where is not None else [arg],
         )
 
-    visit_MultiQuantile = visit_Quantile
+    def visit_Quantile(self, op, *, arg, quantile, where):
+        return self._visit_quantile("quantile", arg, quantile, where)
+
+    def visit_MultiQuantile(self, op, *, arg, quantile, where):
+        return self._visit_quantile("quantiles", arg, quantile, where)
+
+    def visit_ApproxQuantile(self, op, *, arg, quantile, where):
+        return self._visit_quantile("quantileTDigest", arg, quantile, where)
+
+    def visit_ApproxMultiQuantile(self, op, *, arg, quantile, where):
+        return self._visit_quantile("quantilesTDigest", arg, quantile, where)
 
     def visit_Correlation(self, op, *, left, right, how, where):
         if how == "pop":

ibis/backends/sql/compilers/datafusion.py

@@ -51,6 +51,7 @@ class DataFusionCompiler(SQLGlotCompiler):
     )
 
     SIMPLE_OPS = {
+        ops.ApproxQuantile: "approx_percentile_cont",
         ops.ApproxMedian: "approx_median",
         ops.ArrayRemove: "array_remove_all",
         ops.BitAnd: "bit_and",

ibis/backends/sql/compilers/duckdb.py

@@ -532,6 +532,17 @@ def visit_Quantile(self, op, *, arg, quantile, where):
     def visit_MultiQuantile(self, op, *, arg, quantile, where):
         return self.visit_Quantile(op, arg=arg, quantile=quantile, where=where)
 
+    def visit_ApproxQuantile(self, op, *, arg, quantile, where):
+        # duckdb casts the return value back to the input type (so integer
+        # returns are always integral). Casting the result would give an
+        # integral float, we want to cast the values to treat them the same as
+        # if they were a float input.
+        if not op.arg.dtype.is_floating():
+            arg = self.cast(arg, dt.float64)
+        return self.agg.approx_quantile(arg, quantile, where=where)
+
+    visit_ApproxMultiQuantile = visit_ApproxQuantile
+
     def visit_HexDigest(self, op, *, arg, how):
         if how in ("md5", "sha256"):
             return getattr(self.f, how)(arg)

ibis/backends/sql/compilers/exasol.py

@@ -223,6 +223,8 @@ def visit_Quantile(self, op, *, arg, quantile, where):
             expression=sge.Order(expressions=[sge.Ordered(this=arg)]),
         )
 
+    visit_ApproxQuantile = visit_Quantile
+
     def visit_TimestampTruncate(self, op, *, arg, unit):
         short_name = unit.short
         unit_mapping = {"W": "IW"}

ibis/backends/sql/compilers/mssql.py

@@ -219,6 +219,14 @@ def visit_CountDistinct(self, op, *, arg, where):
             arg = self.if_(where, arg, NULL)
         return self.f.count(sge.Distinct(expressions=[arg]))
 
+    def visit_ApproxQuantile(self, op, *, arg, quantile, where):
+        if where is not None:
+            arg = self.if_(where, arg, NULL)
+        return sge.WithinGroup(
+            this=self.f.approx_percentile_cont(quantile),
+            expression=sge.Order(expressions=[sge.Ordered(this=arg, nulls_first=True)]),
+        )
+
     def visit_DayOfWeekIndex(self, op, *, arg):
         return self.f.datepart(self.v.weekday, arg) - 1
 

ibis/backends/sql/compilers/oracle.py

@@ -308,6 +308,15 @@ def visit_Quantile(self, op, *, arg, quantile, where):
         )
         return expr
 
+    def visit_ApproxQuantile(self, op, *, arg, quantile, where):
+        if where is not None:
+            arg = self.if_(where, arg)
+
+        return sge.WithinGroup(
+            this=self.f.approx_percentile(quantile),
+            expression=sge.Order(expressions=[sge.Ordered(this=arg)]),
+        )
+
     def visit_CountDistinct(self, op, *, arg, where):
         if where is not None:
             arg = self.if_(where, arg)

ibis/backends/sql/compilers/postgres.py

@@ -238,6 +238,8 @@ def visit_Quantile(self, op, *, arg, quantile, where):
         return expr
 
     visit_MultiQuantile = visit_Quantile
+    visit_ApproxQuantile = visit_Quantile
+    visit_ApproxMultiQuantile = visit_Quantile
 
     def visit_Correlation(self, op, *, left, right, how, where):
         if how == "sample":

ibis/backends/sql/compilers/pyspark.py

@@ -290,6 +290,15 @@ def visit_Quantile(self, op, *, arg, quantile, where):
 
     visit_MultiQuantile = visit_Quantile
 
+    def visit_ApproxQuantile(self, op, *, arg, quantile, where):
+        if not op.arg.dtype.is_floating():
+            arg = self.cast(arg, dt.float64)
+        if where is not None:
+            arg = self.if_(where, arg, NULL)
+        return self.f.approx_percentile(arg, quantile)
+
+    visit_ApproxMultiQuantile = visit_ApproxQuantile
+
     def visit_Correlation(self, op, *, left, right, how, where):
         if (left_type := op.left.dtype).is_boolean():
             left = self.cast(left, dt.Int32(nullable=left_type.nullable))

ibis/backends/sql/compilers/risingwave.py

@@ -6,7 +6,7 @@
 import ibis.expr.datatypes as dt
 import ibis.expr.operations as ops
 from ibis.backends.sql.compilers import PostgresCompiler
-from ibis.backends.sql.compilers.base import ALL_OPERATIONS
+from ibis.backends.sql.compilers.base import ALL_OPERATIONS, NULL
 from ibis.backends.sql.datatypes import RisingWaveType
 from ibis.backends.sql.dialects import RisingWave
 
@@ -22,6 +22,8 @@ class RisingWaveCompiler(PostgresCompiler):
         ops.DateFromYMD,
         ops.Mode,
         ops.RandomUUID,
+        ops.MultiQuantile,
+        ops.ApproxMultiQuantile,
         *(
             op
             for op in ALL_OPERATIONS
@@ -65,6 +67,17 @@ def visit_Correlation(self, op, *, left, right, how, where):
             op, left=left, right=right, how=how, where=where
         )
 
+    def visit_Quantile(self, op, *, arg, quantile, where):
+        if where is not None:
+            arg = self.if_(where, arg, NULL)
+        suffix = "cont" if op.arg.dtype.is_numeric() else "disc"
+        return sge.WithinGroup(
+            this=self.f[f"percentile_{suffix}"](quantile),
+            expression=sge.Order(expressions=[sge.Ordered(this=arg)]),
+        )
+
+    visit_ApproxQuantile = visit_Quantile
+
     def visit_TimestampTruncate(self, op, *, arg, unit):
         unit_mapping = {
             "Y": "year",

ibis/backends/sql/compilers/snowflake.py

@@ -608,6 +608,12 @@ def visit_Quantile(self, op, *, arg, quantile, where):
         quantile = self.f.percentile_cont(quantile)
         return sge.WithinGroup(this=quantile, expression=order_by)
 
+    def visit_ApproxQuantile(self, op, *, arg, quantile, where):
+        if where is not None:
+            arg = self.if_(where, arg, NULL)
+
+        return self.f.approx_percentile(arg, quantile)
+
     def visit_CountStar(self, op, *, arg, where):
         if where is None:
             return super().visit_CountStar(op, arg=arg, where=where)

ibis/backends/sql/compilers/trino.py

@@ -133,6 +133,17 @@ def visit_Correlation(self, op, *, left, right, how, where):
 
         return self.agg.corr(left, right, where=where)
 
+    def visit_ApproxQuantile(self, op, *, arg, quantile, where):
+        # trino casts the return value back to the input type (so integer
+        # returns are always integral). Casting the result would give an
+        # integral float, we want to cast the values to treat them the same as
+        # if they were a float input.
+        if not op.arg.dtype.is_floating():
+            arg = self.cast(arg, dt.float64)
+        return self.agg.approx_quantile(arg, quantile, where=where)
+
+    visit_ApproxMultiQuantile = visit_ApproxQuantile
+
     def visit_BitXor(self, op, *, arg, where):
         a, b = map(sg.to_identifier, "ab")
         input_fn = combine_fn = sge.Lambda(