Skip to content

Commit

Permalink
feat(clickhouse): implement quantile/multiquantile
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud committed Jan 9, 2023
1 parent bc7fdab commit 96d7d1b
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 2 deletions.
35 changes: 34 additions & 1 deletion ibis/backends/clickhouse/compiler/values.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import functools
from functools import partial
from operator import add, mul, sub
from typing import Any, Literal, Mapping
from typing import Any, Callable, Literal, Mapping

import sqlglot as sg

Expand Down Expand Up @@ -180,6 +180,39 @@ def _not_all(op, **kw):
return translate_val(ops.Not(ops.All(op.arg)), **kw)


def _quantiles(quantiles_translator_func: Callable, func_name: str):
def translate(op, **kw):
quantile = quantiles_translator_func(op.quantile)
args = [_sql(translate_val(op.arg, **kw))]
func = func_name

if (where := op.where) is not None:
args.append(_sql(translate_val(where, **kw)))
func += "If"

return f"{func}({quantile})({', '.join(args)})"

return translate


@translate_val.register(ops.Quantile)
def _quantile(op, **kw):
def quantiles_translator_func(quantiles):
return _sql(translate_val(quantiles, **kw))

return _quantiles(quantiles_translator_func, func_name="quantile")(op, **kw)


@translate_val.register(ops.MultiQuantile)
def _multi_quantile(op, **kw):
def quantiles_translator_func(quantiles):
if not isinstance(quantiles, ops.Literal):
raise TypeError("ClickHouse quantile only accepts a list of Python floats")
return ", ".join(map(str, quantiles.value))

return _quantiles(quantiles_translator_func, func_name="quantiles")(op, **kw)


def _agg_variance_like(func):
variants = {"sample": f"{func}Samp", "pop": f"{func}Pop"}

Expand Down
1 change: 0 additions & 1 deletion ibis/backends/tests/test_aggregation.py
Original file line number Diff line number Diff line change
Expand Up @@ -568,7 +568,6 @@ def test_reduction_ops(
mark.notimpl(
[
"bigquery",
"clickhouse",
"dask",
"datafusion",
"impala",
Expand Down

0 comments on commit 96d7d1b

Please sign in to comment.