assign_coords with mixed DataArray / array args removes coords

[keewis]

I'm not sure if using assign_coords to overwrite the data of coords is the best way to do so, but using mixed args (on current master) turns out to have surprising results:

>>> obj = xr.DataArray(
...     data=[6, 3, 4, 6],
...     coords={"x": list("abcd"), "y": ("x", range(4))},
...     dims="x",
... )
>>> obj
<xarray.DataArray 'obj' (x: 4)>
array([6, 3, 4, 6])
Coordinates:
  * x        (x) <U1 'a' 'b' 'c' 'd'
    y        (x) int64 0 1 2 3
>>> # works as expected
>>> obj.assign_coords(coords={"x": list("efgh"), "y": ("x", [0, 2, 4, 6])})
<xarray.DataArray 'obj' (x: 4)>
array([6, 3, 4, 6])
Coordinates:
  * x        (x) <U1 'e' 'f' 'g' 'h'
    y        (x) int64 0 2 4 6
>>> # works, too (same as .data / .values)
>>> obj.assign_coords(coords={
...     "x": obj.x.copy(data=list("efgh")).variable,
...     "y": ("x", [0, 2, 4, 6]),
... })
<xarray.DataArray 'obj' (x: 4)>
array([6, 3, 4, 6])
Coordinates:
  * x        (x) <U1 'e' 'f' 'g' 'h'
    y        (x) int64 0 2 4 6
>>> # this drops "y"
>>> obj.assign_coords(coords={
...     "x": obj.x.copy(data=list("efgh")),
...     "y": ("x", [0, 2, 4, 6]),
... })
<xarray.DataArray 'obj' (x: 4)>
array([6, 3, 4, 6])
Coordinates:
  * x        (x) <U1 'e' 'f' 'g' 'h'

Passing a DataArray for y, like obj.y * 2 while also changing x (the type does not matter) always results in a MergeError:

>>> obj.assign_coords(x=list("efgh"), y=obj.y * 2)
xarray.core.merge.MergeError: conflicting values for index 'x' on objects to be combined:
first value: Index(['e', 'f', 'g', 'h'], dtype='object', name='x')
second value: Index(['a', 'b', 'c', 'd'], dtype='object', name='x')

I would expect the result to be the same regardless of the type of the new coords.

[dcherian]

The x=DataArray example looks like a bug.

The last example raises a MergeError regardless of the order of x, y which is good I think. It's pretty ambiguous to change both the index x and the y=function(x) at the same time. It will only work if xarray uses magic and infers that everything but x must be assigned before assigning x.

[keewis]

would it not be possible to always convert DataArray args using .variable? Unless I'm missing something, this should keep the information on the dimensions while also allowing to change the data independently:

obj.assign_coords(x=obj.x.copy(data=list("efgh")).variable, y=(obj.y * 2).variable)
obj.assign_coords(x=list("efgh"), y=(obj.y * 2).variable)

both work

[dcherian]

That is a good point. Let's see what @shoyer thinks.

[shoyer]

would it not be possible to always convert DataArray args using .variable? Unless I'm missing something, this should keep the information on the dimensions while also allowing to change the data independently:

It might make sense to change how assign_coords/ds.coords[k] = v] works, but in general we do want assign/ds[k] = v to add in coordinates from assigned DataArray objects rather than dropping them silently.

[keewis]

agreed, my suggestion only applies to assign_coords, and even in that function it might make sense to keep the old behaviour unless the coord was explicitly set.

Edit: this could be fixed by changing how merge_coords and _get_priority_vars_and_indexes work, but they look like magic to me