pyarrow.compute.register_vector_function

pyarrow.compute.register_vector_function(func, function_name, function_doc, in_types, out_type, func_registry=None)

注册一个用户自定义的向量函数。

此 API 是实验性的。

向量函数是一种在数组上执行向量化操作的函数。当计算不适合其他更具体的函数类型（例如标量函数和聚合函数）时，通常会使用向量函数。

参数:

funccallable()

实现该用户自定义函数的可调用对象。第一个参数是 UdfContext 类型的上下文参数。随后，它必须接受数量等于 in_types 定义的参数。它必须返回匹配 out_type 的 Array 或 Scalar。如果所有参数均为标量，则必须返回一个 Scalar，否则必须返回一个 Array。

要定义一个可变参数函数，请传入一个接受 *args 的可调用对象。最后一个 in_type 将作为所有可变参数的类型。

function_namestr

函数名称。在函数注册表中，每个名称只能注册一个函数。

function_docdict

包含键 “summary”（字符串）和 “description”（字符串）的字典对象。

in_typesDict[str, DataType]

一个将函数参数名称映射到其各自 DataType 的字典。参数名称将用于生成函数文档。此处指定的参数数量决定了函数的元数（arity）。

out_typeDataType

函数的输出类型。

func_registryFunctionRegistry

可选的函数注册表，用于替代默认的全局注册表。

示例

>>> import pyarrow as pa
>>> import pyarrow.compute as pc
>>>
>>> func_doc = {}
>>> func_doc["summary"] = "percent rank"
>>> func_doc["description"] = "compute percent rank"
>>>
>>> def list_flatten_udf(ctx, x):
...     return pc.list_flatten(x)
>>>
>>> func_name = "list_flatten_udf"
>>> in_types = {"array": pa.list_(pa.int64())}
>>> out_type = pa.int64()
>>> pc.register_vector_function(list_flatten_udf, func_name, func_doc,
...                   in_types, out_type)
>>>
>>> answer = pc.call_function(func_name, [pa.array([[1, 2], [3, 4]])])
>>> answer
<pyarrow.lib.Int64Array object at ...>
[
  1,
  2,
  3,
  4
]