用户定义函数 (UDFs)¶

虽然Pandas和其他API可以非常表达丰富，但许多数据科学和数据工程的使用案例需要超出直接提供的额外功能。在这些情况下，许多程序员创建用户定义函数，或称UDF，这些是根据上下文设计用于对每一行或一组行进行计算的Python函数。

在Bodo中使用UDF¶

Bodo 用户可以通过定义一个单独的 JIT 函数或在 JIT 函数中创建一个函数（通过 lambda 或闭包）来构造 UDF。例如，这里有两种构造 UDF 的方法，可以将时间戳序列的每个元素推进到当前月份的最后一天。

import pandas as pd
import bodo

@bodo.jit
def jit_udf(x):
    return x + pd.tseries.offsets.MonthEnd(n=0, normalize=True)

@bodo.jit
def jit_example(S):
    return S.map(jit_udf)

@bodo.jit
def lambda_example(S):
    return S.map(lambda x: x + pd.tseries.offsets.MonthEnd(n=0, normalize=True))

S = pd.Series(pd.date_range(start='1/1/2021', periods=100))
pd.testing.assert_series_equal(jit_example(S), lambda_example(S))

UDFs 可以用于每行或每组计算一个值（映射函数）或计算聚合（agg 函数）。Bodo 提供了两者的 API，下面进行了总结。详请请参见 supported Pandas API。

映射函数¶

Series.map
Series.apply
Series.pipe
DataFrame.map
DataFrame.apply
DataFrame.pipe
GroupBy.apply
GroupBy.pipe
GroupBy.transform

聚合函数¶

GroupBy.agg
GroupBy.aggregate

用户定义函数性能¶

Bodo 支持 UDF，而不会像在 Pandas 中那样产生显著的运行时惩罚。一个例子在快速入门指南中展示。

Bodo 在用户定义函数（UDFs）上获得了巨大的性能优势，因为 UDFs 可以像其他任何 JIT 代码一样进行优化。相比之下，基于库的解决方案在优化 UDFs 的能力上通常受到限制。

附加参数¶

我们建议将额外的变量显式传递给 UDF，而不是直接使用定义 UDF 的函数本地变量。后者称为“捕获”变量的情况，这通常容易出错，并可能导致编译错误。

例如，考虑一个用户定义的函数（UDF），它在字符串序列中的每个字符串后附加一个变量后缀。编写这个函数的正确方法是使用 args 参数来 Series.apply()。

import pandas as pd
import bodo

@bodo.jit
def add_suffix(S, suffix):
    return S.apply(lambda x, suf: x + suf, args=(suffix,))

S = pd.Series(["abc", "edf", "32", "Vew3", "er3r2"] * 10)
suffix = "_"
add_suffix(S, suffix)

另外，参数可以按照关键字传递。

@bodo.jit
def add_suffix(S, suffix):
    return S.apply(lambda x, suf: x + suf, suf=suffix)

注意

并非所有API都支持额外的参数。有关预期API使用的更多信息，请参考受支持的Pandas API。

使用Pandas方法和Numpy ufuncs申请¶

除了用户定义函数（UDFs）之外，apply API 也可以用于调用 Pandas 方法和 Numpy ufuncs。要执行 Pandas 方法，您可以将方法名称作为字符串提供。

import pandas as pd
import bodo

@bodo.jit
def ex(S):
    return S.apply("nunique")

S = pd.Series(list(np.arange(100) + list(np.arange(100))))
ex(S)

Numpy通用函数可以通过匹配名称的字符串或函数本身来提供。

import numpy as np
import pandas as pd
import bodo

@bodo.jit
def ex_str(S):
    return S.apply("sin")

def ex_func(S):
    return S.apply(np.sin)

S = pd.Series(list(np.arange(100) + list(np.arange(100))))
pd.testing.assert_series_equal(ex_str(S), ex_func(S))

注意

Numpy ufuncs 目前不支持 DataFrames。

类型稳定性限制¶

Bodo 的类型稳定性要求在使用 DataFrame.apply 处理不同列类型或返回 DataFrame 时可能会遇到一些限制。

不同类型的列¶

DataFrame.apply 将用户提供的 UDF 应用到 DataFrame 的每一行。在 DataFrame 具有不同类型的列的情况下，传递给 UDF 的 Series 将包含不同类型的值。Bodo 在内部将这些表示为异构 Series。此表示在其支持的 Series 操作中存在限制。有关更多信息，请参阅 supported operations for heterogeneous series。

返回一个DataFrame¶

在Pandas中，Series.apply 或 DataFrame.apply 有多种方式可以返回一个DataFrame而不是Series。然而，出于类型稳定性的原因，Bodo只能在返回每一行的大小可以在编译时推断的Series时推断出一个DataFrame。

注意

如果您提供一个索引，则所有索引值必须是编译时常量。

这是一个使用Series.apply返回DataFrame的示例。

import pandas as pd
import bodo

@bodo.jit
def series_ex(S):
    return S.apply(lambda x: pd.Series((1, x)))

S = pd.Series(list(np.arange(100) + list(np.arange(100))))
series_ex(S)

如果使用通过其他方式返回DataFrame的UDF，这种行为在Bodo中可能不匹配，并可能导致编译错误。如果可能，请将您的解决方案转换为支持的方法之一。