有条件地创建熊猫列的最快方法
Posted
技术标签:
【中文标题】有条件地创建熊猫列的最快方法【英文标题】：Fastest way to create a pandas column conditionally 【发布时间】：2018-12-25 13:23:53 【问题描述】：在 Pandas DataFrame 中，我想根据另一列的值有条件地创建一个新列。在我的应用程序中，DataFrame 通常有几百万行，并且唯一条件值的数量很少，大约为单位。性能极其重要：生成新列的最快方法是什么？
我在下面创建了一个示例案例，并且已经尝试并比较了不同的方法。
在示例中，条件填充由
根据label 列的值进行字典查找（这里：1, 2, 3 之一）。
lookup_dict =
1: 100, # arbitrary
2: 200, # arbitrary
3: 300, # arbitrary
然后我希望我的 DataFrame 填充为：
label output
0 3 300
1 2 200
2 3 300
3 3 300
4 2 200
5 2 200
6 1 100
7 1 100
下面是10M行测试的6种不同方法（测试代码中参数Nlines）：
pandas.groupby().apply()
方法二：pandas.groupby().indices.items()
方法3：pandas.Series.map
方法 4：标签上的 for 循环
方法5：numpy.select
方法6：麻木
完整的代码在答案的末尾提供，所有方法的运行时。在比较性能之前，每种方法的输出都被断言为相等。
方法一：pandas.groupby().apply()
我在label 上使用pandas.groupby()，然后使用apply() 用相同的值填充每个块。
def fill_output(r):
''' called by groupby().apply(): all r.label values are the same '''
r.loc[:, 'output'] = lookup_dict[r.iloc[0]['label']]
return r
df = df.groupby('label').apply(fill_output)
我明白了
>>> method_1_groupby ran in 2.29s (average over 3 iterations)
请注意，groupby().apply() 在第一个组上运行两次以确定要使用的代码路径（请参阅Pandas #2936）。这可能会减慢少数群体的速度。我欺骗了方法 1 可以添加第一个虚拟组，但我没有得到太大的改进。
方法二：pandas.groupby().indices.items()
第二个是变体：我没有使用apply，而是使用groupby().indices.items() 直接访问索引。这最终是方法1的两倍，这是我用了很长时间的方法
dgb = df.groupby('label')
for label, idx in dgb.indices.items():
df.loc[idx, 'output'] = lookup_dict[label]
得到：
method_2_indices ran in 1.21s (average over 3 iterations)
方法3：pandas.Series.map
我使用了Pandas.Series.map。
df['output'] = df.label.map(lookup_dict.get)
在查找值的数量与行数相当的类似情况下，我得到了非常好的结果。在本例中，map 最终的速度是方法 1 的两倍。
method_3_map 运行时间为 3.07 秒（平均超过 3 次迭代）
我将其归因于少量查找值，但可能只是我实现它的方式存在问题。
方法 4：标签上的 for 循环
第 4 种方法非常简单：我只是遍历所有标签并选择 DataFrame 的匹配部分。
for label, value in lookup_dict.items():
df.loc[df.label == label, 'output'] = value
但令人惊讶的是，我最终得到的结果比之前的案例快得多。我预计基于groupby 的解决方案会比这个更快，因为Pandas 必须在这里与df.label == label 进行三个比较。结果证明我错了：
method_4_forloop ran in 0.54s (average over 3 iterations)
方法5：numpy.select
第五种方法使用numpy的select函数，基于这个*** answer。
conditions = [df.label == k for k in lookup_dict.keys()]
choices = list(lookup_dict.values())
df['output'] = np.select(conditions, choices)
这会产生最好的结果：
method_5_select ran in 0.29s (average over 3 iterations)
最终，我在方法 6 中尝试了 numba 方法。
方法 6：麻木
仅出于示例的目的，条件填充值是编译函数中的硬编码。我不知道如何给 Numba 一个列表作为运行时常量：
@jit(int64[:](int64[:]), nopython=True)
def hardcoded_conditional_filling(column):
output = np.zeros_like(column)
i = 0
for c in column:
if c == 1:
output[i] = 100
elif c == 2:
output[i] = 200
elif c == 3:
output[i] = 300
i += 1
return output
df['output'] = hardcoded_conditional_filling(df.label.values)
我最终获得了最佳时间，比方法 5 快 50%。
method_6_numba ran in 0.19s (average over 3 iterations)
由于上述原因，我还没有实现这个：我不知道如何给 Numba 一个列表作为运行时常量，而不会显着降低性能。
完整代码
import pandas as pd
import numpy as np
from timeit import timeit
from numba import jit, int64
lookup_dict =
1: 100, # arbitrary
2: 200, # arbitrary
3: 300, # arbitrary
Nlines = int(1e7)
# Generate
label = np.round(np.random.rand(Nlines)*2+1).astype(np.int64)
df0 = pd.DataFrame(label, columns=['label'])
# Now the goal is to assign the look_up_dict values to a new column 'output'
# based on the value of label
# Method 1
# using groupby().apply()
def method_1_groupby(df):
def fill_output(r):
''' called by groupby().apply(): all r.label values are the same '''
#print(r.iloc[0]['label']) # activate to reveal the #2936 issue in Pandas
r.loc[:, 'output'] = lookup_dict[r.iloc[0]['label']]
return r
df = df.groupby('label').apply(fill_output)
return df
def method_2_indices(df):
dgb = df.groupby('label')
for label, idx in dgb.indices.items():
df.loc[idx, 'output'] = lookup_dict[label]
return df
def method_3_map(df):
df['output'] = df.label.map(lookup_dict.get)
return df
def method_4_forloop(df):
''' naive '''
for label, value in lookup_dict.items():
df.loc[df.label == label, 'output'] = value
return df
def method_5_select(df):
''' Based on answer from
https://***.com/a/19913845/5622825
'''
conditions = [df.label == k for k in lookup_dict.keys()]
choices = list(lookup_dict.values())
df['output'] = np.select(conditions, choices)
return df
def method_6_numba(df):
''' This works, but it is hardcoded and i don't really know how
to make it compile with list as runtime constants'''
@jit(int64[:](int64[:]), nopython=True)
def hardcoded_conditional_filling(column):
output = np.zeros_like(column)
i = 0
for c in column:
if c == 1:
output[i] = 100
elif c == 2:
output[i] = 200
elif c == 3:
output[i] = 300
i += 1
return output
df['output'] = hardcoded_conditional_filling(df.label.values)
return df
df1 = method_1_groupby(df0)
df2 = method_2_indices(df0.copy())
df3 = method_3_map(df0.copy())
df4 = method_4_forloop(df0.copy())
df5 = method_5_select(df0.copy())
df6 = method_6_numba(df0.copy())
# make sure we havent modified the input (would bias the results)
assert 'output' not in df0.columns
# Test validity
assert (df1 == df2).all().all()
assert (df1 == df3).all().all()
assert (df1 == df4).all().all()
assert (df1 == df5).all().all()
assert (df1 == df6).all().all()
# Compare performances
Nites = 3
print('Compare performances for 0:.1g lines'.format(Nlines))
print('-'*30)
for method in [
'method_1_groupby', 'method_2_indices',
'method_3_map', 'method_4_forloop',
'method_5_select', 'method_6_numba']:
print('0 ran in 1:.2fs (average over 2 iterations)'.format(
method,
timeit("0(df)".format(method), setup="from __main__ import df0, 0; df=df0.copy()".format(method), number=Nites)/Nites,
Nites))
输出：
Compare performances for 1e+07 lines
------------------------------
method_1_groupby ran in 2.29s (average over 3 iterations)
method_2_indices ran in 1.21s (average over 3 iterations)
method_3_map ran in 3.07s (average over 3 iterations)
method_4_forloop ran in 0.54s (average over 3 iterations)
method_5_select ran in 0.29s (average over 3 iterations)
method_6_numba ran in 0.19s (average over 3 iterations)
我会对任何其他可以产生更好性能的解决方案感兴趣。 我最初是在寻找基于 Pandas 的方法，但我也接受基于 numba/cython 的解决方案。
编辑
添加Chrisb's methods进行比较：
def method_3b_mapdirect(df):
''' Suggested by https://***.com/a/51388828/5622825'''
df['output'] = df.label.map(lookup_dict)
return df
def method_7_take(df):
''' Based on answer from
https://***.com/a/19913845/5622825
Exploiting that labels are continuous integers
'''
lookup_arr = np.array(list(lookup_dict.values()))
df['output'] = lookup_arr.take(df['label'] - 1)
return df
运行时间为：
method_3_mapdirect ran in 0.23s (average over 3 iterations)
method_7_take ran in 0.11s (average over 3 iterations)
这使得#3 比任何其他方法都更快（除了#6），而且也是最优雅的。如果您的用户案例兼容，请使用 #7。
【问题讨论】：
是的，这就是方法#4，但它仍然比#5 慢约2x，比chrisb 建议的方法慢约x3 呃……这很难说，你甚至在问问题吗？坦率地说，这个“问题”的 80% 会成为一个很好的答案。我认为你有很多赞成票，因为这里有有用的信息，但我强烈建议你拿出全部的胆量，把它放在这个问题的答案中，然后把这个问题变成一个实际的问题。现在它只是一个时间统计数据的大转储。另外，我认为将某人的答案拉入问题是非常糟糕的形式，这进一步强化了我的观点。 【参考方案1】：我认为 .map (#3) 是这样做的惯用方式 - 但不要通过 .get - 单独使用字典，应该会看到相当显着的改进。
df = pd.DataFrame('label': np.random.randint(, 4, size=1000000, dtype='i8'))
%timeit df['output'] = df.label.map(lookup_dict.get)
261 ms ± 12.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit df['output'] = df.label.map(lookup_dict)
69.6 ms ± 3.08 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
如果条件的数量很少，并且比较便宜（即整数和您的查找表），直接比较值（4 尤其是 5）比 .map 快，但这并不总是正确的，例如如果你有一组字符串。
如果您的查找标签确实是连续整数，您可以利用这一点并使用take 进行查找，这应该与 numba 一样快。我认为这基本上是最快的——可以在 cython 中编写等效的代码，但不会更快。
%%timeit
lookup_arr = np.array(list(lookup_dict.values()))
df['output'] = lookup_arr.take(df['label'] - 1)
8.68 ms ± 332 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
【讨论】：
我在df['output'] = df.label.map(lookup_dict) 的时间大约是np.select() 的两倍，但在我的测试用例中，您的最后一个解决方案比np.select 快约4 倍。不错:)
很好地发现了方法#3 的get 问题。它比任何其他方法都快，而且绝对更优雅。 take 让它甚至快了一倍，在我的参考案例中是 0.1 秒。在我的应用程序中，我可以构建我的字典来利用这一点：这是值得的！以上是关于有条件地创建熊猫列的最快方法的主要内容，如果未能解决你的问题，请参考以下文章