使用MultiIndex搜索大型DataFrame速度较慢。

Question

我有一个大的Pandas DataFrame（~800M行），我在其中使用两个索引，一个整数和一个日期，对MultiIndex进行了索引。我想根据我拥有的整数列表（约10k）检索DataFrame的子集。整数匹配多重索引的第一个索引。MultiIndex是唯一的。

我尝试的第一件事是按索引对其进行排序，然后使用loc查询：

df = get_my_df()  # 800M rows
ids = [...]       # 10k ints, sorted list

df.set_index(["int_idx", "date_idx"], inplace=True, drop=False)
df.sort_index(inplace=True)

idx = pd.IndexSlice
res = df.loc[idx[ids, :]]

然而这个过程非常缓慢，大约一个小时后我停止了代码的运行。

接下来我尝试的是仅将第一个设置为索引。对我来说，这是次优的，因为索引不是唯一的，而且稍后我还需要按日期进一步筛选：

df.set_index("int_idx", inplace=True, drop=False)
df.sort_index(inplace=True)

idx = pd.IndexSlice
res = df.loc[idx[ids, :]]

令我惊讶的是这有所改善，但仍然非常缓慢。

我的两个问题：

1. 我该如何使查询更快？（无论是使用单索引还是多索引）
2. 为什么排序的多索引仍然很慢？

- Luigi D.

---

我认为只要你使用Python和[pandas]数据框架（除了dask包），就无法解决800M数据的计算速度\并行化问题，也许可以。相反，最好切换到使用像Pyspark、scala、SQL等语言和[Spark Dataframe]的worker。 - Mario

Answer 1

检索包含 800M 行的 DataFrame 的子集可能很困难。以下是一些帮助您加快搜索速度的想法：

1. 使用布尔索引的 .loc() 而不是 pd.IndexSlice：

使用布尔索引和.loc() 来切片您的多级索引。这可以帮助 Pandas 避免在处理大型 DataFrame 时为每个切片建立新的索引对象，从而节省成本。

例如：

res = df.loc[df.index.get_level_values('int_idx').isin(ids)]

避免多次设置索引：

多次设置索引和排序数据可能会很耗费时间。如果可以的话，尽量只设置一次索引，但是尽量避免对其进行排序。

例如：

df.set_index(["int_idx", "date_idx"], inplace=True, drop=False)
res = df[df.index.get_level_values('int_idx').isin(ids)]

使用分块或并行处理： 如果DataFrame太大无法存储在内存中，您可能想将其划分为较小的部分，分别进行处理，然后连接结果。为了加速查询，您还可以使用并行处理。这两种策略都适用于Dask库。 对于您的第二个查询，排序的多重索引应该比未排序的更快，因为它使Pandas能够利用内置在NumPy中的快速搜索方法。但是，如果一个巨大的DataFrame有许多列或排序顺序很复杂，则排序数据可能是昂贵的。一般来说，排序DataFrame是一种昂贵的过程，应尽可能避免。

- Drake

---

嘿@Drake，这确实帮了很多忙，谢谢！使用.isin()生成布尔列表并使用.loc显著加快了查询速度。只是一个小想法：排序的确很耗费资源，但我认为，如果只做一次，它可以帮助显著加快查询速度，特别是对于非常大的DFs。 - Luigi D.

Answer 2

MultiIndices 是一种很方便的工具，但在我的经验中非常慢。这还加上了 pandas 对于单层行和列标签的巨大开销。

如果您的索引/列相对稳定，并且其他所有操作都可以使用 numpy 完成，那么通过单独管理索引并使用 .to_numpy() 转换为 numpy，您将看到巨大的速度提升。根据代码不同，我已经看到了超过100倍的性能提升。首先将您的索引转换为一个 index:iloc 的 dict，然后使用它进行基于整数的行查找。

index_dict = {idx:i for i,idx in enumerate(df.index.tolist())}
n_df = df.to_numpy()
row_ilocs = [index_dict[x] for x in ids]  # get list of 0-based locations in n_df
res = n_df[row_ilocs, :]

如果您需要根据第一个索引级别执行行查找，则索引只是元组列表，因此很容易在pandas之外编写列表推导式。

如果您不想转到numpy，仍然可以通过使用.iloc而不是.loc获得大幅改进（在MultiIndex上可能高达10倍）。例如：

index_dict = {idx:i for i,idx in enumerate(df.index.tolist())}
row_ilocs = [index_dict[x] for x in ids]  # get list of 0-based locations in df
res = df.iloc[row_ilocs]

最好只转换一次并保留到index_dict，或者更好的是在生成初始df时就创建它。

- fantabolous