在 pypika 中自定义 SQL 函数

pypika 是一个在 Python 中构建 SQL 语句的工具。这篇文章提供了一个很棒的入门教程，介绍了许多 pypika 的基本用法。

本文以 ClickHouse 中的 quantileExact 函数为例，介绍如何在 pypika 中自定义函数，尤其是自定义字符串拼接方式的实现方法。

问题背景

假设我们在 ClickHouse 数据库中有一个表格 my_table，其中包含每一天的每一笔账单金额。我们需要计算每一天的所有单笔账单金额的 75 分位数，可以构造如下 SQL 语句：

SQL

SELECT date,quantileExact(0.75)("amount") amount_75th_percentile
FROM my_table
GROUP BY date

其中，quantileExact 是计算分位数的函数，详见文档。

失败方法

首先，定义 CustomFunction，将 ["quantile", "value"] 作为参数。

Python

import sqlparse
from pygments import highlight
from pygments.formatters import TerminalFormatter
from pygments.lexers import SqlLexer
from pypika import Field, Query, Table
from pypika import CustomFunction

QuantileExact = CustomFunction("quantileExact", ["quantile", "value"])


# 定义表
my_table = Table("my_table")

# 定义查询
query = (
    Query.from_(my_table)
    .select(
        my_table.date,
        QuantileExact(0.75, Field("amount")).as_("amount_75th_percentile"),
    )
    .groupby(my_table.date)
)

# 打印生成的 SQL 查询
print(
    highlight(
        sqlparse.format(
            query.get_sql(quote_char=None),
            reindent=True,
            keyword_case="upper",
        ),
        SqlLexer(),
        TerminalFormatter(),
    )
)

SQL

SELECT date,quantileExact(0.75, amount) amount_75th_percentile
FROM my_table
GROUP BY date

但是，得到的 quantileExact 是将两个参数放在一个括号中，即 quantileExact(0.75, amount)。这并不是 ClickHouse 中的 quantileExact(0.75)("amount") 调用用法。

重写 get_function_sql 方法

Python

import sqlparse
from pygments import highlight
from pygments.formatters import TerminalFormatter
from pygments.lexers import SqlLexer
from pypika import Field, Query, Table
from pypika.terms import Function


class QuantileExact(Function):
    def __init__(self, quantile, field, *args, **kwargs):
        super(QuantileExact, self).__init__("quantileExact", *args, **kwargs)
        self.quantile = quantile
        self.field = field

    def get_function_sql(self, **kwargs) -> str:
        return "{name}({quantile})({field})".format(
            name=self.name,
            quantile=self.quantile,
            field=self.field,
        )


# 定义表
my_table = Table("my_table")

# 定义查询
query = (
    Query.from_(my_table)
    .select(
        my_table.date,
        QuantileExact(0.75, Field("amount")).as_("amount_75th_percentile"),
    )
    .groupby(my_table.date)
)

# 打印生成的 SQL 查询
print(
    highlight(
        sqlparse.format(
            query.get_sql(quote_char=None),
            reindent=True,
            keyword_case="upper",
        ),
        SqlLexer(),
        TerminalFormatter(),
    )
)

我们重写了 get_function_sql 方法，将其变成 "{name}({quantile})({field})" 的形式。这样，QuantileExact(0.75, Field("amount")) 就能得到 quantileExact(0.75)("amount") 的 SQL 语句，从而被 ClickHouse 正确识别。

评论