Fix based on sugggestion

Author: ahuang11

xarray/core/dataset.py

@@ -7170,14 +7170,6 @@ def drop_duplicate_coords(
                 )
             elif coord not in new.dims:
                 dims = new[coord].dims
-                for dim in dims:
-                    if any(new.get_index(dim).duplicated()):
-                        raise ValueError(
-                            f"Cannot have duplicate dimension values in "
-                            f"'{dim}' when dropping duplicate coordinate "
-                            f"values for '{coord}' due to ambiguity "
-                            f"in unstack."
-                        )
 
                 # stack the coord's dimensions
                 tmp_dim = "__tmp_dim__"
@@ -7190,22 +7182,29 @@ def drop_duplicate_coords(
                 # can actually find duplicates
                 new[tmp_dim] = new[coord].values.ravel()
 
+                base_coord = coord
+
                 # replace coord with tmp_dim for use with get_index()
                 # as to not repeat the same call in the dimension clause
                 coord = tmp_dim
-
-                unstack_tmp_dim = True
             else:
-                unstack_tmp_dim = False
+                base_coord = None
 
             index = new.get_index(coord).duplicated(keep=keep)
             new = new.isel({coord: ~index})
 
-            if unstack_tmp_dim:
-                # return everything to normal
-                new[coord] = stacked_coord_indices.isel({coord: ~index})
-                new = new.unstack(coord)
+            if base_coord is not None:
+                # remove tmp_dim
+                new = new.swap_dims({tmp_dim: base_coord}).drop(tmp_dim)
+
+                # get associated coordinates with the stacked dim
+                tmp_index = stacked_coord_indices.isel({coord: ~index}).indexes[tmp_dim]
 
+                # unpack the coordinates and add back to dataset
+                keys = tmp_index.names
+                values = list(zip(*tmp_index.values))
+                for key, value in dict(zip(keys, values)).items():
+                    new.coords[key] = base_coord, list(value)
         return new
 
     def curvefit(

xarray/tests/test_dataarray.py

@@ -7371,26 +7371,26 @@ def test_drop_duplicate_coords(keep):
     da.coords["valid"] = (("init", "tau"), np.array([[8, 6, 6], [7, 7, 7]]))
 
     if keep == "first":
-        data = [[1, 2], [4, np.nan]]
-        init = [0, 1]
-        tau = [1, 2]
-        valid = [[8.0, 6.0], [7.0, np.nan]]
+        data = [1, 2, 4]
+        init = [0, 0, 1]
+        tau = [1, 2, 1]
+        valid = [8, 6, 7]
     elif keep == "last":
-        data = [[1, 3], [np.nan, 6]]
-        init = [0, 1]
-        tau = [1, 3]
-        valid = [[8.0, 6.0], [np.nan, 7]]
+        data = [1, 3, 6]
+        init = [0, 0, 1]
+        tau = [1, 3, 3]
+        valid = [8, 6, 7]
     else:
-        data = [[1]]
+        data = [1]
         init = [0]
         tau = [1]
-        valid = [[8]]
+        valid = [8]
 
     result = da.drop_duplicate_coords("valid", keep=keep)
     expected = xr.DataArray(
         data,
-        dims=["init", "tau"],
-        coords={"init": init, "tau": tau, "valid": (("init", "tau"), valid)},
+        dims="valid",
+        coords={"init": ("valid", init), "tau": ("valid", tau), "valid": valid},
     )
     assert_equal(expected, result)
 
@@ -7404,9 +7404,29 @@ def test_drop_duplicate_coords_duplicate_dims(keep):
     )
     da.coords["valid"] = (("init", "tau"), np.array([[8, 6, 6], [7, 7, 7]]))
 
-    with pytest.raises(ValueError):
-        da.drop_duplicate_coords("valid", keep=keep)
+    if keep == "first":
+        data = [1, 2, 4]
+        init = [0, 0, 0]
+        tau = [1, 2, 1]
+        valid = [8, 6, 7]
+    elif keep == "last":
+        data = [1, 3, 6]
+        init = [0, 0, 0]
+        tau = [1, 3, 3]
+        valid = [8, 6, 7]
+    else:
+        data = [1]
+        init = [0]
+        tau = [1]
+        valid = [8]
 
+    result = da.drop_duplicate_coords("valid", keep=keep)
+    expected = xr.DataArray(
+        data,
+        dims="valid",
+        coords={"init": ("valid", init), "tau": ("valid", tau), "valid": valid},
+    )
+    assert_equal(expected, result)
 
 @pytest.mark.parametrize("keep", ["first", "last", False])
 def test_drop_duplicate_coords_missing(keep):

xarray/tests/test_dataset.py

@@ -6910,26 +6910,26 @@ def test_drop_duplicate_coords(keep):
     ds.coords["valid"] = (("init", "tau"), np.array([[8, 6, 6], [7, 7, 7]]))
 
     if keep == "first":
-        data = [[1, 2], [4, np.nan]]
-        init = [0, 1]
-        tau = [1, 2]
-        valid = [[8.0, 6.0], [7.0, np.nan]]
+        data = [1, 2, 4]
+        init = [0, 0, 1]
+        tau = [1, 2, 1]
+        valid = [8, 6, 7]
     elif keep == "last":
-        data = [[1, 3], [np.nan, 6]]
-        init = [0, 1]
-        tau = [1, 3]
-        valid = [[8.0, 6.0], [np.nan, 7]]
+        data = [1, 3, 6]
+        init = [0, 0, 1]
+        tau = [1, 3, 3]
+        valid = [8, 6, 7]
     else:
-        data = [[1]]
+        data = [1]
         init = [0]
         tau = [1]
-        valid = [[8]]
+        valid = [8]
 
     result = ds.drop_duplicate_coords("valid", keep=keep)
     expected = xr.DataArray(
         data,
-        dims=["init", "tau"],
-        coords={"init": init, "tau": tau, "valid": (("init", "tau"), valid)},
+        dims="valid",
+        coords={"init": ("valid", init), "tau": ("valid", tau), "valid": valid},
         name="test",
     ).to_dataset()
     assert_equal(expected, result)
@@ -6938,14 +6938,36 @@ def test_drop_duplicate_coords(keep):
 @pytest.mark.parametrize("keep", ["first", "last", False])
 def test_drop_duplicate_coords_duplicate_dims(keep):
     ds = xr.DataArray(
-        [["a", "b", "c"], ["d", "e", "f"]],
-        coords={"init": [0, 0], "tau": [1, 2, 3]},  # duplicate inits
+        [[1, 2, 3], [4, 5, 6]],
+        coords={"init": [0, 0], "tau": [1, 2, 3]},
         dims=["init", "tau"],
     ).to_dataset(name="test")
     ds.coords["valid"] = (("init", "tau"), np.array([[8, 6, 6], [7, 7, 7]]))
 
-    with pytest.raises(ValueError):
-        ds.drop_duplicate_coords("valid", keep=keep)
+    if keep == "first":
+        data = [1, 2, 4]
+        init = [0, 0, 0]
+        tau = [1, 2, 1]
+        valid = [8, 6, 7]
+    elif keep == "last":
+        data = [1, 3, 6]
+        init = [0, 0, 0]
+        tau = [1, 3, 3]
+        valid = [8, 6, 7]
+    else:
+        data = [1]
+        init = [0]
+        tau = [1]
+        valid = [8]
+
+    result = ds.drop_duplicate_coords("valid", keep=keep)
+    expected = xr.DataArray(
+        data,
+        dims="valid",
+        coords={"init": ("valid", init), "tau": ("valid", tau), "valid": valid},
+        name="test",
+    ).to_dataset()
+    assert_equal(expected, result)
 
 
 @pytest.mark.parametrize("keep", ["first", "last", False])