使用Xarray和Dask加速重塑MFDataset多文件数据集的性能优化问询

你的耗时完全不属于预期范围，核心问题出在分块策略错误、任务图构建逻辑冗余两个方面，优化方案如下：

1. 修复数据分块问题 #

你之前传入open_mfdataset的chunks参数未生效是正常现象：open_mfdataset默认优先按输入文件边界分块，指定的分块参数会被覆盖。加载完成后主动调用rechunk即可：

import xarray as xr
from dask.distributed import Client

c = Client("scheduler:8786")

ds = xr.open_mfdataset(
    "/resource_data/ecmwf/era20c/grid-av/hour-step-3/*.nc",
    parallel=True
)
# 主动重分块，每个块对应单个经纬度的全时间序列，完美匹配输出需求
ds = ds.chunk({"time": -1, "latitude": 1, "longitude": 1})
# 集群内存充足的话可以先持久化重分块后的数据集，避免重复计算
ds = c.persist(ds)

2. 替换手动循环逻辑，避免冗余任务图构建 #

你手动遍历240*121=29040个经纬度点，每次调用sel生成子数据集的操作，会反复修改Dask任务图，光是图构建的开销就会达到几十秒，save_mfdataset处理上万个数据集时图构建耗时更是会指数级上升。
用Xarray的堆叠+分组逻辑可以完全规避这个问题：

import os
import dask
# 把经纬度二维坐标堆叠成一维的latlon坐标
ds_stacked = ds.stack(latlon=["latitude", "longitude"])

# 定义延迟保存函数，单个分组对应单个经纬度点
@dask.delayed
def save_latlon_ds(group):
    lat, lng = group.latlon.item()
    save_path = os.path.join(out_dir, f"lat_{lat:.1f}_lng_{lng:.1f}.nc")
    # 还原维度结构后保存
    group = group.unstack("latlon")
    group.to_netcdf(save_path, engine="netcdf4")
    return save_path

# 生成所有分组的延迟任务
tasks = [save_latlon_ds(group) for _, group in ds_stacked.groupby("latlon")]

os.makedirs(out_dir, exist_ok=True)
# 提交任务执行，资源有限的话可以分批提交控制并发度，比如每次跑2000个
results = dask.compute(tasks, optimize_graph=True)[0]

3. 额外优化建议 #

• 不需要用h5netcdf引擎，默认的netcdf4引擎处理单个小文件的性能更好
• 重分块后可以先检查变量的chunk结构是否符合预期，确认每个变量的块大小为(321664, 1, 1)再执行后续操作
• 如果集群IO性能一般，可以适当调大每次提交的任务批量大小，避免大量小文件写入触发存储端的性能瓶颈

按以上方案调整后，整个流程的耗时应该可以控制在1小时以内，具体取决于集群的计算和IO性能。

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