从 dask 数组生成向量输出

Question

我有一个大型 dask 数组 (labeled_arr)，它实际上是一个带标签的光栅图像（dtype 是 int64）。我想使用 rasterio 将标记的区域转换为多边形并将它们组合成一个多边形列表（或只有一个几何列的 geoseries）。这是单个数组上的一项简单任务，但我无法弄清楚如何告诉 dask 我希望它对每个块执行此操作并返回不是数组的内容。

应用于每个块的函数：

def get_polys(labeled_blocks):
    polys = list(poly[0]['coordinates'][0] for poly in rasterio.features.shapes(
                                labeled_blocks.astype('int32'), transform=trans))[:-1]
    # Note: rasterio.features.shapes returns an iterator, hence the conversion to a list here
    return polys

试图让 dask 执行此操作的代码行：

test_polygons = da.blockwise(get_polys, '', labeled_arr, 'ij')
test_polygons.compute()

labeled_arr 是输入的分块 dask 数组。

按原样运行会返回一个错误，提示我必须为da.blockwise 指定一个数据类型。指定 dtype 会返回 AttributeError，因为输出列表类型没有 dtype 属性。我发现了 meta 关键字，但仍然无法获得正确的语法来将我的输出转换为系列或列表。

我不喜欢上述方法，但我的首要目标是：获取一个标记的、分块的 dask 数据数组（它并不都适合内存），根据每个块的计算提取一个列表，并生成一个连接的列出原始分块数组中所有块的输出（或 pandas 数据对象）。

Answer 1

这可能有效：

import dask
import dask.array as da

# we expect to see 4 blocks here
test_array = da.random.random((4, 4), chunks=(2, 2))

@dask.delayed
def my_func(block):
    # do something fancy
    return list(block)

results = dask.compute([my_func(x) for x in test_array.to_delayed().ravel()])

如您所述，问题在于list 没有dtype。解决此问题的一种方法是将list 转换为np.array，但我不确定这是否适用于所有geometry 对象（对于Points 应该没问题，但由于多边形可能会出现问题到不同的长度）。由于您对将这些几何图形强制放入数组不感兴趣，因此最好将单个块视为 delayed 对象，一次将它们输入您的函数（但在工作人员/进程之间进行缩放）。

Answer 2

这是我最初最终得到的解决方案，尽管考虑到 concatenate=True kwarg，它仍然需要大量 RAM。

poss_list = []
def get_polys(labeled_blocks):
    polys = list(poly[0]['coordinates'][0] for poly in rasterio.features.shapes(
                        labeled_blocks.astype('int32'), transform=trans))[:-1]
    poss_list.append(polys)
        
da.blockwise(get_bergs, '', labeled_arr, 'ij', 
                meta=pd.DataFrame({'c':[]}), concatenate=True).compute()

如果我的解释正确，这不会将这些块提供给我的跨工作程序/进程的函数（现在看来我可以侥幸逃脱）。

更新 - 使用 dask.delayed 改进了答案，基于 @SultanOrazbayev 接受的答案

import dask
# onedem = original_xarray_dataarray
poss_list = []

@dask.delayed
def get_bergs(labeled_blocks, pointer, chunk0, chunk1):
            
    # Note: I'm using this in a CRS (polar stereo) with negative y coordinates - it hasn't been tested for other CRSs
    def getpx(chunkid, chunksz):
       amin = chunkid[0] * chunksz[0][0]
       amax = amin + chunksz[0][0]
       bmin = chunkid[1] * chunksz[1][0]
       bmax = bmin + chunksz[1][0]
       return (amin, amax, bmin, bmax)

    # order of all inputs (and outputs) should be y, x when axis order is used
    chunksz = (onedem.chunks['y'], onedem.chunks['x'])
    ymini, ymaxi, xmini, xmaxi = getpx((chunk0, chunk1), chunksz) 

    # use rasterio Windows and rioxarray to construct transform
    # https://rasterio.readthedocs.io/en/latest/topics/windowed-rw.html#window-transforms
    chwindow = rasterio.windows.Window(xmini, ymini, xmaxi-xmini, ymaxi-ymini) #.from_slices[ymini, ymaxi],[xmini, xmaxi])
    trans = onedem.rio.isel_window(chwindow).rio.transform(recalc=True)

    return list(poly[0]['coordinates'][0] for poly in rasterio.features.shapes(labeled_blocks.astype('int32'), transform=trans))[:-1]

        
for __, obj in enumerate(labeled_arr.to_delayed()):
   for bl in obj:
      piece = dask.delayed(get_bergs)(bl, *bl.key)
      poss_list.append(piece)
        
poss_list = dask.compute(*poss_list)
        
# unnest the list of polygons returned by using dask to polygonize
concat_list = [item for sublist in poss_list for item in sublist if len(item)!=0]