通过存储索引提高 pandas 过滤速度？

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

对于ID1 中的每个id_，将AREA 相加，其中ID1 == id_ 或ID2 == id_， 当df 在ID1 上排序时，它总是运行。

我正在处理的真实数据框是 150,000 条记录，每一行都属于一个唯一的 ID1。 在此数据帧上运行上述内容需要 2.5 小时。由于此操作将重复进行 在可预见的未来，我决定将 True 值的索引存储在 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 表并将它们转换为字典。然后我使用这些字典 在求和循环中获取我需要的记录。

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 表的 ID 列上设置了一个索引，并尝试获取 通过为每个 id 查询 table 来索引，它运行良好，但仍然比上面的时间（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的信息保存在某处：在这里您可以在ID2列中看到id等于2，其中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()

希望它更快