Нахождение лучших N значений для каждой группы, 200 миллионов строк

Question

У меня есть пандас DataFrame, который имеет около 200 миллионов строк и выглядит так:

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» (на моей машине 8 ГБ ОЗУ)

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

Есть предложения?

Answer 1

Если вы хотите продолжать работать с пандами, вы можете разделить ваши данные на пакеты - например, по 10 тысяч строк.Вы можете разделить данные либо после загрузки исходных данных в DF, либо, что еще лучше, загрузить данные по частям.
Вы можете сохранять результаты каждой итерации (партии) в словаре, сохраняя только количество фильмов, которые вас интересуют:

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

и обновлять вложенный словарь на каждой итерации, сохраняя толькоN лучших фильмов на пользователя.

Таким образом, вы сможете анализировать данные, намного превышающие объем памяти вашего компьютера

Answer 2

Быстрый и грязный ответ

Учитывая, что сортировка работает, вы можете запищать с помощью следующего: в ней используется эффективная альтернатива памяти на основе Numpy для панд 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 использует для хранения данныхпод капотом).Если вы используете соответствующий структурированный массив , вы сможете уместить все ваши данные в один массив размером примерно:

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

, что означает, что он (и параВременные вычисления) должны легко поместиться в системную память объемом 8 ГБ.

Вот некоторые подробности:

• Самая сложная часть - инициализация структурированного массива.Возможно, вам удастся обойтись без ручной инициализации массива и последующего копирования данных:

  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) без потери каких-либо данных.