Skip to content

Fix concat with variable or dataarray as dim (propagate attrs) #6387

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 20, 2022
Merged

Fix concat with variable or dataarray as dim (propagate attrs) #6387

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion xarray/core/concat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
11 changes: 9 additions & 2 deletions xarray/tests/test_concat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down