⌘+k ctrl+k
1.1.3 (stable)
Search Shortcut cmd + k | ctrl + k
Python Function API

你可以从Python函数创建一个DuckDB用户定义函数(UDF),以便它可以在SQL查询中使用。 与常规的函数类似,它们需要有一个名称、返回类型和参数类型。

这是一个使用调用第三方库的Python函数的示例。

import duckdb
from duckdb.typing import *
from faker import Faker

def generate_random_name():
    fake = Faker()
    return fake.name()

duckdb.create_function("random_name", generate_random_name, [], VARCHAR)
res = duckdb.sql("SELECT random_name()").fetchall()
print(res)
[('Gerald Ashley',)]

创建函数

要注册一个Python UDF,请使用DuckDB连接中的create_function方法。以下是语法:

import duckdb
con = duckdb.connect()
con.create_function(name, function, parameters, return_type)

create_function 方法接受以下参数:

  1. name 一个字符串,表示连接目录中UDF的唯一名称。
  2. function 您希望注册为UDF的Python函数。
  3. parameters 标量函数可以操作一个或多个列。此参数接受用作输入的列类型列表。
  4. return_type 标量函数每行返回一个元素。此参数指定函数的返回类型。
  5. type(可选):DuckDB 支持内置的 Python 类型和 PyArrow 表。默认情况下,假定使用内置类型,但您可以指定 type = 'arrow' 来使用 PyArrow 表。
  6. null_handling (可选): 默认情况下,NULL 值会自动处理为 NULL-in NULL-out。用户可以通过设置 null_handling = 'special' 来指定 NULL 值的期望行为。
  7. exception_handling (可选): 默认情况下,当从Python函数抛出异常时,它将在Python中重新抛出。用户可以通过将此参数设置为'return_null'来禁用此行为,并返回NULL
  8. side_effects(可选):默认情况下,函数预期在相同输入下产生相同的结果。如果函数的结果受到任何类型的随机性影响,side_effects 必须设置为 True

要注销一个UDF,你可以调用remove_function方法并传入UDF名称:

con.remove_function(name)

类型注解

当函数有类型注解时,通常可以省略所有可选参数。 使用DuckDBPyType我们可以隐式地将许多已知类型转换为DuckDB的类型系统。 例如:

import duckdb

def my_function(x: int) -> str:
    return x

duckdb.create_function("my_func", my_function)
print(duckdb.sql("SELECT my_func(42)"))
┌─────────────┐
│ my_func(42) │
│   varchar   │
├─────────────┤
│ 42          │
└─────────────┘

如果只能推断出参数列表类型,你需要将None作为parameters传入。

NULL 处理

默认情况下,当函数接收到NULL值时,它会立即返回NULL,这是默认的NULL处理的一部分。 如果不希望这样,你需要显式地将此参数设置为"special"

import duckdb
from duckdb.typing import *

def dont_intercept_null(x):
    return 5

duckdb.create_function("dont_intercept", dont_intercept_null, [BIGINT], BIGINT)
res = duckdb.sql("SELECT dont_intercept(NULL)").fetchall()
print(res)
[(None,)]

使用 null_handling="special":

import duckdb
from duckdb.typing import *

def dont_intercept_null(x):
    return 5

duckdb.create_function("dont_intercept", dont_intercept_null, [BIGINT], BIGINT, null_handling="special")
res = duckdb.sql("SELECT dont_intercept(NULL)").fetchall()
print(res)
[(5,)]

异常处理

默认情况下,当从Python函数抛出异常时,我们会转发(重新抛出)该异常。 如果你想禁用此行为,并改为返回NULL,你需要将此参数设置为"return_null"

import duckdb
from duckdb.typing import *

def will_throw():
    raise ValueError("ERROR")

duckdb.create_function("throws", will_throw, [], BIGINT)
try:
    res = duckdb.sql("SELECT throws()").fetchall()
except duckdb.InvalidInputException as e:
    print(e)

duckdb.create_function("doesnt_throw", will_throw, [], BIGINT, exception_handling="return_null")
res = duckdb.sql("SELECT doesnt_throw()").fetchall()
print(res)
Invalid Input Error: Python exception occurred while executing the UDF: ValueError: ERROR

At:
  ...(5): will_throw
  ...(9): <module>
[(None,)]

副作用

默认情况下,DuckDB会假设创建的函数是一个函数,这意味着在给定相同输入时,它将产生相同的输出。 如果你的函数不遵循这个规则,例如当你的函数使用了随机性时,那么你需要将这个函数标记为具有side_effects

例如,这个函数每次调用都会生成一个新的计数。

def count() -> int:
    old = count.counter;
    count.counter += 1
    return old

count.counter = 0

如果我们创建这个函数而不将其标记为具有副作用,结果将是以下内容:

con = duckdb.connect()
con.create_function("my_counter", count, side_effects=False)
res = con.sql("SELECT my_counter() FROM range(10)").fetchall()
print(res)
[(0,), (0,), (0,), (0,), (0,), (0,), (0,), (0,), (0,), (0,)]

这显然不是我们想要的结果,当我们添加side_effects=True时,结果如我们所预期的那样:

con.remove_function("my_counter")
count.counter = 0
con.create_function("my_counter", count, side_effects=True)
res = con.sql("SELECT my_counter() FROM range(10)").fetchall()
print(res)
[(0,), (1,), (2,), (3,), (4,), (5,), (6,), (7,), (8,), (9,)]

Python 函数类型

目前支持两种函数类型,native(默认)和arrow

Arrow

如果函数预期接收箭头数组,请将type参数设置为'arrow'

这将让系统知道向函数提供最多STANDARD_VECTOR_SIZE元组的箭头数组,并且期望从函数返回相同数量的元组数组。

Native

当函数类型设置为native时,函数将一次提供一个单一的元组,并且只期望返回一个单一的值。 这对于与不操作Arrow的Python库(如faker)进行交互非常有用:

import duckdb

from duckdb.typing import *
from faker import Faker

def random_date():
    fake = Faker()
    return fake.date_between()

duckdb.create_function("random_date", random_date, [], DATE, type="native")
res = duckdb.sql("SELECT random_date()").fetchall()
print(res)
[(datetime.date(2019, 5, 15),)]