Map_blocks error: zip argument #2 must support iteration

[tlsw231]

I am trying to use map_blocks to: ingest a multi-dimensional array as input, reduce along one dimension and add extra dimensions to the output. Is this possible? I am attaching a simple MRE below that gives me an zip argument #2 must support iteration error. Any pointers on what I might be doing wrong?
[My real example is a 3d-dataset with (time,lat,lon) dimensions and I am trying to reduce along time while adding two new dimensions to the output. I tried so many things and got so many errors, including the one in the title, that I thought it is better to first understand how map_blocks works!]

# The goal is to feed in a 2d array, reduce along one dimension and add two new dimensions to the output. 

chunks={}
dummy = xr.DataArray(data=np.random.random([8,100]),dims=['dim1','dim2']).chunk(chunks)

def some_func(func):
        
    dims=func.dims

    n1 = len(func[func.dims[1]])  # This is 'dim2', we will average along 'dim1' below in the for loop
    
    newdim1 = 2; newdim2 = 5;
    
    output = xr.DataArray(np.nan*np.ones([n1,newdim1,newdim2]),dims=[dims[1],'new1','new2'])
    
    for n in range(n1):
        fmean = func.isel(dim2=n).mean(dims[0]).compute()
        for i in range(newdim1):
            for j in range(newdim2):
                output[n,i,j] = fmean
                 
    return output
    
#out = some_func(dummy)  # This works
    
template=xr.DataArray(np.nan*np.ones([len(dummy.dim2),2,5]),
                      dims=['dim2','new1','new2'])
out = xr.map_blocks(some_func,dummy,template=template).compute()  # gives me the error message in the title

[Edit: Fixed a typo in the n1 = len(func[func.dims[1]]) line, of course getting the same error.]

[dcherian]

template.chunks is None. It needs to have the expected output chunking.

A PR raising a nice error message here when if not template.chunksizes: would be very welcome

We can also delete the if clause here:

xarray/xarray/core/parallel.py

Lines 383 to 388 in ce40b93

	if isinstance(template, DataArray):
	output_chunks = dict(
	zip(template.dims, template.chunks) # type: ignore[arg-type]
	)
	else:
	output_chunks = dict(template.chunks)

and use the new chunksizes property: output_chunks = template.chunksizes

[tlsw231]

Thanks! that worked. A couple of follow-up questions:

1. The input and the template need to have the same chunks for any common dimensions (dim2 in my case) other than the one along which I am doing the reducing (dim1), am I correct?
2. Can I chunk the dimension along which I am doing the reducing (dim1)? I tried doing that for the input and the map_blocks step goes through but fails the np.testing.assert_allclose test.

I will check the docs on how to submit a PR.

[tlsw231]

[More follow-up questions. Start a new thread or continue here?]
Moving beyond the toy example, the good news is I am able to get map_blocks to work with small chunks of the original netcdf dataset that has dimensions (ncol,time) where ncol is a linear index running through all (lat,lon) values on a global grid and time is the dimension along which I am doing the reduction within the function. The size of ncol is huge (~1 million), so here is my plan:

1. Split the problem into smaller ncol intervals.
2. Run map_blocks successively on each ncol interval and then stitch the all output files later using xr.open_mfdataset(...,concat_dim='ncol')
3. I can get map_blocks to work with a size of 5000 for ncol but this is impractical for a total size of 1 million.
4. I would like to be able to run it for a size of 50000 but am hitting a snag right away as xarray is taking forever to generate the task graph. I have now been waiting for 40 minutes after running the command and there is zero activity in the task bar. Here is the relevant code:

[I started with chunksize=100 and have also tried chunksize=50, no difference]

# Select only a subset of `ncol` values
l1 = 0
l2 = 50000
subset=dict(ncol=slice(l1,l2))

# Specify the chunksize (set to 20)
chunks=dict(ncol=chunksize,time=-1)

# Subset the input data and chunk using the chunks specified above
field = data.isel(subset).chunk(chunks)

# Specify chunks again as our output will not have the time dimension
chunks=dict(ncol=chunksize)  
# The dimensions newdim1 and newdim2 are the new dimensions that are being added by the function
# similar to what the toy function does
output = np.nan * xr.DataArray(np.ones([len(field.ncol),newdim1,newdim2]),
                                                   dims=['ncol','newdim1','newdim2'] ).chunk(chunks)
template = xr.full_like(output,fill_value=np.nan).chunk(chunks)

# The step below takes forever for an ncol size of 50000 and I just have to kill the kernel
output = xr.map_blocks(function,field,template=output).compute()
output['ncol'] = np.arange(l1,l2)

[tlsw231]

[Update: Actually the tasks do start after 30(!) minutes and I can finally see the progress on the dashboard. Still feels like less than optimal.]

I tried field = field.persist() before the map_blocks call. Same result. Nothing is happening in the dashboard and I have to kill the kernel. I do get this warning in the jupyter window, not sure how to address this as I have tried different values of chunksize and even persist().

UserWarning: Large object of size 4.06 MiB detected in task graph: 
  ([('time',), <xarray.IndexVariable 'time' (time: 7 ... =object), {}],)
Consider scattering large objects ahead of time
with client.scatter to reduce scheduler burden and 
keep data on workers

[dcherian]

Please make your example reproducible.

map_blocks isn't the most performant way to do things in general, though it can be most convenient. Take a look at apply_ufunc (https://docs.xarray.dev/en/stable/user-guide/dask.html#automatic-parallelization-with-apply-ufunc-and-map-blocks)

[tlsw231]

Thanks for suggesting apply_ufunc. In the past I have had trouble with apply_ufunc even for simple tasks, e.g, this unresolved thread with a MVE (#6322). I also have some questions on how to use apply_ufunc with dask=allowed, will post a separate query.