Skip to content

Commit 5da6eea

Browse files
committed
adapt tests as per review request, fix ensure_common_dims
1 parent 1cd20c3 commit 5da6eea

File tree

2 files changed

+172
-460
lines changed

2 files changed

+172
-460
lines changed

xarray/core/concat.py

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -536,7 +536,7 @@ def _dataset_concat(
536536

537537
# we've already verified everything is consistent; now, calculate
538538
# shared dimension sizes so we can expand the necessary variables
539-
def ensure_common_dims(vars):
539+
def ensure_common_dims(vars, concat_dim_lengths):
540540
# ensure each variable with the given name shares the same
541541
# dimensions and the same shape for all of them except along the
542542
# concat dimension
@@ -573,23 +573,24 @@ def get_indexes(name):
573573
if name in concat_over and name not in result_indexes:
574574
variables = []
575575
variable_index = []
576+
var_concat_dim_length = []
576577
for i, ds in enumerate(datasets):
577578
if name in ds.variables:
578579
variables.append(ds[name].variable)
579580
# add to variable index, needed for reindexing
580-
variable_index.extend(
581-
[
582-
sum(concat_dim_lengths[:i]) + k
583-
for k in range(concat_dim_lengths[i])
584-
]
585-
)
581+
var_idx = [
582+
sum(concat_dim_lengths[:i]) + k
583+
for k in range(concat_dim_lengths[i])
584+
]
585+
variable_index.extend(var_idx)
586+
var_concat_dim_length.append(len(var_idx))
586587
else:
587588
# raise if coordinate not in all datasets
588589
if name in coord_names:
589590
raise ValueError(
590591
f"coordinate {name!r} not present in all datasets."
591592
)
592-
vars = ensure_common_dims(variables)
593+
vars = ensure_common_dims(variables, var_concat_dim_length)
593594

594595
# Try to concatenate the indexes, concatenate the variables when no index
595596
# is found on all datasets.

0 commit comments

Comments
 (0)