在聚合二进制列的大型数据集上优化 pandas groupby()答案

【问题标题】：Optimising pandas groupby() on a large dataset which aggregates binary columns在聚合二进制列的大型数据集上优化 pandas groupby()
【发布时间】：2020-06-12 17:58:14
【问题描述】：

我有一个格式如下的数据框：

   template          is_a is_b is_c is_d is_e
0  cv_template        0    1     0    0    0
1  topic_template     1    0     0    0    0
2  model_template     1    0     0    0    0
3  model_template     0    1     0    0    0

我想按template 分组并聚合is_ 列，这些列是每个template 的二进制值。

即在上面的示例中，输出将是：

   template          is_a is_b is_c is_d is_e
0  cv_template        0    1     0    0    0
1  topic_template     1    0     0    0    0
2  model_template     1    1     0    0    0

我目前的解决方案是这样做：

df.groupby('template', as_index=False)['is_a', 'is_b', 'is_c', 'is_d'].max()

但是，在处理大型数据集时，分组依据很慢。我想知道是否有更好的方法可以加快速度。

【问题讨论】：

你的解决方案是合理的。你的数据集有多大？它运行多长时间？
@QuangHoang，它很大。 ~50m 行。任何一点点优化都会很方便。

标签： python pandas numpy group-by pandas-groupby

【解决方案1】：

我不能确定这会快得多。但是，我把它和 Numba 放在一起

import pandas as pd
import numpy as np
from numba import njit

@njit
def max_at(i, a, shape):
    out = np.zeros(shape, np.int64)
    for j in range(len(a)):
        row = a[j]
        pos = i[j]
        cur = out[pos]
        out[pos, :] = np.maximum(cur, row)
    return out

i, t = df['template'].factorize()
cols = ['is_a', 'is_b', 'is_c', 'is_d', 'is_e']
is_ = np.column_stack([df[c].to_numpy() for c in cols])

result = max_at(i, is_, (len(t), len(cols)))

pd.DataFrame(result, t, cols).reset_index()

            index  is_a  is_b  is_c  is_d  is_e
0     cv_template     0     1     0     0     0
1  topic_template     1     0     0     0     0
2  model_template     1     1     0     0     0

【讨论】：