map_blocks: allow user function to add new unindexed dimension.

doc/whats-new.rst

@@ -39,6 +39,8 @@ New Features
   often means a user is attempting to pass multiple dimensions to group over
   and should instead pass a list.
   By `Maximilian Roos <https://github.com/max-sixty>`_
+- :py:func:`map_blocks` can now apply functions that add new unindexed dimensions.
+  By `Deepak Cherian <https://github.com/dcherian>`_
 - The new ``Dataset._repr_html_`` and ``DataArray._repr_html_`` (introduced
   in 0.14.1) is now on by default. To disable, use
   ``xarray.set_options(display_style="text")``.
@@ -51,7 +53,6 @@ New Features
   (:issue:`3843`, :pull:`3844`) 
   By `Aaron Spring <https://github.com/aaronspring>`_.
 
-
 Bug fixes
 ~~~~~~~~~
 - Fix :py:meth:`Dataset.interp` when indexing array shares coordinates with the

xarray/core/parallel.py

@@ -386,6 +386,9 @@ def _wrapper(func, obj, to_array, args, kwargs):
                 var_chunks.append(input_chunks[dim])
             elif dim in indexes:
                 var_chunks.append((len(indexes[dim]),))
+            elif dim in template.dims:
+                # new unindexed dimension
+                var_chunks.append((template.sizes[dim],))
 
         data = dask.array.Array(
             hlg, name=gname_l, chunks=var_chunks, dtype=template[name].dtype

xarray/tests/test_dask.py

@@ -1147,6 +1147,7 @@ def test_map_blocks_to_array(map_ds):
         lambda x: x.to_dataset(),
         lambda x: x.drop_vars("x"),
         lambda x: x.expand_dims(k=[1, 2, 3]),
+        lambda x: x.expand_dims(k=3),
         lambda x: x.assign_coords(new_coord=("y", x.y * 2)),
         lambda x: x.astype(np.int32),
         # TODO: [lambda x: x.isel(x=1).drop_vars("x"), map_da],
@@ -1167,6 +1168,7 @@ def test_map_blocks_da_transformations(func, map_da):
         lambda x: x.drop_vars("a"),
         lambda x: x.drop_vars("x"),
         lambda x: x.expand_dims(k=[1, 2, 3]),
+        lambda x: x.expand_dims(k=3),
         lambda x: x.rename({"a": "new1", "b": "new2"}),
         # TODO: [lambda x: x.isel(x=1)],
     ],