索引非常大的 Pandas 数据帧的最快方法

Question

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

这个数据框 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 和 obj 的行。

a) 我尝试通过使用isin() 函数生成布尔系列来索引KG：

KG[KG['subj'].isin(['BART', 'NEWYORK']) & KG['obj'].isin(['USA', 'HOMMER'])]

b) 我还尝试使用query() 函数对KG 进行索引：

KG = KG.set_index(['subj','obj'], drop=True)
KG = KG.sort_index()
subj_substitution = ['BART', 'NEWYORK']
obj_substitution= ['USA', 'HOMMER']    
KG.query(f"subj in {subj_substitution} & obj in {obj_substitution}

c) 我还尝试使用merge() 加入两个 DataFrame，如下所示。

subj_df

      subj
0     BART
1  NEWYORK


obj_df

      obj
0     USA
1  HOMMER

merge_result = pd.merge(KG, subj_df, on = ['subj']).drop_duplicates()
merge_result = pd.merge(merge_result, obj_df, on = ['obj']).drop_duplicates()

这些方法的结果如下：

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

我使用timeit 函数来检查每个时间，如下所示。

timeit.timeit(lambda: KG[(KG['subj'].isin(['BART', 'NEWYORK']) & (KG['obj'].isin(['USA', 'HOMMER'])))] , number=10)

运行时是：

function	runtime
isin()	35.6s
query()	155.2s
merge()	288.9s

我认为isin() 是索引一个非常大的数据框的最快方法。 如果您能告诉我比这更快的方法，我将不胜感激。

Answer 1

我个人会选择isin 或query with in。

熊猫文档说：

query() 的性能

对于大帧，使用 numexpr 的 DataFrame.query() 比 Python 稍快。 注意：只有当您的框架有超过大约 200,000 行时，您才会看到将 numexpr 引擎与 DataFrame.query() 结合使用的性能优势。

关于query的详细信息可以找到here

在您的示例中，当我使用query 和isin 测试shape (50331648, 3) - 50M+ rows and 3 column 和shape (50331648, 3) - 50M+ rows and 3 column 时，性能结果几乎相同。

不正常

%timeit KG[KG['subj'].isin(['BART', 'NEWYORK']) & KG['obj'].isin(['USA', 'HOMMER'])]
4.14 s ± 83.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

使用 in 运算符查询

%timeit KG.query("(subj in ['BART', 'NEWYORK']) and (obj in ['USA', 'HOMMER'])")
4.08 s ± 82.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

用 isin 查询

%timeit KG.query("(subj.isin(['BART', 'NEWYORK']))& (obj.isin(['USA', 'HOMMER']))")
4.99 s ± 210 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

测试数据

d="""pred,subj,obj
nationality,BART,USA
placeOfBirth,BART,NEWYORK
locatedIn,NEWYORK,USA
hasFather,BART,HOMMER
nationality,HOMMER,USA
placeOfBirth,HOMMER,NEWYORK"""
KG = pd.read_csv(StringIO(d))
for i in range(23):
    KG = pd.concat([KG,KG])
KG.shape # (50331648, 3)

如果关注performance + code readability(maintenance)，那么至少对于复杂的查询我会选择query function。