diff --git a/xarray/core/concat.py b/xarray/core/concat.py index 8ee4672c49a..e2525c435af 100644 --- a/xarray/core/concat.py +++ b/xarray/core/concat.py @@ -429,6 +429,7 @@ def _dataset_concat( """ Concatenate a sequence of datasets along a new or existing dimension """ + from .dataarray import DataArray from .dataset import Dataset datasets = list(datasets) @@ -438,6 +439,13 @@ def _dataset_concat( "The elements in the input list need to be either all 'Dataset's or all 'DataArray's" ) + if isinstance(dim, DataArray): + dim_var = dim.variable + elif isinstance(dim, Variable): + dim_var = dim + else: + dim_var = None + dim, index = _calc_concat_dim_index(dim) # Make sure we're working on a copy (we'll be loading variables) @@ -582,7 +590,11 @@ def get_indexes(name): if index is not None: # add concat index / coordinate last to ensure that its in the final Dataset - result[dim] = index.create_variables()[dim] + if dim_var is not None: + index_vars = index.create_variables({dim: dim_var}) + else: + index_vars = index.create_variables() + result[dim] = index_vars[dim] result_indexes[dim] = index # TODO: add indexes at Dataset creation (when it is supported) diff --git a/xarray/tests/test_concat.py b/xarray/tests/test_concat.py index 8abede64761..44393dbcb1a 100644 --- a/xarray/tests/test_concat.py +++ b/xarray/tests/test_concat.py @@ -459,8 +459,15 @@ def test_concat_do_not_promote(self) -> None: def test_concat_dim_is_variable(self) -> None: objs = [Dataset({"x": 0}), Dataset({"x": 1})] - coord = Variable("y", [3, 4]) - expected = Dataset({"x": ("y", [0, 1]), "y": [3, 4]}) + coord = Variable("y", [3, 4], attrs={"foo": "bar"}) + expected = Dataset({"x": ("y", [0, 1]), "y": coord}) + actual = concat(objs, coord) + assert_identical(actual, expected) + + def test_concat_dim_is_dataarray(self) -> None: + objs = [Dataset({"x": 0}), Dataset({"x": 1})] + coord = DataArray([3, 4], dims="y", attrs={"foo": "bar"}) + expected = Dataset({"x": ("y", [0, 1]), "y": coord}) actual = concat(objs, coord) assert_identical(actual, expected)