python - 查找每个组的前 N ​​个值，2 亿行

Question

我有一个 pandas DataFrame，它有大约 2 亿行，看起来像这样:

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)

有什么建议吗？

Answer 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，你就可以按照你的目标进行 groupby:

  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) 而不会丢失任何数据。