Dask如何基于现有索引列[A,B]高效分组（无需重置索引）

一、利用Parquet预定义的复合索引直接分组 #

既然你的Parquet文件已经将A和B标记为索引，读取时直接指定复合索引即可避免reset_index的开销：

import dask.dataframe as dd
# 读取时把A和B设为复合索引
df = dd.read_parquet("path/to/your/file.parquet", index=["A", "B"])
# 直接按复合索引分组，无需重置索引
result = df.groupby(level=[0, 1]).apply(f)
# 或者更简洁的写法
result = df.groupby(df.index).apply(f)

这种方式的优势在于：Parquet的复合索引通常意味着数据在存储层面已按[A,B]排序或分区，Dask可以直接利用这种结构，减少跨分区的shuffle操作，大幅提升分组效率，同时完整保留原有索引。

二、已读入单索引DataFrame的转换方案 #

如果你的DataFrame已经以A为索引读入，无需重新读取文件，可通过追加索引的方式构建复合索引：

# 将B列追加到现有索引，形成[A,B]复合索引
df = df.set_index("B", append=True)
# 按复合索引分组
result = df.groupby(level=[0, 1]).apply(f)

这里append=True会保留原有的A索引，同时将B加入索引，全程不需要重置索引，避免数据重排的开销。

三、解决CPU使用率低的瓶颈问题 #

CPU使用率低通常是IO等待或shuffle开销过大导致的，可从以下方向优化：

1. 优化IO性能 #

• 合并小文件：如果Parquet文件数量过多（每个文件<100MB），会导致频繁的小文件读写，CPU处于等待状态。可通过重新分区合并文件：

  # 重新分区为合适的数量（比如按每分区100-500MB调整）
  df = df.repartition(npartitions=20)
  # 重新写入Parquet，后续读取即可用优化后的文件
  df.to_parquet("path/to/optimized_file.parquet", index=["A", "B"])
  

• 列裁剪与谓词下推：读取时只加载需要的列，减少IO数据量：

  df = dd.read_parquet("path/to/file.parquet", index=["A", "B"], columns=["C", "D"])
  

• 存储介质升级：如果使用机械硬盘，换成SSD可大幅提升随机读写速度，降低IO等待。

2. 减少Shuffle开销 #

• 利用预分区：如果Parquet是按[A,B]分区存储的，读取时指定partition_on参数，Dask会直接按分区分组，无需shuffle：

  df = dd.read_parquet("path/to/file.parquet", partition_on=["A", "B"], index=["A", "B"])
  

• 调整Shuffle配置：对于必须shuffle的场景，可尝试调整Dask的shuffle方法，默认disk适合大数据，tasks适合内存足够的场景：

  import dask.config
  # 设置为基于磁盘的shuffle（默认）
  dask.config.set({"dataframe.shuffle.method": "disk"})
  

• 控制分组基数：如果分组键[A,B]的基数极高（比如数百万个不同分组），会导致大量小分组，shuffle开销剧增。此时可考虑调整分组策略，比如先按A分组处理，再在每个A分组内按B处理。

内容的提问来源于stack exchange，提问作者Hanan Shteingart