Add template xarray object kwarg to map_blocks

xarray/core/parallel.py

@@ -107,6 +107,7 @@ def map_blocks(
     obj: Union[DataArray, Dataset],
     args: Sequence[Any] = (),
     kwargs: Mapping[str, Any] = None,
+    template: T_DSorDA = None,
 ) -> T_DSorDA:
     """Apply a function to each chunk of a DataArray or Dataset. This function is
     experimental and its signature may change.
@@ -119,14 +120,10 @@ def map_blocks(
         corresponding to one chunk along each chunked dimension. ``func`` will be
         executed as ``func(obj_subset, *args, **kwargs)``.
 
-        The function will be first run on mocked-up data, that looks like 'obj' but
-        has sizes 0, to determine properties of the returned object such as dtype,
-        variable names, new dimensions and new indexes (if any).
-
         This function must return either a single DataArray or a single Dataset.
 
-        This function cannot change size of existing dimensions, or add new chunked
-        dimensions.
+        This function cannot add a new chunked dimension.
+
     obj: DataArray, Dataset
         Passed to the function as its first argument, one dask chunk at a time.
     args: Sequence
@@ -135,6 +132,12 @@ def map_blocks(
     kwargs: Mapping
         Passed verbatim to func after unpacking. xarray objects, if any, will not be
         split by chunks. Passing dask collections is not allowed.
+    template: (optional) DataArray, Dataset
+        xarray object representing the final result after compute is called. If not provided,
+        the function will be first run on mocked-up data, that looks like 'obj' but
+        has sizes 0, to determine properties of the returned object such as dtype,
+        variable names, new dimensions and new indexes (if any).
+        'template' must be provided if the function changes the size of existing dimensions.
 
     Returns
     -------
@@ -198,18 +201,26 @@ def map_blocks(
         * time     (time) object 1990-01-31 00:00:00 ... 1991-12-31 00:00:00
     """
 
-    def _wrapper(func, obj, to_array, args, kwargs):
+    def _wrapper(func, obj, to_array, args, kwargs, expected_shapes):
+        check_shapes = dict(obj.dims)
+        check_shapes.update(expected_shapes)
+
         if to_array:
             obj = dataset_to_dataarray(obj)
 
         result = func(obj, *args, **kwargs)
 
+        missing_dimensions = set(expected_shapes) - set(result.sizes)
+        if missing_dimensions:
+            raise ValueError(
+                f"Dimensions {missing_dimensions} missing on returned object."
+            )
+
         for name, index in result.indexes.items():
-            if name in obj.indexes:
-                if len(index) != len(obj.indexes[name]):
+            if name in check_shapes:
+                if len(index) != check_shapes[name]:
                     raise ValueError(
-                        "Length of the %r dimension has changed. This is not allowed."
-                        % name
+                        f"Received dimension {name} of length {len(index)}. Expected length {expected_shapes[name]}."
                     )
 
         return make_dict(result)
@@ -245,8 +256,24 @@ def _wrapper(func, obj, to_array, args, kwargs):
         input_is_array = False
 
     input_chunks = dataset.chunks
+    dataset_indexes = set(dataset.indexes)
+    if template is None:
+        # infer template by providing zero-shaped arrays
+        template = infer_template(func, obj, *args, **kwargs)
+        template_indexes = set(template.indexes)
+        preserved_indexes = template_indexes & dataset_indexes
+        new_indexes = template_indexes - dataset_indexes
+        indexes = {dim: dataset.indexes[dim] for dim in preserved_indexes}
+        indexes.update({k: template.indexes[k] for k in new_indexes})
+        output_chunks = input_chunks
+
+    else:
+        # template xarray object has been provided with proper sizes and chunk shapes
+        template_indexes = set(template.indexes)
+        indexes = {dim: dataset.indexes[dim] for dim in dataset_indexes}
+        indexes.update({k: template.indexes[k] for k in template_indexes})
+        output_chunks = template.chunks
 
-    template: Union[DataArray, Dataset] = infer_template(func, obj, *args, **kwargs)
     if isinstance(template, DataArray):
         result_is_array = True
         template_name = template.name
@@ -258,13 +285,6 @@ def _wrapper(func, obj, to_array, args, kwargs):
             f"func output must be DataArray or Dataset; got {type(template)}"
         )
 
