使用Dask并行计算栅格分段统计时client.scatter调用崩溃问题

我有一幅4波段NAIP栅格图像，以及一幅通过scipy SLIC算法得到的、包含表示图像分段整数的栅格。工作流的下一步是计算每个分段内所有像素的统计信息，由于分段数量超过30万个，尝试用Dask实现并行计算。

最初编写的每个分段ID计算函数：

def get_features(id):
    segment_pixels = img[segments == id]
    return segment_features(segment_pixels)

用Dask bag处理少量分段时：

import dask.bag as db
b = db.from_sequence(segment_ids[:80], npartitions=4)
b1 = b.map(get_features).compute()

发现内存占用快速增长，原因是每个任务都在传递img和segments这两个栅格数据，属于不合理模式。

尝试用client.scatter向工作节点传递对象，代码如下：

scattered_img = client.scatter(img, broadcast=True)
scattered_segments = client.scatter(segments, broadcast=True)

def get_features(id, img=scattered_img, segments=scattered_segments):
    segment_pixels = img[segments == id]
    return segment_features(segment_pixels)
    
b1 = b.map(get_features).compute()

但导致会话崩溃，请问操作哪里有误？

核心问题 #

1. client.scatter返回的是Future对象而非实际数据：直接把Future对象作为函数默认参数传入，任务执行时无法正确解析为栅格数据。默认参数在函数定义时就完成绑定，而Future需要在工作节点通过.result()获取实际数据，绑定时机错误导致会话崩溃。
2. 逐ID遍历模式效率极低：哪怕解决了数据传递问题，30万个分段逐ID执行segments == id的布尔索引，本质是对整个栅格做30万次全量扫描，计算和内存开销都会爆炸，这才是性能问题的根源。

优化方案 #

方案1：修正scatter使用方式（仅解决数据传递，不解决效率问题）

不要将Future作为默认参数，在函数内部显式获取实际数据：

scattered_img = client.scatter(img, broadcast=True)
scattered_segments = client.scatter(segments, broadcast=True)

def get_features(id):
    # 在工作节点上获取实际栅格数据
    local_img = scattered_img.result()
    local_segments = scattered_segments.result()
    segment_pixels = local_img[local_segments == id]
    return segment_features(segment_pixels)
    
b = db.from_sequence(segment_ids[:80], npartitions=4)
b1 = b.map(get_features).compute()

但这种方式依然存在逐ID遍历的低效问题，不适合30万分段的大规模场景。

方案2：改用向量化分组计算（推荐）

利用numpy/xarray的分组统计能力，结合Dask并行特性，彻底避免逐ID遍历的开销：

1. 基于Numpy+Dask实现

import numpy as np
from dask import array as da, delayed

# 将栅格转为Dask数组（按需设置分块大小）
dask_img = da.from_array(img, chunks=(1000, 1000, 4))
dask_segments = da.from_array(segments, chunks=(1000, 1000))

# 获取唯一分段ID
unique_ids = da.unique(dask_segments).compute()

def compute_segment_stats():
    # 用bincount实现高效分组统计，避免逐ID索引
    results = {}
    segments_flat = dask_segments.flatten().compute()
    for band_idx in range(4):
        band_data = dask_img[:, :, band_idx].flatten().compute()
        # 计算每个分段的像素数和像素值总和
        counts = np.bincount(segments_flat)
        sums = np.bincount(segments_flat, weights=band_data)
        # 计算均值
        means = sums / counts
        # 映射到对应ID
        for idx, seg_id in enumerate(unique_ids):
            if seg_id not in results:
                results[seg_id] = {}
            results[seg_id][f"band_{band_idx+1}_mean"] = means[idx]
    return results

# 用Dask延迟执行计算
delayed_task = delayed(compute_segment_stats)()
final_results = delayed_task.compute()

2. 基于Xarray+Dask实现（更简洁）

import xarray as xr
import dask.array as da

# 构建Xarray Dataset，自动适配Dask分块
ds = xr.Dataset(
    {
        "band1": (("y", "x"), img[:, :, 0]),
        "band2": (("y", "x"), img[:, :, 1]),
        "band3": (("y", "x"), img[:, :, 2]),
        "band4": (("y", "x"), img[:, :, 3]),
        "segments": (("y", "x"), segments)
    }
).chunk(chunks={"y": 1000, "x": 1000})

# 按分段ID分组计算均值（可扩展为其他统计量，如std、min、max）
segment_stats = ds.groupby("segments").mean().compute()

这种向量化分组方式，将全量扫描次数从30万次降到1次，内存和计算效率会提升几个数量级。

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