使用Dask读取大型HDF文件时sorted_index=True致分区数恒为1的问题

我碰到过类似的问题，这其实是Dask处理sorted_index=True时的一个常见坑——当你直接设置这个参数时，Dask的逻辑是先尝试读取整个索引来确认它的排序状态，然后基于索引的连续性划分分区。但如果你的HDF5文件里没有预先对索引做分块存储，Dask根本没法拆分索引，最后只能生成一个大分区，自然会爆内存。

下面是两个可行的解决方案：

方案1：先按chunksize拆分，再标记索引为已排序 #

先不指定sorted_index=True，让Dask通过chunksize把数据拆分成多个分区，之后再确认索引是排序的，最后重新设置索引并标记为已排序，这样Dask会保留现有的分区结构：

import dask.dataframe as dd

# 第一步：不设置sorted_index，用chunksize拆分数据为多个分区
df = dd.read_hdf('large.h5', key='data', chunksize=10000, mode='r')

# 可选但建议：验证索引确实是单调递增的（确保符合sorted_index的前提）
assert df.index.is_monotonic_increasing.compute()

# 第二步：重新设置索引并标记为已排序，保留现有分区
df = df.set_index(df.index, sorted=True)

# 现在计算操作会并行处理多个分区，不会耗尽内存
print(df.mean().compute())

方案2：手动指定分区边界（divisions） #

如果你已经知道索引的范围和想要的分区数量，可以直接手动传入divisions参数，让Dask按你定义的边界拆分数据，同时配合sorted_index=True：

import dask.dataframe as dd

# 假设你的索引是从0到100000，想要分成10个分区，每个分区10000条数据
divisions = list(range(0, 100001, 10000))

# 手动传入divisions，同时设置sorted_index=True
df = dd.read_hdf('large.h5', key='data', divisions=divisions, sorted_index=True)

# 此时npartitions会等于你定义的分区数，计算时会并行处理
print(df.npartitions)
print(df.mean().compute())

为什么原来的代码不行？ #

当你直接设置sorted_index=True时，Dask会尝试从HDF5文件中读取完整的索引元数据来确定分区边界，但如果HDF5文件的索引是作为连续数组存储的（没有分块），Dask无法将索引拆分成多个段，只能创建一个分区。而先通过chunksize读取再标记排序的方式，会让Dask基于已有的数据分区来处理索引，避免了一次性读取整个索引的问题。

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