Python 带有自定义agg函数的pandas groupby速度太慢或占用太多内存

我在15M行数据帧上运行groupby，按2个键（每个键最多30个字符）进行分组，并应用一个返回多个值的自定义聚合函数，然后写入CSV。我有两种方法，一种耗尽内存并失败，另一种太慢（运行DoFar需要24小时以上…）。简化代码如下

是否有： -有没有办法减少方法1的内存使用 或 -在方法2中有没有加速迭代器的方法

import pandas as pd
import numpy as np

def myfunct(x):
    # test function
    return 1,2,3

# fake the dataset
df = pd.DataFrame(np.random.randn(1500,3), columns=['a', 'b', 'c'])

df['key1'] = np.random.choice(['A','B','C','D','E'], df.shape[0])
df['key2'] = np.random.choice(['A','B','C','D','E'], df.shape[0])

# group and aggregate
grouped = df.groupby(['key1', 'key2'], sort=False)


#
#  APPROACH 1.
#
# Works but runs out of memory on 15M row datafram with 30 char key1/2
#
f = {'a':['mean', 'count'], 'b':['mean'], 'c':myfunct}
grouped = grouped.aggregate(f)

grouped.columns = ['_'.join(col).strip() for col in grouped.columns.values]

#print grouped.head(5)

# split out columns - expensive?
grouped[['c_myfunct_1', 'c_myfunct_2', 'c_myfunct_3']] = grouped['c_myfunct'].apply(pd.Series)
grouped.drop('c_myfunct', axis=1, inplace=True)

#print grouped.head(5)

# write to file
grouped.to_csv('test1.csv')


#
#    APPROACH 2.
#
# Alternate approach works fine but extremely slow!
#
grouped_again = df.groupby(['key1', 'key2'], sort=False)

header_ = True

with open('test2.csv', 'a') as f:

    for keys, group in grouped_again:

        (k1, k2) = keys
        temp_df = df.loc[(df['key1'] == k1) & (df['key2'] == k2)].groupby(['key1', 'key2'], sort=False)
        temp_df = temp_df.aggregate({'a':['mean', 'count'], 'b':['mean'], 'c':myfunct})
        temp_df.columns = ['_'.join(col).strip() for col in temp_df.columns.values]

        if header_ == True:
            temp_df.to_csv(f, float_format='%.3f', header = header_)
            header_ = False

        else:
            temp_df.to_csv(f, float_format='%.3f', header = header_)

---

刚刚用您的代码和15行测试了方法1，效果很好。我认为这样做是正确的。我怀疑问题来自自定义函数。感谢测试。我用上面的代码切换了我的自定义函数，它仍然非常快。我认为自定义函数不是问题。在对它进行了更多的测试之后，问题似乎出在groupby上，因为它使用了大量不同的键。我会看看是否可以得到一些不同的测试代码来解决这个问题。是的，当groupby速度较慢时，它通常与大量不同的组相关，而这不是您设置示例数据的方式。您可能希望将range与np.tile或np.repeat组合使用，以获得大量组。另外，您可能希望使myfunct更现实一些，因为它不会像编写的那样进行任何实际聚合。可能没关系，但正如所写的那样，您甚至不需要通过groupby运行它，因此，请务必测试它正在做什么。