python - 如何正确使用类似 numpy 的矢量化来加速 pandas 数据框应用函数中的复杂条件评估

Question

numpy/pandas 因其底层加速(即矢量化)而闻名。

条件评估是常见的表达式，出现在任何地方的代码中。

但是，直观地使用 pandas dataframe apply 函数时，条件评估似乎非常慢。

我的apply代码示例如下:

 def condition_eval(df):
        x=df['x']
        a=df['a']
        b=df['b']
        if x <= a:
            d = round((x-a)/0.01)-1
            if d <- 10:
                d = -10
        elif x >= b:
            d = round((x-b)/0.01)+1
            if d > 10:
                d = 10
        else:
            d = 0 
        return d
df['eval_result'] = df.apply(condition_eval, axis=1)

此类问题的特征可能是:

1. 仅使用其自己的行数据即可计算结果，并且始终使用多个列。
2. 每一行都有相同的计算算法。
3. 算法可能包含复杂的条件分支。

numpy/pandas 解决此类问题的最佳实践是什么？

---

还有一些想法。

在我看来，矢量化加速之所以有效的原因之一是因为底层CPU有某种矢量指令(例如SIMD、intel avx)，它依赖于一个事实:计算指令具有确定性行为，即无论输入数据如何，都可以在固定数量的CPU周期后获得结果。因此，并行化此类操作很容易。

然而，cpu中的分支执行要复杂得多。首先，相同条件评估的不同分支具有不同的执行路径，因此它们可能会导致不同的CPU周期。现代 CPU 甚至利用了很多技巧，例如分支预测，这会产生更多的不确定性。

所以我想知道 pandas 是否以及如何尝试加速此类向量条件评估操作，以及它们是否是处理此类计算工作负载的更好实践。

Answer 1

这应该是等效的:

import pandas as pd
import numpy as np

def get_eval_result(df):
    conditions = (
        df.x.le(df.a),
        df.x.gt(df.b),
    )
    choices = (
        np.where((d := df.x.sub(df.a).div(0.01).round().sub(1)).lt(-10), -10, d),
        np.where((d := df.x.sub(df.b).div(0.01).round().add(1)).gt(10), 10, d), 
    )
    return np.select(conditions, choices, 0)

df = df.assign(eval_result=get_eval_result)

我的答案基本上计算每个分支的结果，然后使用 numpy 语法来指定应使用哪些结果。这可以稍微优化，但由于它使用纯矢量化函数，它应该比使用 .apply 快得多。