Funciones definidas por el usuario (UDF) de Python

chDB permite registrar funciones de Python como UDF invocables desde SQL. Se ejecutan de forma nativa en el mismo proceso, sin iniciar subprocesos ni añadir sobrecarga de serialización. Las funciones tienen tipado seguro, admiten la inferencia automática de tipos a partir de anotaciones de Python y permiten configurar el manejo de NULL y de excepciones.

Inicio rápido

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

Métodos de registro

Decorador `@func`

La forma más sencilla de registrar una UDF. El atributo __name__ de la función se convierte en el nombre de la función 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}!"

La función decorada sigue pudiendo invocarse como una función normal de Python:

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

`create_function`

Registre cualquier objeto invocable (lambda, función, método) con un nombre explícito:

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`

Elimina una UDF registrada. Eliminar un nombre que no está registrado no tiene ningún efecto, por lo que se puede llamar incondicionalmente de forma segura:

from chdb import drop_function

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

Sistema de tipos

Tipos disponibles

Todos los tipos se pueden importar desde 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,
)

Especificar tipos

Los tipos se pueden proporcionar de cuatro formas:

Método	Ejemplo	Descripción
constante `ChdbType`	`INT64`, `STRING`	Se importa desde `chdb.sqltypes`
cadena de tipo ClickHouse	`"Int64"`, `"String"`	Nombres de tipo estándar de ClickHouse
cadena parametrizada	`"DateTime('UTC')"`, `"DateTime64(6)"`	Para tipos con parámetros
tipo de Python	`int`, `str`, `float`	Se pasa directamente en `arg_types`/`return_type` o se usa como anotación de tipo en la firma de la función

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

Inferencia automática de tipos

Cuando se omiten arg_types o return_type, chDB infiere los tipos a partir de las anotaciones de tipo de Python:

Tipo de Python	Tipo de 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)

Siempre se requiere un tipo de retorno: si se omite return_type y la función no tiene una anotación de retorno, se produce un error al registrarla. En cambio, los tipos de los argumentos son opcionales: un parámetro sin tipo explícito ni anotación acepta dinámicamente cualquier tipo de entrada compatible.

Manejo de NULL

El parámetro on_null controla el comportamiento cuando alguno de los argumentos de entrada es NULL.

Valor	Comportamiento
`"skip"` (predeterminado)	Devuelve NULL inmediatamente sin llamar a la función
`"pass"`	Convierte NULL en `None` de Python y llama a la función normalmente

También puedes usar el enum: chdb.NullHandling.SKIP / chdb.NullHandling.PASS.

Ejemplo: predeterminado (omitir)

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

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

Ejemplo: pasar NULL como `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

Ejemplo: varios argumentos

@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

Manejo de excepciones

El parámetro on_error controla el comportamiento cuando la función de Python genera una excepción.

Valor	Comportamiento
`"propagate"` (predeterminado)	Propaga la excepción como un error SQL
`"ignore"`	Captura la excepción y devuelve NULL para esa fila

También puede usar el enum: chdb.ExceptionHandling.PROPAGATE / chdb.ExceptionHandling.IGNORE.

Ejemplo: predeterminado (propagar)

@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

Ejemplo: ignorar errores

@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

Combinación de NULL y manejo de excepciones

Las opciones on_null y on_error se pueden combinar:

on_null	on_error	Entrada NULL	Excepción
`"skip"`	`"propagate"`	Devuelve NULL	Genera un error
`"skip"`	`"ignore"`	Devuelve NULL	Devuelve NULL
`"pass"`	`"propagate"`	Llama con `None`	Genera un error
`"pass"`	`"ignore"`	Llama con `None`	Devuelve 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)

Compatibilidad con DateTime y zonas horarias

Las UDF admiten plenamente tipos de fecha y hora compatibles con zonas horarias.

Tipos de 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 con zonas horarias

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 (alta precisión)

DATETIME64 tiene una escala predeterminada de 6 (microsegundos):

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

Uso de UDFs con sesiones

Las UDFs se registran de forma global y están disponibles en todas las sesiones del mismo proceso:

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