如何在DuckDB中创建操作PyArrow对象的简单UDF？

错误原因拆解 #

• type参数格式错误：create_function要求type传入duckdb.PythonUDFType枚举值，不是字符串'arrow'
• 类型定义不兼容：parameters和return_type不能用Python原生的int/float，得用DuckDB内置类型或者PyArrow类型
• SQL语法错误：update语句里多写了table关键字，正确写法是update t set b = test(a)

修正后的完整代码 #

import pyarrow as pa
import duckdb
from duckdb import PythonUDFType

def test(x: pa.Array) -> pa.Array:
    # 先把整数数组转成浮点数，避免PyArrow整数类型调用sqrt报错
    return x.cast(pa.float64()).compute.sqrt()

duckdb.create_function(
    'test',
    test,
    type=PythonUDFType.ARROW,
    parameters=[pa.int64()],  # 也可以用 duckdb.types.INTEGER
    return_type=pa.float64(), # 也可以用 duckdb.types.DOUBLE
)

N = 1000
duckdb.sql(f"""
    create table t (a int);
    insert into t (a) select * from generate_series(1, {N});
    alter table t add column b float;
    update t set b = test(a);
""")

# 验证结果示例
print(duckdb.sql("select a, b from t limit 5").fetchall())

核心修改说明 #

1. 引入PythonUDFType枚举，明确指定UDF类型为Arrow批处理模式
2. 用PyArrow类型定义参数和返回值，保证和DuckDB的类型系统匹配
3. 函数内部做类型转换，因为PyArrow的整数数组不能直接调用sqrt方法
4. 修正SQL语句里的语法错误

内容的提问来源于stack exchange，提问作者Test