-    template_indexes = set(template.indexes)
-    dataset_indexes = set(dataset.indexes)
-    preserved_indexes = template_indexes & dataset_indexes
-    new_indexes = template_indexes - dataset_indexes
-    indexes = {dim: dataset.indexes[dim] for dim in preserved_indexes}
-    indexes.update({k: template.indexes[k] for k in new_indexes})
-
     # We're building a new HighLevelGraph hlg. We'll have one new layer
     # for each variable in the dataset, which is the result of the
     # func applied to the values.
@@ -284,7 +304,7 @@ def _wrapper(func, obj, to_array, args, kwargs):
 
     # iterate over all possible chunk combinations
     for v in itertools.product(*ichunk.values()):
-        chunk_index_dict = dict(zip(dataset.dims, v))
+        input_chunk_index = dict(zip(dataset.dims, v))
 
         # this will become [[name1, variable1],
         #                   [name2, variable2],
@@ -299,7 +319,7 @@ def _wrapper(func, obj, to_array, args, kwargs):
                 # recursively index into dask_keys nested list to get chunk
                 chunk = variable.__dask_keys__()
                 for dim in variable.dims:
-                    chunk = chunk[chunk_index_dict[dim]]
+                    chunk = chunk[input_chunk_index[dim]]
 
                 chunk_variable_task = (f"{gname}-{chunk[0]}",) + v
                 graph[chunk_variable_task] = (
@@ -311,8 +331,8 @@ def _wrapper(func, obj, to_array, args, kwargs):
                 # index into variable appropriately
                 subsetter = {}
                 for dim in variable.dims:
-                    if dim in chunk_index_dict:
-                        which_chunk = chunk_index_dict[dim]
+                    if dim in input_chunk_index:
+                        which_chunk = input_chunk_index[dim]
                         subsetter[dim] = slice(
                             chunk_index_bounds[dim][which_chunk],
                             chunk_index_bounds[dim][which_chunk + 1],
@@ -333,6 +353,11 @@ def _wrapper(func, obj, to_array, args, kwargs):
             else:
                 data_vars.append([name, chunk_variable_task])
 
+        # input chunk 0 along a dimension maps to output chunk 0 along the same dimension
+        # even if length of dimension is changed by the applied function
+        # expected_shapes is used to raise nice error messages in _wrapper
+        expected_shapes = {k: output_chunks[k][v] for k, v in input_chunk_index.items()}
+
         from_wrapper = (gname,) + v
         graph[from_wrapper] = (
             _wrapper,
@@ -341,6 +366,7 @@ def _wrapper(func, obj, to_array, args, kwargs):
             input_is_array,
             args,
             kwargs,
+            expected_shapes,
         )
 
         # mapping from variable name to dask graph key
@@ -353,8 +379,12 @@ def _wrapper(func, obj, to_array, args, kwargs):
 
             key: Tuple[Any, ...] = (gname_l,)
             for dim in variable.dims:
-                if dim in chunk_index_dict:
-                    key += (chunk_index_dict[dim],)
+                if dim in input_chunk_index:
+                    key += (input_chunk_index[dim],)
+                elif dim in output_chunks:
+                    raise ValueError(
+                        f"Function is attempting to add a new chunked dimension {dim}. This is not allowed."
+                    )
                 else:
                     # unchunked dimensions in the input have one chunk in the result
                     key += (0,)
@@ -379,10 +409,13 @@ def _wrapper(func, obj, to_array, args, kwargs):
         dims = template[name].dims
         var_chunks = []
         for dim in dims:
-            if dim in input_chunks:
-                var_chunks.append(input_chunks[dim])
+            if dim in output_chunks:
+                var_chunks.append(output_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

@@ -1039,7 +1039,7 @@ def test_map_blocks_error(map_da, map_ds):
     def bad_func(darray):
         return (darray * darray.x + 5 * darray.y)[:1, :1]
 
-    with raises_regex(ValueError, "Length of the.* has changed."):
+    with raises_regex(ValueError, "Received dimension.*"):
         xr.map_blocks(bad_func, map_da).compute()
 
     def returns_numpy(darray):
@@ -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)],
     ],