首页 > 解决方案 > 比 Pandas 的 isin 函数更快地提取条件行的方法

问题描述

我有一个非常大的熊猫数据框格式的知识图,如下所示。

这个数据框KG有超过 1 亿行。

公斤:

                   pred     subj      obj
        0   nationality     BART      USA
        1  placeOfBirth     BART  NEWYORK
        2     locatedIn  NEWYORK      USA
      ...           ...      ...      ...
116390740     hasFather     BART   HOMMER
116390741   nationality   HOMMER      USA
116390743  placeOfBirth   HOMMER  NEWYORK

我试图从这个 KG 中得到一个具有特定主题值的行。

使用 subj 列作为系列,我尝试通过使用isin()如下所示的函数生成布尔系列来索引 KG。

KG[KG['subj'].isin(['BART', 'NEWYORK'])]

我想要的输出是

                   pred     subj      obj
        0   nationality     BART      USA
        1  placeOfBirth     BART  NEWYORK
        2     locatedIn  NEWYORK      USA
116390740     hasFather     BART   HOMMER

我必须重复以上

但上述方法耗时较长。有没有比这种方法更有效地减少时间的方法?

谢谢!

标签: pythonpandas

解决方案


您可以set/sort index然后选择所需的值:根据索引值查找行比根据列值查找行更快。对索引进行排序时速度更快。

df = df.set_index('subj')
df = df.sort_index()
result = df.loc[['BART', 'NEWYORK']] 

您可以在设置多索引后尝试查询:

df = df.set_index(['subj','obj'])
df = df.sort_index()
df.query("subj in ['BART','NEWYORK'] & obj in ['USA','HOMMER']")

推荐阅读