带MultiIndex的8亿行DataFrame查询优化及慢查询原因求助

1. 如何加快查询速度？ #

以下是几种针对大体积DataFrame的高效查询方案，覆盖Pandas原生及替代工具：

方案一：使用Merge半连接 #

利用Pandas的merge操作做内连接，借助其底层优化的归并排序算法，比loc切片更适合大表批量匹配：

import pandas as pd

# 将目标ID列表转换为单列DataFrame
ids_df = pd.DataFrame({"int_idx": ids})
# 执行内连接，仅保留匹配int_idx的行
res = pd.merge(df.reset_index(drop=True), ids_df, on="int_idx", how="inner")
# 若后续需按日期过滤，可重新设置MultiIndex
res.set_index(["int_idx", "date_idx"], inplace=True)

方案二：基于排序索引的批量位置切片 #

利用排序后MultiIndex的有序性，通过get_indexer和get_loc批量定位每个ID的行范围，避免逐个ID的串行查找：

# 获取第一层级索引的所有唯一值
unique_int_ids = df.index.get_level_values(0).unique()
# 批量获取目标ID在唯一值列表中的位置
id_positions = unique_int_ids.get_indexer(ids)

res_list = []
for pos in id_positions:
    if pos != -1:
        current_id = unique_int_ids[pos]
        # 获取当前ID对应的首末行位置
        start_loc = df.index.get_loc((current_id, df.index.get_level_values(1).min()))
        end_loc = df.index.get_loc((current_id, df.index.get_level_values(1).max()))
        # 切片并收集结果
        res_list.append(df.iloc[start_loc:end_loc + 1])

# 合并所有结果
res = pd.concat(res_list)

方案三：切换到列式数据库工具（Polars/Dask） #

对于8亿行的超大规模数据，Pandas的内存模型和单线程处理会遇到瓶颈，Polars（列式内存型）或Dask（分布式）能提供数量级的速度提升：

import polars as pl

# 转换为Polars DataFrame
pl_df = pl.from_pandas(df)
# 快速过滤匹配ID的行
res_pl = pl_df.filter(pl.col("int_idx").is_in(ids))
# 如需转回Pandas格式
res = res_pl.to_pandas()

2. 为何排序后的MultiIndex查询仍如此缓慢？ #

核心原因在于Pandas对MultiIndex列表切片的内部实现限制：

• 串行处理开销：当使用loc[idx[ids, :]]时，Pandas会逐个遍历列表中的每个ID，单独查找该ID对应的所有日期索引切片，而非批量处理。对于10k个ID，这个串行过程会累积大量的索引定位开销。
• MultiIndex的结构损耗：MultiIndex的层级索引是分开存储的，查找时需要同时定位两个层级的位置，比单索引的二分查找多了一层计算，尤其是当第一层级的唯一值数量极多（8亿行场景下必然如此），每个ID的定位成本都会被放大。
• 优化缺失：Pandas对单索引的列表查询做了批量优化（如一次性定位所有匹配的索引范围），但MultiIndex的类似场景并未得到同等优化，导致查询效率远低于单索引。

内容的提问来源于stack exchange，提问作者Luigi D.