【问题标题】：Finding top N values for each group, 200 million rows为每组查找前 N 个值，2 亿行
【发布时间】：2019-05-17 16:41:45
【问题描述】：

我有一个大约 2 亿行的 pandas DataFrame，如下所示：

UserID  MovieID  Rating
1       455      5
2       411      4
1       288      2
2       300      3
2       137      5
1       300      3

...

我想按评分降序排列每个用户的前 N 部电影，因此对于 N=2，输出应如下所示：

UserID  MovieID  Rating
1       455      5
1       300      3
2       137      5
2       411      4

当我尝试这样做时，我得到一个由“groupby”引起的“内存错误”（我的机器上有 8gb 的 RAM）

df.sort_values(by=['rating']).groupby('userID').head(2)

有什么建议吗？

【问题讨论】：

排序会很昂贵，而且一切都在内存中，如果可能的话，使用 minHeap （显然你需要以某种方式离开 DF）？
@SMA 排序实际上在 10 分钟内完成并且不是问题（除了相对较慢），只有当我添加 groupby 部分时它才会崩溃。
2 亿远远超出了 pandas 的处理能力。尝试使用 dask 或 spark。
鉴于排序有效，您可以使用 Numpy 操作来替换 groupby （请参阅我的答案的第一部分）。如果失败了，那么您几乎会被批处理解决方案（如 Avi 所建议的那样）或重型解决方案（如 Coldspeed 所建议的那样）所困扰。

标签： python python-3.x pandas numpy pandas-groupby

【解决方案1】：

快速而肮脏的答案

鉴于排序有效，您可能可以使用以下方法来吱吱作响，它使用基于 Numpy 的内存高效替代 Pandas groupby：

import pandas as pd

d = '''UserID  MovieID  Rating
1       455      5
2       411      4
3       207      5
1       288      2
3        69      2
2       300      3
3       410      4
3       108      3
2       137      5
3       308      3
1       300      3'''
df = pd.read_csv(pd.compat.StringIO(d), sep='\s+', index_col='UserID')

df = df.sort_values(['UserID', 'Rating'])

# carefully handle the construction of ix to ensure no copies are made
ix = np.zeros(df.shape[0], np.int8)
np.subtract(df.index.values[1:], df.index.values[:-1], out=ix[:-1])

# the above assumes that UserID is the index of df. If it's just a column, use this instead
#np.subtract(df['UserID'].values[1:], df['UserID'].values[:-1], out=ix[:-1])

ix[:-1] += ix[1:]
ix[-2:] = 1
ix = ix.view(np.bool)
print(df.iloc[ix])

输出：

        MovieID  Rating
UserID                 
1           300       3
1           455       5
2           411       4
2           137       5
3           410       4
3           207       5

更节省内存的答案

而不是 Pandas 数据框，对于这么大的东西，您应该只使用 Numpy 数组（Pandas 用于在后台存储数据）。如果您使用适当的structured array，您应该能够将所有数据放入大致大小的单个数组中：

2 * 10**8 * (4 + 2 + 1)
1,400,000,000 bytes
or ~1.304 GB

这意味着它（和几个临时计算）应该很容易适应您的 8 GB 系统内存。

这里有一些细节：

最棘手的部分是初始化结构化数组。您可能能够摆脱手动初始化数组，然后复制数据：

dfdtype = np.dtype([('UserID', np.uint32), ('MovieID', np.uint16), ('Rating', np.uint8)])
arr = np.empty(df.shape[0], dtype=dfdtype)
arr['UserID'] = df.index.values
for n in dfdtype.names[1:]:
    arr[n] = df[n].values

如果上述情况导致内存不足错误，则从程序一开始，您就必须构建和填充结构化数组而不是数据帧：

arr = np.empty(rowcount, dtype=dfdtype)
...
adapt the code you use to populate the df and put it here
...

拥有arr 后，您可以按照以下方式进行分组：

arr.sort(order=['UserID', 'Rating'])

ix = np.zeros(arr.shape[0], np.int8)
np.subtract(arr['UserID'][1:], arr['UserID'][:-1], out=ix[:-1])
ix[:-1] += ix[1:]
ix[-2:] = 1
ix = ix.view(np.bool)
print(arr[ix])

上述大小计算和dtype假设没有UserID大于4,294,967,295，没有MovieID大于65535，也没有一个等级大于255。这意味着您的数据框的列可以是 (np.uint32, np.uint16, np.uint8) 而不会丢失任何数据。

【讨论】：

【解决方案2】：

如果您想继续使用 pandas，您可以将数据分成多个批次 - 例如，一次 10K 行。您可以在将源数据加载到 DF 之后拆分数据，或者更好的是，分批加载数据。
您可以将每次迭代（批次）的结果保存到字典中，只保留您感兴趣的电影数量：

{userID: {MovieID_1: score1, MovieID_2: s2, ... MovieID_N: sN}, ...}

并在每次迭代时更新嵌套字典，每个用户只保留最好的 N 部电影。

这样您就可以分析比计算机内存大得多的数据

【讨论】：