基于Dask Array实现Rasterio线程安全写入GeoTIFF的方案问询

嘿，我刚好之前踩过同样的坑，给你几个经过实践验证的方案，既能保证文件不损坏，又能最大化利用Dask的并行性提升速度：

方案一：用Dask Delayed + Rasterio窗口写入（最灵活可靠） #

这个思路是让每个Dask块独立打开、写入、关闭文件，完全避免共享文件句柄的线程安全问题，并行效率也很高。

步骤很清晰：

1. 先创建一个空的GeoTIFF文件，提前写好所有元数据（CRS、分辨率、尺寸等）；
2. 给每个Dask块分配对应的Rasterio窗口，用delayed封装独立的写入任务；
3. 最后用da.compute执行所有并行任务。

代码示例：

import dask.array as da
import rasterio
from dask.delayed import delayed
from rasterio.windows import Window

# 示例Dask数组，替换成你的实际数据
arr = da.random.random((1, 10000, 10000), chunks=(1, 1000, 1000))
# 替换成你的实际GeoTIFF元数据
meta = {
    'driver': 'GTiff',
    'height': arr.shape[1],
    'width': arr.shape[2],
    'count': arr.shape[0],
    'dtype': arr.dtype,
    'crs': 'EPSG:4326',
    'transform': rasterio.transform.from_origin(-180, 90, 0.0001, 0.0001),
    'nodata': None
}

# 第一步：创建空的输出文件
with rasterio.open('output.tif', 'w', **meta) as dst:
    pass

# 第二步：生成每个块的写入任务
def write_window(window, chunk_data):
    # 每个任务独立打开文件，写入后立刻关闭
    with rasterio.open('output.tif', 'r+') as dst:
        dst.write(chunk_data, window=window)

tasks = []
# 遍历所有块的索引
for band_idx in range(arr.numblocks[0]):
    for row_idx in range(arr.numblocks[1]):
        # 计算当前块的切片范围
        band_slice = slice(band_idx * arr.chunks[0][0], (band_idx+1)*arr.chunks[0][0])
        row_slice = slice(row_idx * arr.chunks[1][0], (row_idx+1)*arr.chunks[1][0])
        chunk_slice = (band_slice, row_slice, slice(0, arr.shape[2]))
        # 转换为Rasterio的Window格式（注意行列顺序）
        window = Window.from_slices(row_slice, band_slice)
        # 获取当前块的Dask数组
        chunk = arr[chunk_slice]
        # 添加延迟任务
        tasks.append(delayed(write_window)(window, chunk))

# 执行所有并行任务
da.compute(*tasks)

这个方案的优势是完全避免了线程安全冲突，每个任务独立操作文件，并行效率拉满，而且自定义程度高，适合复杂的元数据或特殊需求场景。

方案二：用dask-rasterio库（最简洁高效） #

如果你的环境允许安装第三方库，dask-rasterio是专门为Dask和Rasterio集成设计的，内部已经处理了所有线程安全和并行写入的细节，代码超简洁，速度也经过优化。

先安装：

pip install dask-rasterio

然后直接调用：

import dask.array as da
import dask_rasterio

# 替换成你的Dask数组和元数据
arr = da.random.random((1, 10000, 10000), chunks=(1, 1000, 1000))
meta = {
    'driver': 'GTiff',
    'height': arr.shape[1],
    'width': arr.shape[2],
    'count': arr.shape[0],
    'dtype': arr.dtype,
    'crs': 'EPSG:4326',
    'transform': rasterio.transform.from_origin(-180, 90, 0.0001, 0.0001),
    'nodata': None
}

# 一行搞定
dask_rasterio.to_raster(arr, 'output.tif', **meta)

这绝对是最优解，官方维护的库，踩坑更少，性能也有保障。

方案三：用da.store + 自定义锁（妥协方案） #

如果你坚持要用da.store，那必须通过全局锁强制写入操作互斥，否则文件必然损坏（GDAL/Rasterio的文件句柄完全不支持多线程并发写入）。

代码示例：

import dask.array as da
import rasterio
from threading import Lock

# 创建全局锁，确保同一时间只有一个线程写入
write_lock = Lock()

def safe_rasterio_writer(chunk, path, **kwargs):
    with write_lock:
        with rasterio.open(path, 'r+') as dst:
            dst.write(chunk)

# 先创建空文件
with rasterio.open('output.tif', 'w', **meta) as dst:
    pass

# 使用自定义writer，关闭da.store自带的锁（我们自己处理了）
da.store(arr, 'output.tif', writer=safe_rasterio_writer, lock=False)

这个方案的缺点是锁会限制并行性，速度不如前两个方案，适合块数量较少或者块尺寸很大的场景。

总结 #

• 优先选dask-rasterio.to_raster，简洁高效，无额外代码；
• 不能用第三方库的话，选Delayed+窗口写入，灵活可靠；
• 绝对不要用da.store(lock=False)，GDAL/Rasterio的文件句柄天生不是线程安全的，多线程并发写入一定会破坏文件结构。

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