Повысить скорость фильтра панд, сохранив индексы?

Question

У меня есть следующий df:

df = pd.DataFrame({'ID1':[1,2,3,4,5,6],'ID2':[2,6,6,2,1,2],'AREA':[1,1,1,1,1,1]})
...

    ID1 ID2 AREA
0   1   2   1
1   2   6   1
2   3   6   1
3   4   2   1
4   5   1   1
5   6   2   1

Я накапливаю столбец AREA следующим образом:

for id_ in df.ID1:   
    id1_filter = df.ID1 == id_
    id2_filter = (df.ID1 == id_) | (df.ID2 == id_)
    df.loc[id1_filter, 'AREA'] = df.loc[id2_filter].AREA.sum()

print(df)
...
ID1 ID2 AREA
0   1   2   2
1   2   6   5
2   3   6   1
3   4   2   1
4   5   1   1
5   6   2   7

Для каждого id_ в ID1, AREA суммируется, где ID1 == id_ или ID2 == id_, и он всегда запускается, когда df отсортирован по ID1.

Реальный фрейм данных, над которым я работаю, - это 150 000 записей, каждая строка принадлежит уникальному ID1. Выполнение вышеуказанного на этом информационном кадре занимает 2,5 часа. Так как эта операция будет проходить повторно в обозримом будущем я решил сохранить индексы истинных значений в id1_filter и id2_filter в БД со следующей схемой.

Таблица ID1:

ID_,INDEX_
1  ,   0
2  ,   1
etc, ect

Таблица ID2:

ID_,INDEX_
1  ,   0
1  ,   4
2  ,   0
2  ,   1
2  ,   3
2  ,   5
etc, etc

При следующем запуске накопления в столбце AREA (теперь заполнены различными значениями AREA) Я читаю в таблицах sql и конвертирую их в dicts. Я тогда использую эти диктанты чтобы получить нужные записи во время цикла суммирования.

id1_dict = pd.read_sql('select * from ID1',db_engine).groupby('ID_').INDEX_.unique().to_dict()
id2_dict = pd.read_sql('select * from ID2',db_engine).groupby('ID_').INDEX_.unique().to_dict()

# print indices for id1_filter and id2_fillter for id 1
print(id1_dict[1])
print(id2_dict[1])
...
[0]
[0, 4]

 for id_ in df.ID1:
        df.loc[id1_dict[id_], 'AREA'] = df.loc[id2_dict[id_]].AREA.sum()

При таком запуске это займет всего 6 минут!

Мой вопрос: есть ли лучший / стандартный способ справиться с этим сценарием, т. Е. Сохранить выборки данных для позже использовать? Примечание: я установил индекс для столбцов идентификаторов таблицы SQL и попытался получить индексирует, запрашивая таблицу для каждого идентификатора, и это работает хорошо, но все же занимает немного больше времени, чем указано выше (9 минут).

Answer 1

Один из способов сделать это так:

df = df.set_index('ID1') 
for row in df.join(df.groupby('ID2')['AREA'].apply(lambda x: x.index.tolist()),rsuffix='_').dropna().itertuples():
    df.loc[row[0],'AREA'] += df.loc[row[3],'AREA'].sum()
df = df.reset_index()

и вы получите ожидаемый результат

   ID1  ID2  AREA
0    1    2     2
1    2    6     5
2    3    6     1
3    4    2     1
4    5    1     1
5    6    2     7

Теперь больше df вроде:

df = pd.DataFrame( {'ID1':range(1,1501),'ID2': np.random.randint(1,1501,(1500,)),'AREA':[1]*1500}, 
                   columns = ['ID1','ID2','AREA'])

Метод, представленный здесь, оборачивается на моем компьютере примерно за 0,76 с, а ваш первый работает за 6,5 с.

В конечном итоге вы можете создать df_list, например:

df_list = (df.set_index('ID1')
             .join(df.set_index('ID1').groupby('ID2')['AREA']
                     .apply(lambda x: x.index.tolist()),rsuffix='_ID2')
             .dropna().drop(['AREA','ID2'],1))

чтобы где-то хранить информацию, связывающую ID1 и ID2: здесь вы можете увидеть, что идентификатор равен 2 в столбце ID2, где значение ID1 = 1, 4 и 6

      AREA_ID2
ID1           
1          [5]
2    [1, 4, 6]
6       [2, 3]

и затем вы можете запустить, чтобы не создавать заново df_list, с небольшой разницей в коде:

df = df.set_index('ID1') 
for row in df_list.itertuples():
    df.loc[row[0],'AREA'] += df.loc[row[1],'AREA'].sum()
df = df.reset_index()

Надеюсь, это быстрее