Пользовательские функции Python (UDF)

chDB позволяет регистрировать функции Python как пользовательские функции (UDF), вызываемые из SQL. Они выполняются нативно в том же процессе — без запуска подпроцессов и накладных расходов на сериализацию. Функции типобезопасны, поддерживают автоматический вывод типов на основе аннотаций Python и позволяют настраивать обработку NULL и исключений.

Быстрый старт

from chdb import query, func
from chdb.sqltypes import INT64

@func([INT64, INT64], INT64)
def add(a, b):
    return a + b

result = query("SELECT add(2, 3)")
print(result)  # 5

Способы регистрации

Декоратор `@func`

Самый простой способ зарегистрировать UDF. Значение атрибута __name__ функции становится именем SQL-функции.

from chdb import func
from chdb.sqltypes import INT64, STRING

# Explicit types
@func([INT64, INT64], INT64)
def add(a, b):
    return a + b

# Types inferred from annotations
@func()
def multiply(a: int, b: int) -> int:
    return a * b

# Explicit return_type, arg_types inferred from annotations
@func(return_type=STRING)
def greet(name: str):
    return f"Hello, {name}!"

Декорированную функцию по-прежнему можно вызывать как обычную функцию Python:

add(2, 3)       # 5 (Python call)
query("SELECT add(2, 3)")  # 5 (SQL call)

`create_function`

Зарегистрируйте любой вызываемый объект (lambda, функцию, method) под явно заданным именем:

from chdb import create_function, query
from chdb.sqltypes import INT64, STRING

create_function("strlen", len, arg_types=[STRING], return_type=INT64)
query("SELECT strlen('hello')")  # 5

create_function("double", lambda x: x * 2, arg_types=[INT64], return_type=INT64)
query("SELECT double(21)")  # 42

`drop_function`

Удаляет зарегистрированную UDF. Если удалить незарегистрированное имя, ничего не произойдёт, поэтому этот метод можно безопасно вызывать безусловно:

from chdb import drop_function

drop_function("strlen")
# query("SELECT strlen('hello')")  # Error: function not found

Система типов

Доступные типы

Все типы доступны для импорта из chdb.sqltypes:

from chdb.sqltypes import (
    # Boolean
    BOOL,
    # Signed integers
    INT8, INT16, INT32, INT64, INT128, INT256,
    # Unsigned integers
    UINT8, UINT16, UINT32, UINT64, UINT128, UINT256,
    # Floating point
    FLOAT32, FLOAT64,
    # String
    STRING,
    # Date and time
    DATE, DATE32, DATETIME, DATETIME64,
)

Указание типов

Типы можно указывать четырьмя способами:

Способ	Пример	Описание
Константа `ChdbType`	`INT64`, `STRING`	Импортируется из `chdb.sqltypes`
Строковое представление типа ClickHouse	`"Int64"`, `"String"`	Стандартные имена типов ClickHouse
Параметризованная строка	`"DateTime('UTC')"`, `"DateTime64(6)"`	Для типов с параметрами
Тип Python	`int`, `str`, `float`	Передаётся напрямую в `arg_types`/`return_type` или используется в качестве аннотации типа в сигнатуре функции

from chdb import create_function, func
from chdb.sqltypes import INT64

# All equivalent:
create_function("f1", lambda x: x * 2, arg_types=[INT64], return_type=INT64)
create_function("f2", lambda x: x * 2, arg_types=["Int64"], return_type="Int64")
create_function("f3", lambda x: x * 2, arg_types=[int], return_type=int)

@func()
def f4(x: int) -> int:
    return x * 2

Автоматический вывод типов

Если arg_types или return_type не указаны, chDB выводит типы на основе аннотаций типов Python:

Тип Python	Тип ClickHouse
`bool`	`Bool`
`int`	`Int64`
`float`	`Float64`
`str`	`String`
`bytes`	`String`
`bytearray`	`String`
`datetime.date`	`Date`
`datetime.datetime`	`DateTime64(6)`

@func()
def process(name: str, age: int) -> str:
    return f"{name} is {age} years old"

# Equivalent to:
# @func([STRING, INT64], STRING)

Тип возвращаемого значения обязателен всегда: если return_type не указан и у функции нет аннотации возвращаемого типа, регистрация завершится ошибкой. Типы аргументов, напротив, необязательны: параметр без явно указанного типа и аннотации может динамически принимать любой поддерживаемый входной тип.

Обработка NULL

Параметр on_null определяет поведение, если какой-либо входной аргумент имеет значение NULL.

Значение	Поведение
`"skip"` (по умолчанию)	Немедленно вернуть NULL, не вызывая функцию
`"pass"`	Преобразовать NULL в Python `None` и вызвать функцию как обычно

