Python 函数 API – DuckDB - DuckDB 数据库

搜索快捷键 cmd + k | ctrl + k

文档 / 客户端API / Python

Python 函数 API

你可以通过 Python 函数创建 DuckDB 用户自定义函数 (UDF)，从而在 SQL 查询中使用它们。与普通函数类似，它们需要具备名称、返回类型和参数类型。

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

import duckdb
from duckdb.sqltypes import VARCHAR
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 方法包含以下参数：

name 一个字符串，表示连接目录中 UDF 的唯一名称。
function 你希望注册为 UDF 的 Python 函数。
parameters 标量函数可以对一个或多个列进行操作。此参数接受一个用作输入的列类型列表。
return_type 标量函数每行返回一个元素。此参数指定函数的返回类型。
type（可选）：DuckDB 同时支持原生 Python 类型和 PyArrow 数组。默认情况下，假设为 type = 'native'，但你可以指定 type = 'arrow' 来使用 PyArrow 数组。通常，使用 Arrow UDF 比原生 UDF 高效得多，因为它能够以批处理方式进行操作。
null_handling（可选）：默认情况下，NULL 值会自动按照“输入为 NULL，则输出为 NULL”的逻辑处理。用户可以通过设置 null_handling = 'special' 来指定对 NULL 值的自定义行为。
exception_handling（可选）：默认情况下，当 Python 函数抛出异常时，该异常会重新在 Python 中抛出。用户可以通过将此参数设置为 'return_null' 来禁用此行为，改为返回 NULL。
side_effects（可选）：默认情况下，函数被视为对于相同的输入会产生相同的结果。如果函数的结果受到任何形式的随机性影响，则必须将 side_effects 设置为 True。

要注销 UDF，可以使用 UDF 名称调用 remove_function 方法：

con.remove_function(name)

使用偏函数 (Partial Functions)

DuckDB UDF 也可以使用 Python 偏函数创建。

在下面的示例中，我们展示了自定义记录器如何返回 ISO 格式的执行日期时间拼接字符串，随后总是跟上在 UDF 创建时传递的参数以及函数调用时提供的输入参数。

from datetime import datetime
import duckdb
import functools


def get_datetime_iso_format() -> str:
    return datetime.now().isoformat()


def logger_udf(func, arg1: str, arg2: int) -> str:
    return ' '.join([func(), arg1, str(arg2)])
    
    
with duckdb.connect() as con:
    con.sql("select * from range(10) tbl(id)").to_table("example_table")
    
    con.create_function(
        'custom_logger',
        functools.partial(logger_udf, get_datetime_iso_format, 'logging data')
    )
    rel = con.sql("SELECT custom_logger(id) from example_table;")
    rel.show()

    con.create_function(
        'another_custom_logger',
        functools.partial(logger_udf, get_datetime_iso_format, ':')
    )
    rel = con.sql("SELECT another_custom_logger(id) from example_table;")
    rel.show()

┌───────────────────────────────────────────┐
│             custom_logger(id)             │
│                  varchar                  │
├───────────────────────────────────────────┤
│ 2025-03-27T12:07:56.811251 logging data 0 │
│ 2025-03-27T12:07:56.811264 logging data 1 │
│ 2025-03-27T12:07:56.811266 logging data 2 │
│ 2025-03-27T12:07:56.811268 logging data 3 │
│ 2025-03-27T12:07:56.811269 logging data 4 │
│ 2025-03-27T12:07:56.811270 logging data 5 │
│ 2025-03-27T12:07:56.811271 logging data 6 │
│ 2025-03-27T12:07:56.811272 logging data 7 │
│ 2025-03-27T12:07:56.811274 logging data 8 │
│ 2025-03-27T12:07:56.811275 logging data 9 │
├───────────────────────────────────────────┤
│                  10 rows                  │
└───────────────────────────────────────────┘

┌────────────────────────────────┐
│   another_custom_logger(id)    │
│            varchar             │
├────────────────────────────────┤
│ 2025-03-27T12:07:56.812106 : 0 │
│ 2025-03-27T12:07:56.812116 : 1 │
│ 2025-03-27T12:07:56.812118 : 2 │
│ 2025-03-27T12:07:56.812119 : 3 │
│ 2025-03-27T12:07:56.812121 : 4 │
│ 2025-03-27T12:07:56.812122 : 5 │
│ 2025-03-27T12:07:56.812123 : 6 │
│ 2025-03-27T12:07:56.812124 : 7 │
│ 2025-03-27T12:07:56.812126 : 8 │
│ 2025-03-27T12:07:56.812127 : 9 │
├────────────────────────────────┤
│            10 rows             │
└────────────────────────────────┘

类型注解

当函数具有类型注解时，通常可以省略所有可选参数。使用 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          │
└─────────────┘

如果只能推断参数列表类型，则需要将 parameters 设置为 None。

`NULL` 处理

默认情况下，当函数接收到 NULL 值时，作为默认 NULL 处理的一部分，它会直接返回 NULL。当这不是你想要的预期行为时，需要显式地将此参数设置为 "special"。

import duckdb
from duckdb.sqltypes import BIGINT

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.sqltypes import BIGINT

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,)]

当函数可能返回 NULL 时，请务必使用 null_handling="special"。

import duckdb
from duckdb.sqltypes import VARCHAR


def return_str_or_none(x: str) -> str | None:
    if not x:
        return None
    
    return x

duckdb.create_function(
    "return_str_or_none",
    return_str_or_none,
    [VARCHAR],
    VARCHAR,
    null_handling="special"
)
res = duckdb.sql("SELECT return_str_or_none('')").fetchall()
print(res)

[(None,)]

异常处理

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

import duckdb
from duckdb.sqltypes import BIGINT

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,)]

副作用 (Side Effects)

默认情况下，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

如果预期函数接收 arrow 数组，请将 type 参数设置为 'arrow'。

这将告知系统向函数提供最多 STANDARD_VECTOR_SIZE 个元组的 arrow 数组，并期望函数返回相同数量元组的数组。

通常，使用 Arrow UDF 比原生 UDF 高效得多，因为它能够以批处理方式进行操作。

import duckdb
import pyarrow as pa
from duckdb.sqltypes import VARCHAR
from pyarrow import compute as pc


def mirror(strings: pa.Array, sep: pa.Array) -> pa.Array:
    assert isinstance(strings, pa.ChunkedArray)
    assert isinstance(sep, pa.ChunkedArray)
    return pc.binary_join_element_wise(strings, pc.ascii_reverse(strings), sep)


duckdb.create_function(
    "mirror",
    mirror,
    [VARCHAR, VARCHAR],
    return_type=VARCHAR,
    type="arrow",
)

duckdb.sql(
    "CREATE OR REPLACE TABLE strings AS SELECT 'hello' AS str UNION ALL SELECT 'world' AS str;"
)
print(duckdb.sql("SELECT mirror(str, '|') FROM strings;").fetchall())

[('hello|olleh',), ('world|dlrow',)]

Native（原生）

当函数类型设置为 native 时，函数将一次接收一个元组，并期望只返回单个值。这对于与不支持 Arrow 操作的 Python 库（例如 faker）交互非常有用。

import duckdb

from duckdb.sqltypes import DATE
from faker import Faker

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

duckdb.create_function(
    "random_date",
    random_date,
    parameters=[],
    return_type=DATE,
    type="native",
)
res = duckdb.sql("SELECT random_date()").fetchall()
print(res)

[(datetime.date(2019, 5, 15),)]

关于此页面

行为准则商标使用指南

关于此页面

本文中