使用Datashader结合Dask Xarray时内存占用过高问题咨询

你的内存暴涨问题根源有两点：

1. Dask任务调度并行度过高：
   16核虚拟机下，Dask默认启用与CPU核心数一致的线程数（16线程）。你的数据块为8192×8192（单块float64数据约512MB），16线程同时运行时会一次性加载大量数据块，总容量快速逼近100GB内存上限。

2. raster聚合的计算特性：
   ds.Canvas.raster()是将大尺寸输入重采样到300×300画布的过程，需要遍历所有输入数据块（每个输出像素可能覆盖多个输入块）。若并行度无限制，Dask会短时间内加载绝大多数数据块，直接占满内存。

注意：错误栈中shade调用的compute()是针对聚合后的300×300小结果，这一步本身不会占用大量内存，真正的内存消耗发生在计算聚合结果的过程中。

1. 限制Dask并行线程数 #

通过减少线程数，控制同时加载的数据块数量，避免内存过载：

import dask.config

# 根据内存情况调整线程数（比如4-8）
dask.config.set(scheduler='threads', num_workers=4)

# 执行绘图代码
tf.shade(ds.Canvas(plot_height=300, plot_width=300).raster(dask_xarray))

2. 使用Dask分布式本地集群（更推荐） #

通过本地集群可精细控制内存使用，Dask会自动调度任务，避免内存溢出：

from dask.distributed import Client, LocalCluster

# 创建本地集群，限制每个工作进程内存（示例：8个进程，每个10GB）
cluster = LocalCluster(memory_limit='10GB', n_workers=8)
client = Client(cluster)

# 执行绘图代码
result = tf.shade(ds.Canvas(plot_height=300, plot_width=300).raster(dask_xarray))

# 关闭集群
client.close()
cluster.close()

3. 优化数据块大小（可选） #

若当前块大小导致单块内存占用过高，可调整为更小尺寸（比如4096×4096），降低单块内存压力：

# 重新分块Dask数组
dask_array = dask_array.rechunk(chunks=(4096, 4096))
dask_xarray = dask_xarray.chunk(chunks={"x":4096, "y":4096})

import numpy as np
import dask.array as da
import datashader as ds
from datashader import transfer_functions as tf
import xarray as xr
import dask.config

# 设置Dask线程数
dask.config.set(scheduler='threads', num_workers=4)

# 创建大型Dask数组
dask_array = da.random.random((100000, 100000), chunks=(1000, 1000))
# 转换为Dask Xarray
dask_xarray = xr.DataArray(
    dask_array,
    dims=["x", "y"], 
    coords={"x": np.arange(100000), "y": np.arange(100000)},
    name="example_data"
)
# 绘制图像
tf.shade(ds.Canvas(plot_height=300, plot_width=300).raster(dask_xarray))

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