如何将形状与位置可变的Dask delayed对象赋值给Dask数组

需要将多个**形状未知（延迟获取）**的小型Dask数组，赋值到提前知晓形状的大型Dask数组的非顺序2D位置。直接使用dask.delayed包装小型数组和位置时触发错误，而使用普通Dask数组则可正常运行。

示例代码 #

import dask.array as da
import dask

## 初始化大型数组
big_array = da.zeros([5, 6])  # 提前知晓该形状

# 模拟小型数组
aa_shape = dask.delayed((2,3))  # 无法提前知晓该形状
aa = dask.delayed(1 * da.ones(aa_shape))
aa_loc = dask.delayed((slice(0,2), slice(0,3)))  # 无法提前知晓该位置
          
bb_shape = dask.delayed((3,3))
bb = dask.delayed(2 * da.ones(bb_shape))
bb_loc = dask.delayed((slice(0,3), slice(3,6)))

cc_shape = dask.delayed((3,3))
cc = dask.delayed(3 * da.ones(cc_shape))
cc_loc = dask.delayed((slice(2,5), slice(0,3)))

dd_shape = dask.delayed((2,3))
dd = dask.delayed(4 * da.ones(dd_shape))
dd_loc = dask.delayed((slice(3,5), slice(3,6)))

# 手动填充大型数组
big_array[aa_loc] = aa
big_array[bb_loc] = bb
big_array[cc_loc] = cc
big_array[dd_loc] = dd

big_array.compute()

理想输出 #

array([[1., 1., 1., 2., 2., 2.],
       [1., 1., 1., 2., 2., 2.],
       [3., 3., 3., 2., 2., 2.],
       [3., 3., 3., 4., 4., 4.],
       [3., 3., 3., 4., 4., 4.]])

报错信息 #

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Input In [7], in <cell line: 21>()
     18 locs = [aa_loc, bb_loc, cc_loc, dd_loc]
     20 # Manually populate big array
---> 21 big_array[aa_loc] = aa
     22 big_array[bb_loc] = bb
     23 big_array[cc_loc] = cc

File ~/.conda-envs/daskenv202301/lib/python3.9/site-packages/dask/array/core.py:1893, in Array.__setitem__(self, key, value)
   1890 value = asanyarray(value)
   1892 out = "setitem-" + tokenize(self, key, value)
-> 1893 dsk = setitem_array(out, self, key, value)
   1895 meta = meta_from_array(self._meta)
   1896 if np.isscalar(meta):

File ~/.conda-envs/daskenv202301/lib/python3.9/site-packages/dask/array/slicing.py:1754, in setitem_array(out_name, array, indices, value)
   1752 array_shape = array.shape
   1753 value_shape = value.shape
-> 1754 value_ndim = len(value_shape)
   1756 # Reformat input indices
   1757 indices, implied_shape, reverse, implied_shape_positions = parse_assignment_indices(
   1758     indices, array_shape
   1759 )

File ~/.conda-envs/daskenv202301/lib/python3.9/site-packages/dask/delayed.py:591, in Delayed.__len__(self)
    589 def __len__(self):
    590     if self._length is None:
-> 591         raise TypeError("Delayed objects of unspecified length have no len()")
    592     return self._length

TypeError: Delayed objects of unspecified length have no len()

解决方案 #

错误根源是Dask数组的赋值操作需要提前知晓形状和位置的元数据，而dask.delayed对象无法提供这些元信息（无法直接获取len()）。以下是三种可行思路：

方法1：提前解析延迟的形状和位置 #

如果形状、位置的计算开销低，可以先compute()得到实际值，再用普通Dask数组执行赋值：

import dask.array as da
import dask

# 初始化大型数组
big_array = da.zeros([5, 6])

# 解析延迟的形状和位置
aa_shape = dask.delayed((2,3)).compute()
aa = 1 * da.ones(aa_shape)
aa_loc = dask.delayed((slice(0,2), slice(0,3))).compute()

bb_shape = dask.delayed((3,3)).compute()
bb = 2 * da.ones(bb_shape)
bb_loc = dask.delayed((slice(0,3), slice(3,6))).compute()

cc_shape = dask.delayed((3,3)).compute()
cc = 3 * da.ones(cc_shape)
cc_loc = dask.delayed((slice(2,5), slice(0,3))).compute()

dd_shape = dask.delayed((2,3)).compute()
dd = 4 * da.ones(dd_shape)
dd_loc = dask.delayed((slice(3,5), slice(3,6))).compute()

# 执行赋值
big_array[aa_loc] = aa
big_array[bb_loc] = bb
big_array[cc_loc] = cc
big_array[dd_loc] = dd

print(big_array.compute())

方法2：用延迟函数封装全部赋值逻辑 #

如果形状和位置必须延迟计算（依赖其他重计算），可以将整个赋值逻辑封装为延迟函数，最后转为Dask数组：

import dask.array as da
import dask
import numpy as np

def fill_big_array():
    # 内部用numpy完成所有动态赋值逻辑
    big_np = np.zeros([5,6])
    
    # 模拟动态获取形状、数组和位置
    aa_shape = (2,3)  # 替换为实际延迟计算逻辑
    aa = 1 * np.ones(aa_shape)
    aa_loc = (slice(0,2), slice(0,3))
    big_np[aa_loc] = aa
    
    bb_shape = (3,3)
    bb = 2 * np.ones(bb_shape)
    bb_loc = (slice(0,3), slice(3,6))
    big_np[bb_loc] = bb
    
    cc_shape = (3,3)
    cc = 3 * np.ones(cc_shape)
    cc_loc = (slice(2,5), slice(0,3))
    big_np[cc_loc] = cc
    
    dd_shape = (2,3)
    dd = 4 * np.ones(dd_shape)
    dd_loc = (slice(3,5), slice(3,6))
    big_np[dd_loc] = dd
    
    return big_np

# 将延迟函数转为Dask数组，指定输出形状和类型
big_array = da.from_delayed(dask.delayed(fill_big_array()), shape=(5,6), dtype=np.float64)
print(big_array.compute())

方法3：分块处理赋值（适合分布式场景） #

如果大型数组采用分块存储，可针对每个分块编写赋值逻辑，用map_blocks执行：

import dask.array as da
import numpy as np

def assign_block(block, block_info):
    # 获取当前分块在大型数组中的起始位置
    block_start = block_info[0]['chunk-location']
    
    # 根据分块位置匹配对应小型数组的赋值逻辑
    if block_start == (0,0):
        # 填充aa和bb的部分区域
        block[:2, :3] = 1 * np.ones((2,3))
        block[:2, 3:] = 2 * np.ones((2,3))
    elif block_start == (0,1):
        # 填充bb的剩余区域
        block[:3, :] = 2 * np.ones((3,3))
    elif block_start == (1,0):
        # 填充cc的区域
        block[:, :3] = 3 * np.ones((3,3))
    elif block_start == (1,1):
        # 填充cc的剩余区域和dd的区域
        block[:2, :3] = 3 * np.ones((2,3))
        block[:2, 3:] = 4 * np.ones((2,3))
    return block

# 初始化分块的大型数组
big_array = da.zeros([5,6], chunks=(2,3))
# 对每个分块应用赋值函数
big_array = big_array.map_blocks(assign_block, meta=np.float64)
print(big_array.compute())

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