produce vector output from a dask array

Question

I have a large dask array (labeled_arr) that is actually a labeled raster image (dtype is int64). I want to use rasterio to turn the labeled regions into polygons and combine them into a single list of polygons (or geoseries with just a geometry column). This is a straightforward task on a single array, but I'm having trouble figuring out how to tell dask that I want it to do this operation on each chunk and return something that is not an array.

function to apply to each chunk:

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

line of code trying to get dask to do this:

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

where labeled_arr is the input chunked dask array.

Running as is returns an error saying I have to specify a dtype for da.blockwise. Specifying a dtype returns an AttributeError since the output list type does not have a dtype attribute. I discovered the meta keyword, but still have been unable to get the right syntax to turn my output into a Series or list.

I'm not attached to the above approach, but my overarching goal is: take a labeled, chunked dask dataarray (which does not all fit in memory), extract a list based on computations for each chunk, and generate a concatenated list (or pandas data object) with the outputs from all the chunks in my original chunked array.

Answer 1

This might work:

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()])

As you noted, the problem is that list has no dtype. A way around this would be to convert the list into a np.array, but I'm not sure if this will work with all geometry objects (it should be OK for Points, but polygons might be problematic due to varying length). Since you are not interested in forcing these geometries into an array, it's best to treat individual blocks as delayed objects feeding them into your function one at a time (but scaled across workers/processes).

Answer 2

Here's the solution I ended up with initially, though it still requires a lot of RAM given the concatenate=True kwarg.

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()

If I'm interpreting correctly, this doesn't feed the chunks into my function across workers/processes though (which it seems I can get away with for now).

Update - improved answer using dask.delayed, building on the accepted answer by @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]