我需要基于两列 DataFrame 生成新列，如何才能更快？

I need to generate new column based on two columns of dataframe, how can it be faster?

提问人：jaried 提问时间：6/10/2022 最后编辑：jaried 更新时间：6/11/2022 访问量：97

问：

我需要根据列和列生成列，如果两者都大于 0，则分配值 1，如果两者且小于 0，则分配值 -1，我使用 double 。a_babdfaba_baba_bnp.where

我的代码如下，其中生成并用于，需要运行的地方：generate_datademo dataget_resultproductionget_result4 million times

import numpy as np
import pandas as pd

rand = np.random.default_rng(seed=0)
pd.set_option('display.max_columns', None)


def generate_data() -> pd.DataFrame:
    _df = pd.DataFrame(rand.uniform(-1, 1, size=(10,7)), columns=['a', 'b1', 'b2', 'b3', 'b4', 'b5', 'b6'])
    return _df


def get_result(_df: pd.DataFrame) -> pd.DataFrame:
    a = _df.a.to_numpy()
    for col in ['b1', 'b2', 'b3', 'b4', 'b5', 'b6']:
        b = _df[col].to_numpy()
        _df[f'a_{col}'] = np.where(
            (a > 0) & (b > 0), 1., np.where(
                (a < 0) & (b < 0), -1., 0.)
        )
    return _df


def main():
    df = generate_data()
    print(df)
    df = get_result(df)
    print(df)


if __name__ == '__main__':
    main()

generate_data生成的数据：

          a        b1        b2        b3        b4        b5        b6
0  0.273923 -0.460427 -0.918053 -0.966945  0.626540  0.825511  0.213272
1  0.458993  0.087250  0.870145  0.631707 -0.994523  0.714809 -0.932829
2  0.459311 -0.648689  0.726358  0.082922 -0.400576 -0.154626 -0.943361
3 -0.751433  0.341249  0.294379  0.230770 -0.232645  0.994420  0.961671
4  0.371084  0.300919  0.376893 -0.222157 -0.729807  0.442977  0.050709
5 -0.379516 -0.028329  0.778976  0.868087 -0.284410  0.143060 -0.356261
6  0.188600 -0.324178 -0.216762  0.780549 -0.545685  0.246374 -0.831969
7  0.665288  0.574197 -0.521261  0.752968 -0.882864 -0.327766 -0.699441
8 -0.099321  0.592649 -0.538716 -0.895957 -0.190896 -0.602974 -0.818494
9  0.160665 -0.402608  0.343990 -0.600969  0.884226 -0.269780 -0.789009

我想要的结果：


          a        b1        b2        b3        b4        b5        b6  a_b1  \
0  0.273923 -0.460427 -0.918053 -0.966945  0.626540  0.825511  0.213272   0.0   
1  0.458993  0.087250  0.870145  0.631707 -0.994523  0.714809 -0.932829   1.0   
2  0.459311 -0.648689  0.726358  0.082922 -0.400576 -0.154626 -0.943361   0.0   
3 -0.751433  0.341249  0.294379  0.230770 -0.232645  0.994420  0.961671   0.0   
4  0.371084  0.300919  0.376893 -0.222157 -0.729807  0.442977  0.050709   1.0   
5 -0.379516 -0.028329  0.778976  0.868087 -0.284410  0.143060 -0.356261  -1.0   
6  0.188600 -0.324178 -0.216762  0.780549 -0.545685  0.246374 -0.831969   0.0   
7  0.665288  0.574197 -0.521261  0.752968 -0.882864 -0.327766 -0.699441   1.0   
8 -0.099321  0.592649 -0.538716 -0.895957 -0.190896 -0.602974 -0.818494   0.0   
9  0.160665 -0.402608  0.343990 -0.600969  0.884226 -0.269780 -0.789009   0.0   

   a_b2  a_b3  a_b4  a_b5  a_b6  
0   0.0   0.0   1.0   1.0   1.0  
1   1.0   1.0   0.0   1.0   0.0  
2   1.0   1.0   0.0   0.0   0.0  
3   0.0   0.0  -1.0   0.0   0.0  
4   1.0   0.0   0.0   1.0   1.0  
5   0.0   0.0  -1.0   0.0  -1.0  
6   0.0   1.0   0.0   1.0   0.0  
7   0.0   1.0   0.0   0.0   0.0  
8  -1.0  -1.0  -1.0  -1.0  -1.0  
9   1.0   0.0   1.0   0.0   0.0

绩效评估：

%timeit get_result(df)
1.56 ms ± 54.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

怎么可能更快？

Python Pandas DataFrame numpy 性能

评论