Также можно использовать enum: chdb.NullHandling.SKIP / chdb.NullHandling.PASS.

Пример: default (пропуск)

@func(return_type="Int64")
def increment(x: int) -> int:
    return x + 1

query("SELECT increment(NULL)")  # NULL
query("SELECT increment(5)")     # 6

Пример: передача NULL в виде `None`

@func(return_type="Int64", on_null="pass")
def null_to_zero(x):
    return 0 if x is None else x + 1

query("SELECT null_to_zero(NULL)")  # 0
query("SELECT null_to_zero(5)")     # 6

Пример: несколько аргументов

@func(arg_types=["Int64", "Int64"], return_type="Int64", on_null="pass")
def add_or_zero(a, b):
    return (a or 0) + (b or 0)

query("SELECT add_or_zero(NULL, 5)")    # 5
query("SELECT add_or_zero(NULL, NULL)") # 0
query("SELECT add_or_zero(3, 7)")       # 10

Обработка исключений

Параметр on_error определяет поведение, если функция Python вызывает исключение.

Значение	Поведение
`"propagate"` (по умолчанию)	Возбудить исключение как ошибку SQL
`"ignore"`	Перехватить исключение и вернуть NULL для этой строки

Также можно использовать enum: chdb.ExceptionHandling.PROPAGATE / chdb.ExceptionHandling.IGNORE.

Пример: default (передать дальше)

@func(arg_types=["Int64", "Int64"], return_type="Int64")
def divide(a, b):
    return a // b

query("SELECT divide(10, 2)")  # 5
query("SELECT divide(1, 0)")   # Error: ZeroDivisionError

Пример: игнорирование ошибок

@func(arg_types=["Int64", "Int64"], return_type="Int64", on_error="ignore")
def safe_divide(a, b):
    return a // b

query("SELECT safe_divide(10, 2)")  # 5
query("SELECT safe_divide(1, 0)")   # NULL

Совместная обработка NULL и исключений

Параметры on_null и on_error можно комбинировать:

on_null	on_error	Входной NULL	Исключение
`"skip"`	`"propagate"`	Вернуть NULL	Сгенерировать ошибку
`"skip"`	`"ignore"`	Вернуть NULL	Вернуть NULL
`"pass"`	`"propagate"`	Вызвать с `None`	Сгенерировать ошибку
`"pass"`	`"ignore"`	Вызвать с `None`	Вернуть NULL

@func(
    arg_types=["Int64", "Int64"],
    return_type="Int64",
    on_null="pass",
    on_error="ignore",
)
def robust_divide(a, b):
    if a is None or b is None:
        return -1
    return a // b

query("SELECT robust_divide(10, 2)")     # 5
query("SELECT robust_divide(NULL, 2)")   # -1
query("SELECT robust_divide(1, 0)")      # NULL (exception caught)

Поддержка DateTime и часовых поясов

Пользовательские функции (UDF) полностью поддерживают типы даты и времени с поддержкой часовых поясов.

Типы данных Date

from datetime import date, timedelta

@func()
def next_day(d: date) -> date:
    return d + timedelta(days=1)

@func()
def get_year(d: date) -> int:
    return d.year

query("SELECT next_day(toDate('2024-06-15'))")  # 2024-06-16
query("SELECT get_year(toDate('2024-06-15'))")  # 2024

DateTime с часовыми поясами

from datetime import timedelta

@func(arg_types=["DateTime('UTC')"], return_type="DateTime('UTC')")
def add_one_hour(dt):
    return dt + timedelta(hours=1)

query("SELECT add_one_hour(toDateTime('2024-01-01 12:00:00', 'UTC'))")  # 2024-01-01 13:00:00

DateTime64 (высокая точность)

Для DATETIME64 по умолчанию задан scale 6 (микросекунды):

from datetime import timedelta

@func(arg_types=["DateTime64(6, 'UTC')"], return_type="DateTime64(6, 'UTC')")
def add_microsecond(dt):
    return dt + timedelta(microseconds=1)

query("SELECT add_microsecond(toDateTime64('2024-01-01 12:00:00.000000', 6, 'UTC'))")  # 2024-01-01 12:00:00.000001

Использование пользовательских функций (UDF) в сеансах

Пользовательские функции (UDF) регистрируются глобально и доступны во всех сеансах в рамках одного процесса:

from chdb import session as chs, func
from chdb.sqltypes import INT64

@func([INT64], INT64)
def double(x):
    return x * 2

sess = chs.Session()
sess.query("CREATE TABLE t (x Int64) ENGINE = Memory")
sess.query("INSERT INTO t VALUES (1), (2), (3)")
result = sess.query("SELECT double(x) FROM t ORDER BY x", "CSV")
print(result)
# 2
# 4
# 6