提高处理大熊猫数据帧的性能

Question

我有一个大熊猫数据框（100 万行），我的代码需要更好的性能来处理这些数据。

下面是我的代码，也提供了profiling分析。

数据集的标题：

key_id, date, par1, par2, par3, par4, pop, price, value

对于每个键，我们有一行包含 5000 个可能的日期中的每一个

有 200 个 key_id * 5000 个日期 = 1000000 行

使用不同的变量 var1, ..., var4，我为每一行计算一个值，我想为每个 key_id 提取具有最佳值的前 20 个日期，然后计算所用变量集的流行度。

最后，我想找到优化这种受欢迎程度的变量。

def compute_value_col(dataset, val1=0, val2=0, val3=0, val4=0):
    dataset['value'] = dataset['price'] + val1 * dataset['par1'] \
        + val2 * dataset['par2'] + val3 * dataset['par3'] \
        + val4 * dataset['par4']

    return dataset

def params_to_score(dataset, top=10, val1=0, val2=0, val3=0, val4=0):
    dataset = compute_value_col(dataset, val1, val2, val3, val4)
    dataset = dataset.sort(['key_id','value'], ascending=True)
    dataset = dataset.groupby('key_id').head(top).reset_index(drop=True)
    return dataset['pop'].sum()

def optimize(dataset, top):
    for i,j,k,l in product(xrange(10),xrange(10),xrange(10),xrange(10)):
        print i, j, k, l, params_to_score(dataset, top, 10*i, 10*j, 10*k, 10*l)

optimize(my_dataset, 20)

我需要提高性能

这是运行 49 params_to_score 后的 %prun 输出

 ncalls  tottime  percall  cumtime  percall filename:lineno(function)
       98    2.148    0.022    2.148    0.022 {pandas.algos.take_2d_axis1_object_object}
       49    1.663    0.034    9.852    0.201 <ipython-input-59-88fc8127a27f>:150(params_to_score)
       49    1.311    0.027    1.311    0.027 {method 'get_labels' of 'pandas.hashtable.Float64HashTable' objects}
       49    1.219    0.025    1.223    0.025 {pandas.algos.groupby_indices}
       49    0.875    0.018    0.875    0.018 {method 'get_labels' of 'pandas.hashtable.PyObjectHashTable' objects}
      147    0.452    0.003    0.457    0.003 index.py:581(is_unique)
      343    0.193    0.001    0.193    0.001 {method 'copy' of 'numpy.ndarray' objects}
        1    0.136    0.136   10.058   10.058 <ipython-input-59-88fc8127a27f>:159(optimize)
      147    0.122    0.001    0.122    0.001 {method 'argsort' of 'numpy.ndarray' objects}
      833    0.112    0.000    0.112    0.000 {numpy.core.multiarray.empty}
       49    0.109    0.002    0.109    0.002 {method 'get_labels_groupby' of 'pandas.hashtable.Int64HashTable' objects}
       98    0.083    0.001    0.083    0.001 {pandas.algos.take_2d_axis1_float64_float64}
       49    0.078    0.002    1.460    0.030 groupby.py:1014(_cumcount_array)

我想我可以通过 key_id 将大数据帧拆分为小数据帧，以缩短排序时间，因为我想为每个 key_id 取前 20 个具有最佳值的日期，所以按键排序只是为了分隔不同的键.

但我需要任何建议，因为我需要运行数千个 params_to_score，我该如何提高此代码的效率？

编辑：@杰夫

非常感谢您的帮助！

我尝试使用 nsmallest 代替 sort & head，但奇怪的是，当我对以下两个函数进行基准测试时，它慢了 5-6 倍：

def to_bench1(dataset):
    dataset = dataset.sort(['key_id','value'], ascending=True)
    dataset = dataset.groupby('key_id').head(50).reset_index(drop=True)
    return dataset['pop'].sum()

def to_bench2(dataset):
    dataset = dataset.set_index('pop')
    dataset = dataset.groupby(['key_id'])['value'].nsmallest(50).reset_index()
    return dataset['pop'].sum()

在大约 100000 行的样本中，to_bench2 的执行时间为 0.5 秒，而 to_bench1 平均只需 0.085 秒。

在分析 to_bench2 后，我注意到与以前相比，isinstance 调用更多，但我不知道它们来自哪里......

Answer 1

显着加快速度的方法是这样的。

创建一些示例数据

In [148]: df = DataFrame({'A' : range(5), 'B' : [1,1,1,2,2] })

像你一样定义compute_val_column

In [149]: def f(p):
    return DataFrame({ 'A' : df['A']*p, 'B' : df.B })
   .....: 

这些是案例（这你可能想要一个元组列表），例如您要输入上述函数的所有情况的笛卡尔积

In [150]: parms = [1,3]

创建一个具有完整值集的新数据框，由每个参数键控）。这基本上是一个广播操作。

In [151]: df2 = pd.concat([ f(p) for p in parms ],keys=parms,names=['parm','indexer']).reset_index()

In [155]: df2
Out[155]: 
   parm  indexer   A  B
0     1        0   0  1
1     1        1   1  1
2     1        2   2  1
3     1        3   3  2
4     1        4   4  2
5     3        0   0  1
6     3        1   3  1
7     3        2   6  1
8     3        3   9  2
9     3        4  12  2

这就是魔法。按您想要的任何列分组，包括作为第一个（或可能多个）的 parm。然后进行部分排序（这是 nlargest 所做的）；这比排序和头更有效（这取决于组密度）。最后求和（同样是我们所讨论的石斑鱼，因为您正在进行“部分”减少）

In [153]: df2.groupby(['parm','B']).A.nlargest(2).sum(level=['parm','B'])
Out[153]: 
parm  B
1     1     3
      2     7
3     1     9
      2    21
dtype: int64