1赞 SultanOrazbayev 6/10/2022
真实数据的形态是什么？在一小部分上进行基准测试可能会产生误导，因此在最终用例中了解数据大小会很棒。
0赞 jaried 6/10/2022
@SultanOrazbayev 实数的形状也是 （10,7）

答：

2赞 SultanOrazbayev 6/10/2022 #1

对于较小的数据帧 （10,7），矢量化的收益很少，因此我不确定在那里可以获得多少收益。但是，您可以重写代码以使其更具可读性（尽管这可能是主观的）：

def get_result2(_df: pd.DataFrame) -> pd.DataFrame:
    
    bcols = [c for c in _df.columns if c.startswith('b')]
    bcols_names = [f'a_{c}' for c in bcols]

    a_sign = np.sign(df['a']).values.reshape(-1,1)
    b_signs = np.sign(df[bcols])

    _df[bcols_names] = ( b_signs == a_sign ) * a_sign

    return _df

您可以使用以下方法检查这是否给出了相同的结果：

x = get_result(df)
y = get_result2(df)

print(x.equals(y))
# True

但是，在我的测试中，此函数不会在运行时中产生一致的改进。我猜它可能更适合更大的数据集。

评论

1赞 jaried 6/10/2022
3.29 ms ± 205 μs （平均±标准开发 7 次，每次 100 次循环）
0赞 jaried 6/10/2022
这种方法对于纯 numpy 来说更快，请参阅我的代码答案。
0赞 SultanOrazbayev 6/10/2022
当然，这个问题并没有明确说明一个纯粹的麻木解决方案会起作用。
0赞 jaried 6/10/2022
由于我需要使用 df，我接受另一个答案。我也投了这个答案。
2赞 I'mahdi 6/10/2022 #2

因为你标记了性能，我建议你，使用和并行计算，如下所示：（如果我们直接输入值到并行函数，我们可以达到 3.35 μs）numba

import numpy as np
import numba as nb
import pandas as pd


@nb.njit( parallel=True )
def parallel_fun(vals):
    a = vals[:,0]
    new_vals = np.empty((10,6))
    for i in nb.prange(6):
        b = vals[:,i+1]
        for j in nb.prange(10):
            val = 0
            if (a[j] >0) and (b[j]>0): val =1
            elif (a[j] <0) and (b[j]<0) : val= -1
            new_vals[j,i] = val
    return new_vals

def get_result_3(_df: pd.DataFrame) -> pd.DataFrame:
    vals = _df[['a','b1', 'b2', 'b3', 'b4', 'b5', 'b6']].to_numpy()
    new_vals = parallel_fun(vals)
    return pd.DataFrame(new_vals, columns=[f'a_{b}' for b in ['b1', 'b2', 'b3', 'b4', 'b5', 'b6']])


_df = pd.DataFrame(np.random.uniform(-1, 1, size=(10,7)), columns=['a', 'b1', 'b2', 'b3', 'b4', 'b5', 'b6'])
vals = _df[['a','b1', 'b2', 'b3', 'b4', 'b5', 'b6']].to_numpy()

colab 基准测试：

%timeit get_result_3(_df)
# 658 µs per loop
%timeit parallel_fun(vals)
# 3.35 µs per loop

评论

1赞 jaried 6/10/2022
每个环路 1.11 ms ± 42.3 μs（平均±标准开发，7 次运行，每次 1000 个循环）
0赞 I'mahdi 6/10/2022
@jaried，你跑到哪里去或检查过？
1赞 jaried 6/10/2022
在我自己的电脑上查看。
1赞 I'mahdi 6/10/2022
@jaried，好吧，祝你;)
2赞 Jérôme Richard 6/10/2022
请注意，这并不是一个好主意，充其量是没有用的。它用于并行化操作，但如果没有 .无论如何，编译器都会展开循环。顺便说一句，大部分开销应该来自 Pandas，它创建小数据帧的速度非常慢（仅 Numba 代码就应该不超过我们几个）。另请注意，编译时间会减慢函数的首次执行速度。nb.prangeparallel=True
1赞 jaried 6/10/2022 #3

有人用纯粹的麻木回答我：

import numpy as np

rand = np.random.default_rng(seed=0)
a = rand.uniform(low=-1, high=1, size=(10, 1))
b = rand.uniform(low=-1, high=1, size=(10, 6))


def signs():
    sa = np.sign(a)
    return sa * (sa == np.sign(b))


def main():
    signs()
    return


if __name__ == '__main__':
    main()


%timeit signs()
10.2 µs ± 678 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

上一个：为什么子类的浮点变量与父类的浮点变量是相同的内存地址？

下一个：如何进行 DataFrame 列的分类和计数？[复制]