Pandas的apply替代方案 - 基于多列创建新列

Question

我有一个Pandas数据帧，我想根据其他列的值添加新列。下面是一个说明我的用例的最小示例。

df = pd.DataFrame([[4,5,19],[1,2,0],[2,5,9],[8,2,5]], columns=['a','b','c'])
df

    a   b   c
---------------
0   4   5   19
1   1   2   0
2   2   5   9
3   8   2   5

x = df.sample(n=2)
x

    a   b   c
---------------
3   8   2   5
1   1   2   0

def get_new(row):
    a, b, c = row
    return random.choice(df[(df['a'] != a) & (df['b'] == b) & (df['c'] != c)]['c'].values)

y = x.apply(lambda row: get_new(row), axis=1)
x['new'] = y
x

    a   b   c   new
--------------------
3   8   2   5   0
1   1   2   0   5

注意：原始数据框包含约4百万行和6列。样本中的行数可能在50到500之间变化。我正在一台8 GB RAM的64位机器上运行。

以上方法可行，但速度相对较慢（我的情况需要大约15秒）。我还尝试使用x.itertuples()代替apply，但在这种情况下几乎没有改善。

以下是需要解决的问题：

1.看起来apply（使用axis = 1）很慢，因为它没有利用向量化操作。有没有更快的方法？

2.与使用条件布尔变量相比，能否修改或使筛选（在get_new函数中）更有效？

3.我能在某种程度上使用numpy来加速吗？

编辑：df.sample()也非常慢，而且我不能使用.iloc或.loc，因为我还要修改样本，并且不希望影响原始数据框。

- swathis

---

那么 dask 呢？ - Pierluigi

@Pierluigi - 我希望在查看其他库之前，Pandas或Numpy已经有更好（更有效）的方法来完成这个任务。但是感谢您指出这一点。我一定会去看看的。 - swathis

https://dev59.com/h2Qn5IYBdhLWcg3w5qlg#47149876 - Scott Boston

我创建了 tabel, 这是一个简化版的 pandas。它完全基于 numpy，但抽象化了处理多种数据类型的细节。而且在像你的 get/set 场景中，numpy 通常更快。https://pypi.org/project/tabel/ - Bastiaan

Answer 1

我发现使用.loc而不是链式索引可以获得合理的性能提升：

import random, pandas as pd, numpy as np

df = pd.DataFrame([[4,5,19],[1,2,0],[2,5,9],[8,2,5]], columns=['a','b','c'])

df = pd.concat([df]*1000000)

x = df.sample(n=2)

def get_new(row):
    a, b, c = row
    return random.choice(df[(df['a'] != a) & (df['b'] == b) & (df['c'] != c)]['c'].values)

def get_new2(row):
    a, b, c = row
    return random.choice(df.loc[(df['a'] != a) & (df['b'] == b) & (df['c'] != c), 'c'].values)


%timeit x.apply(lambda row: get_new(row), axis=1)   # 159ms
%timeit x.apply(lambda row: get_new2(row), axis=1)  # 119ms

- jpp

---

尝试了这个方法，但不幸的是在我的情况下没有改善。它仍然需要相同的时间。 - swathis

@swathis。很有趣。您能复制粘贴我的代码并确认这是否与您的数据不同有关吗？ - jpp

你的代码块有效。使用.loc 也让我加速了约50毫秒。我的问题可能是因为我在函数中基于其他条件多次进行索引吗？此外，df.sample 对我来说非常慢。 - swathis

@swathis，是的，也许这甚至不是瓶颈。我的建议是将您的代码拆分成许多小函数，并通过性能分析器运行代码以查找性能下降的原因。您可以使用 cProfile.run 进行分析，并且它会按函数输出结果。 - jpp

使用cProfile非常有用。甚至有一个PyCharm的插件可以使它更容易。谢谢 :-) - swathis