From 14b0460c7364f1da7f2a07c428c244eae0a91d2a Mon Sep 17 00:00:00 2001 From: dcherian Date: Tue, 25 Feb 2020 20:35:46 -0700 Subject: [PATCH] map_blocks: allow user function to add new unindexed dimension. --- doc/whats-new.rst | 2 ++ xarray/core/parallel.py | 3 +++ xarray/tests/test_dask.py | 2 ++ 3 files changed, 7 insertions(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 1deb77eecfc..2b9b8ebdc4f 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -35,6 +35,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 `_ +- :py:func:`map_blocks` can now apply functions that add new unindexed dimensions. + By `Deepak Cherian `_ Bug fixes ~~~~~~~~~ diff --git a/xarray/core/parallel.py b/xarray/core/parallel.py index facfa06b23c..e294b25dab0 100644 --- a/xarray/core/parallel.py +++ b/xarray/core/parallel.py @@ -383,6 +383,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 diff --git a/xarray/tests/test_dask.py b/xarray/tests/test_dask.py index 8fb54c4ee84..cc7315fa3d7 100644 --- a/xarray/tests/test_dask.py +++ b/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)], ],