fillna removes non-indexed coordinates in some cases

What happened?

fillna (and other functions which need aligning) remove non-indexed extra coordinates (along a dimensions which also have an indexed coordinates)

Below I have an example with fillna, But I have experienced similar things with combine_first and assume it will happen whenever alignment is needed. More on what I think is happening below.

What did you expect to happen?

I expected the extra coordinate to remain in place as it was defined consistently in both datasets. Alternatively an error message would also help to understand what's going on. Currently it's dropped silently.

Minimal Complete Verifiable Example

import xarray as xr
import numpy as np

# create a dataset wit a few dimensions and random data
area_iso3 = np.array(["COL", "ARG", "MEX", "BOL"])

test_ds = xr.Dataset(
    { "CO2": 
          xr.DataArray(data=np.ones(len(area_iso3)),
                       coords={
                           "area (ISO3)": area_iso3,
                       },
                       dims=["area (ISO3)"])
    }
) 

# attach an additional coordinate the existing dimensions
country_names = ["Colombia", "Argentina", "Mexico", "Bolovia"]
test_ds = test_ds.assign_coords(country_name=("area (ISO3)", country_names))

#### use a loc to fill - dim with additional coordinate involved
test_ds_loc = test_ds.loc[{'area (ISO3)': ['COL', 'ARG']}]
print(f"test_ds_loc coords: {test_ds_loc.coords}")

# set some values to nan to fill later
test_ds["CO2"].loc[{'area (ISO3)': ['COL', 'ARG']}] = np.nan

# fill
test_ds = test_ds.fillna(test_ds_loc)
print(f"test_ds coords after fillna: {test_ds.coords}\n")
# additional coordinate gone

MVCE confirmation

• Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray.
• Complete example — the example is self-contained, including all data and the text of any traceback.
• Verifiable example — the example copy & pastes into an IPython prompt or Binder notebook, returning the result.
• New issue — a search of GitHub Issues suggests this is not a duplicate.
• Recent environment — the issue occurs with the latest version of xarray and its dependencies.

Relevant log output

No error is raise and no log output generated. My output of the example code is:

test_ds_loc coords: Coordinates:
  * area (ISO3)   (area (ISO3)) <U3 24B 'COL' 'ARG'
    country_name  (area (ISO3)) <U9 72B 'Colombia' 'Argentina'
test_ds coords after fillna: Coordinates:
  * area (ISO3)  (area (ISO3)) <U3 48B 'COL' 'ARG' 'MEX' 'BOL'

Anything else we need to know?

The problem only occurs when the .loc is done on the dimension with the additional coordinate. I think the reason for the problem is the following:

align fills the additional coordinate from the smaller dataset (same for dataarray) with np.nan to expand it to the larger dataset. Later in the process the aligned coordinates are passed to merge_coordinates_without_align which removes the additional coordinate as it contains np.nan in one of the datasets where the other dataset has values.

So the non-index coordinates are neither combined like the indexed coordinates, nor filled like the data variables.

Below a short example involving only align and merge_coordinates_without_align.

import xarray as xr
from xarray.core.merge import merge_coordinates_without_align
import numpy as np

# create a dataset wit a few dimensions and random data
area_iso3 = np.array(["COL", "ARG", "MEX", "BOL"])

test_ds = xr.Dataset(
    { "CO2":
          xr.DataArray(data=np.ones(4),
                       coords={
                           "area (ISO3)": area_iso3,
                       },
                       dims=["area (ISO3)"])
      }
)

# attach an additional coordinate to one of the existing dimensions
country_names = ["Colombia", "Argentina", "Mexico", "Bolovia"]
test_ds = test_ds.assign_coords(country_name=("area (ISO3)", country_names))

test_ds_loc = test_ds.loc[{'area (ISO3)': ['COL', 'ARG']}]
print(f"test_ds_loc coords: {test_ds_loc.coords}")

aligned = xr.align(test_ds,test_ds_loc,join='outer')
merged = merge_coordinates_without_align(aligned)
print(f"merged coords: {merged[0]}")

Environment

Reproduce the minimal example, then start with xr.align and xarray.core.merge.merge_coordinates_without_align, the entry points identified in the issue. Trace how the aligned coordinates are merged and add a regression test covering fillna after the .loc selection; done means the non-indexed country_name coordinate is retained without silent loss.