From 045ae2b1bf939515e0a38c960d0cdc7974bcfa37 Mon Sep 17 00:00:00 2001 From: dcherian Date: Tue, 25 Feb 2020 20:35:46 -0700 Subject: [PATCH] Allow user function to add new unindexed dimension. --- xarray/core/parallel.py | 3 +++ xarray/tests/test_dask.py | 2 ++ 2 files changed, 5 insertions(+) diff --git a/xarray/core/parallel.py b/xarray/core/parallel.py index 17f13576417..2f4de742c4a 100644 --- a/xarray/core/parallel.py +++ b/xarray/core/parallel.py @@ -401,6 +401,9 @@ def _wrapper(func, obj, to_array, args, kwargs, expected_shapes): 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 diff --git a/xarray/tests/test_dask.py b/xarray/tests/test_dask.py index 857b37672b7..76920a1498a 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)], ],