Pandas.groupby.apply()中的内存泄漏？

Question

我现在在做一个项目，使用Pandas处理大约600mb的csv文件。在分析过程中，我把csv文件读入一个数据框（dataframe），然后根据某一列进行分组，并对分组后的数据框应用一个简单的函数。我发现这个过程中我的系统开始使用交换内存（Swap Memory），所以我做了一个简单的测试：

首先，我在命令行中创建了一个相当大的数据框：

import pandas as pd
import numpy as np
df = pd.DataFrame(np.random.randn(3000000, 3),index=range(3000000),columns=['a', 'b', 'c'])

我定义了一个没什么用的函数，叫做do_nothing()：

def do_nothing(group):
    return group

然后我运行了以下命令：

df = df.groupby('a').apply(do_nothing)

我的系统有16GB的内存，运行的是Debian（Mint）。在创建数据框后，我大约使用了600MB的内存。当apply方法开始执行时，这个内存使用量开始飙升。它稳步上升到大约7GB（！），然后在命令完成后又降回到5.4GB（当命令行仍然在运行时）。问题是，我的工作需要做的不止是“什么都不做”，所以在执行真正的程序时，我的16GB内存被用满了，开始使用交换内存，这让程序变得无法使用。这是正常现象吗？我不明白为什么Pandas需要7GB的内存来有效地“什么都不做”，即使它需要存储分组后的对象。

有没有人知道是什么原因导致的，或者怎么解决这个问题？

谢谢，

.P

Answer 1

我的解决方案：

result_df = None
        for i, (account_id, grp) in enumerate(grouped_df):
            grp.name = account_id
            if i % 500 == 0:
                print(f"\rStep {i}", end="", flush=True)
                gc.collect()

            series = partial_func(grp)
            if (
                series is not None
            ):  
                dataframed = series.to_frame().transpose()

                if result_df is None:
                    result_df = dataframed
                else:
                    result_df.append(dataframed)

            else:
                print("Cleaning dropped row.")

       
        grouped_df = result_df
        del result_df
        gc.collect()

Answer 2

使用0.14.1版本，我觉得没有内存泄漏（占用你框架的1/3大小）。

In [79]: df = DataFrame(np.random.randn(100000,3))

In [77]: %memit -r 3 df.groupby(df.index).apply(lambda x: x)
maximum of 3: 1365.652344 MB per loop

In [78]: %memit -r 10 df.groupby(df.index).apply(lambda x: x)
maximum of 10: 1365.683594 MB per loop

对于如何处理这样的问题，有两个一般性的建议：

1) 如果可能的话，使用cython级别的函数，这样会快得多，而且会占用更少的内存。换句话说，通常把一个groupby表达式拆开来用是值得的，尽量避免使用函数（当然，有些事情太复杂了，没办法，但你要的就是把事情拆解开）。比如：

不要这样做：

df.groupby(...).apply(lambda x: x.sum() / x.mean())

这样做要好得多：

g = df.groupby(...)
g.sum() / g.mean()

2) 你可以通过手动进行聚合来轻松“控制”groupby（这样还可以在需要时进行定期输出和垃圾回收）。

results = []
for i, (g, grp) in enumerate(df.groupby(....)):

    if i % 500 == 0:
        print "checkpoint: %s" % i
        gc.collect()


    results.append(func(g,grp))

# final result
pd.concate(results)