Pandas .loc занимает очень много времени

Question

У меня есть 10 ГБ CSV-файл со 170,000,000 строками и 23 столбцами, которые я считываю в фрейм данных следующим образом:

import pandas as pd 

d = pd.read_csv(f, dtype = {'tax_id': str})

У меня также есть список строк с почти 20 000 уникальных элементов:

h = ['1123787', '3345634442', '2342345234', .... ]

Я хочу создать новый столбец с именем class в кадре данных d. Я хочу присвоить d['class'] = 'A' всякий раз, когда d['tax_id'] имеет значение, которое находится в списке строк h. В противном случае я хочу d['class'] = 'B'.

Следующий код работает очень быстро на 1% выборке моего фрейма данных d:

d['class'] = 'B'
d.loc[d['tax_num'].isin(h), 'class'] = 'A' 

Однако на полном кадре данных d этот код занимает более 48 часов (и считается) для запуска на 32-ядерном сервере в пакетном режиме. Я подозреваю, что индексирование с помощью loc замедляет код, но я не уверен, что это может быть на самом деле.

В сумме: Есть ли более эффективный способ создания столбца class?

Answer 1

«У меня также есть список строк с почти 20 000 уникальных элементов»

Ну, для начала, вам следует сделать этот список set, если вы собираетесь использовать его для тестирования членства. list объекты имеют линейное тестирование на членство по времени, set объекты имеют очень оптимизированную производительность при постоянном времени для тестирования на членство. Это самый низкий висящий фрукт здесь. Так что используйте

h = set(h) # convert list to set
d['class'] = 'B'
d.loc[d['tax_num'].isin(h), 'class'] = 'A' 

Answer 2

Если ваши налоговые номера уникальны, я бы порекомендовал установить индекс tax_num, а затем индексировать его. В нынешнем виде вы называете isin, что является линейной операцией. Какой бы быстрой ни была ваша машина, она не может выполнить линейный поиск по 170 миллионам записей за разумное время.

df.set_index('tax_num', inplace=True) # df = df.set_index('tax_num')
df['class'] = 'B'
df.loc[h,  'class'] = 'A'

Если у вас все еще проблемы с производительностью, я бы рекомендовал перейти на распределенную обработку с dask.