如何在vaex中注册支持额外参数的自定义函数？

Vaex原生支持注册带额外参数的自定义函数，不需要通过df.apply实现，直接按如下方式开发即可，全程走Vaex原生优化链路，支持JIT编译，性能远高于外部调用的apply写法。

正确实现示例 #

1. 函数注册 #

注意修正原示例中的变量名笔误（原代码中txt应为输入参数field）：

import vaex

@vaex.register_function()
def abc(field, arg2):
    split_arg2 = set(arg2.split(' '))
    # 用Vaex矢量化操作实现，避免纯Python遍历
    return field.str.split(' ').apply(lambda x: len(set(x) & split_arg2))

如果是数值类计算逻辑，可额外加@vaex.jit装饰器触发JIT编译，进一步提升性能：

import vaex
from vaex import jit

@vaex.register_function()
@jit
def 数值计算函数(field, arg2):
    # 仅支持numba可编译的逻辑
    return field * arg2 + 10

2. 调用方式 #

完全符合预期的调用写法：

# 构造测试数据
df = vaex.from_dict({"field": ["a b c", "d e f", "a e c"]})

# 直接传入额外参数即可
result = df.func.field.abc("a c f")
print(result)
# 输出结果为 [2, 1, 2]

额外说明 #

• 额外参数既可以传固定字面量，也可以传入DataFrame的其他列表达式，比如df.func.field.abc(df.另一列名)即可实现按行匹配不同的arg2值
• 函数定义时第一个参数固定为调用时df.func.xxx对应的列/表达式，后续自定义参数按正常顺序声明即可

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