如何按所属layer为Dask高低层级图的任务节点统一着色

该功能可以实现，核心思路是先从HighLevelGraph中提取每个低层级任务对应的layer归属，再自定义着色映射传入可视化接口即可。

步骤1：提取低层级任务与Layer的对应关系 #

import dask
from itertools import chain

# 沿用示例中的Dask对象生成逻辑
df = dask.datasets.timeseries(end='2000-04-01', partition_freq='1m')
df['x'] = df['x'] * 2

# 提取「任务key-所属layer名称」映射
hlg = df.dask
key_to_layer = {}
for layer_name, layer in hlg.layers.items():
    for key in layer.keys():
        key_to_layer[key] = layer_name

# 生成统一的layer-颜色映射，可自定义配色
unique_layers = list(hlg.layers.keys())
# 可替换为你需要的颜色列表，支持十六进制颜色值或标准颜色名
color_palette = ["#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", "#8c564b"]
layer_to_color = {layer: color_palette[i % len(color_palette)] for i, layer in enumerate(unique_layers)}

# 生成低层级任务专属的颜色映射
task_color_map = {key: layer_to_color[layer] for key, layer in key_to_layer.items()}

步骤2：调用可视化接口传入着色规则 #

高层级图（HighLevelGraph）可视化 #

# 按layer名称着色，和低层级使用同一套配色规则
hlg.visualize(color=layer_to_color)

低层级图可视化 #

# 低层级任务按所属layer着色，和高层级颜色完全对齐
df.visualize(size="7", color=task_color_map)

补充说明 #

• 可通过修改color_palette自定义配色，也可以调用matplotlib的colormap自动生成适配更多layer的配色
• 复杂工作流可视化时可额外搭配collapse_outputs=True、rankdir="LR"参数优化布局
• 该方法对DataFrame、Array、Bag、自定义Delayed工作流等所有支持可视化的Dask对象都生效

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