使用 MultiIndex 缓慢搜索大型 DataFrame

Question

我有一个大的 Pandas DataFrame（~800M 行），我在一个MultiIndex上建立了索引，它有两个索引，一个 int 和一个日期。 我想根据我拥有的整数列表（大约 10k）检索 DataFrame 行的子集。 整数与多索引的第一个索引匹配。 多索引是唯一的。

我尝试的第一件事是对索引进行排序，然后使用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. 为什么排序的多索引仍然这么慢？

Answer 1

检索包含 800M 行的 DataFrame 的子集可能很困难。 这里有一些想法可以帮助您更快地搜索 go：

1. Use.loc() 与 boolean 索引而不是 pd.IndexSlice：

使用 boolean 索引 with.loc() 而不是 pd.IndexSlice 来切片你的多索引。 这可以帮助 Pandas 避免在处理巨大的 DataFrame 时为每个切片建立新索引 object 的昂贵做法。

例如：

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

2. 避免多次设置索引：

多次设置索引和排序数据的成本可能很高。 如果可以，尽量只设置一次索引，但尽量避免对其进行排序。

例如：

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

3. 使用分块或并行处理：

您可能需要考虑将 DataFrame 分成更小的部分，分别处理每个部分，如果结果太大而无法存储在 memory 中，则将结果连接起来。要加快查询速度，您也可以使用并行处理。 这两种策略都适用于 Dask 库。

在响应您的第二个查询时，排序的多索引应该比未排序的多索引更快，因为它使 Pandas 能够利用 NumPy 内置的快速搜索方法。但是，如果一个巨大的 DataFrame 有很多列或排序顺序很复杂，对数据进行排序可能很昂贵。 一般来说，排序 DataFrame 是一个昂贵的过程，应该尽可能避免。

Answer 2

MultiIndices非常方便，但根据我的经验，速度非常慢。 这是pandas已经为单深度行和列标签添加超过numpy的巨大开销之上的。

如果您的索引/列相当稳定并且其他所有事情都可以在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 生成过程中创建它。