Fix assignment to Datasets with non-unique indexes (#945)

* Fix assignment to Datasets with non-unique indexes

Fixes GH943

* Fix failing test

Author: shoyer

doc/whats-new.rst

@@ -13,6 +13,17 @@ What's New
     import xarray as xr
     np.random.seed(123456)
 
+.. _whats-new.0.8.1:
+
+v0.8.1 (unreleased)
+-------------------
+
+Bug fixes
+~~~~~~~~~
+
+- Fix regression in v0.8.0 that broke assignment to Datasets with non-unique
+  indexes (:issue:`943`).
+
 .. _whats-new.0.8.0:
 
 v0.8.0 (2 August 2016)

xarray/core/alignment.py

@@ -96,10 +96,23 @@ def partial_align(*objects, **kwargs):
     exclude = kwargs.pop('exclude', None)
     if exclude is None:
         exclude = set()
+    skip_single_target = kwargs.pop('skip_single_target', False)
     if kwargs:
         raise TypeError('align() got unexpected keyword arguments: %s'
                         % list(kwargs))
 
+    if len(objects) == 1:
+        obj, = objects
+        if (indexes is None or
+                (skip_single_target and
+                 all(obj.indexes[k].equals(v) for k, v in indexes.items()
+                     if k in obj.indexes))):
+            # We don't need to align, so don't bother with reindexing, which
+            # fails for non-unique indexes.
+            # `skip_single_target` is a hack so we can skip alignment of a
+            # single object in merge.
+            return (obj.copy() if copy else obj,)
+
     joined_indexes = _join_indexes(join, objects, exclude=exclude)
     if indexes is not None:
         joined_indexes.update(indexes)
@@ -109,14 +122,16 @@ def partial_align(*objects, **kwargs):
         valid_indexers = dict((k, v) for k, v in joined_indexes.items()
                               if k in obj.dims)
         result.append(obj.reindex(copy=copy, **valid_indexers))
+
     return tuple(result)
 
 
 def is_alignable(obj):
     return hasattr(obj, 'indexes') and hasattr(obj, 'reindex')
 
 
-def deep_align(list_of_variable_maps, join='outer', copy=True, indexes=None):
+def deep_align(list_of_variable_maps, join='outer', copy=True, indexes=None,
+               skip_single_target=False):
     """Align objects, recursing into dictionary values.
     """
     if indexes is None:
@@ -145,7 +160,8 @@ def deep_align(list_of_variable_maps, join='outer', copy=True, indexes=None):
         else:
             out.append(variables)
 
-    aligned = partial_align(*targets, join=join, copy=copy, indexes=indexes)
+    aligned = partial_align(*targets, join=join, copy=copy, indexes=indexes,
+                            skip_single_target=skip_single_target)
 
     for key, aligned_obj in zip(keys, aligned):
         if isinstance(key, tuple):

xarray/core/merge.py

@@ -315,7 +315,8 @@ def align_and_merge_coords(objs, compat='minimal', join='outer',
     """Align and merge coordinate variables."""
     _assert_compat_valid(compat)
     coerced = coerce_pandas_values(objs)
-    aligned = deep_align(coerced, join=join, copy=False, indexes=indexes)
+    aligned = deep_align(coerced, join=join, copy=False, indexes=indexes,
+                         skip_single_target=True)
     expanded = expand_variable_dicts(aligned)
     priority_vars = _get_priority_vars(aligned, priority_arg, compat=compat)
     variables = merge_variables(expanded, priority_vars, compat=compat)
@@ -370,7 +371,8 @@ def merge_core(objs, compat='broadcast_equals', join='outer', priority_arg=None,
     _assert_compat_valid(compat)
 
     coerced = coerce_pandas_values(objs)
-    aligned = deep_align(coerced, join=join, copy=False, indexes=indexes)
+    aligned = deep_align(coerced, join=join, copy=False, indexes=indexes,
+                         skip_single_target=True)
     expanded = expand_variable_dicts(aligned)
 
     coord_names, noncoord_names = determine_coords(coerced)

xarray/test/test_dataset.py

@@ -1453,6 +1453,13 @@ def test_setitem_auto_align(self):
         expected = Dataset({'x': ('y', [4, 5, 6])})
         self.assertDatasetIdentical(ds, expected)
 
+    def test_setitem_align_new_indexes(self):
+        ds = Dataset({'foo': ('x', [1, 2, 3])}, {'x': [0, 1, 2]})
+        ds['bar'] = DataArray([2, 3, 4], [('x', [1, 2, 3])])
+        expected = Dataset({'foo': ('x', [1, 2, 3]),
+                            'bar': ('x', [np.nan, 2, 3])})
+        self.assertDatasetIdentical(ds, expected)
+
     def test_assign(self):
         ds = Dataset()
         actual = ds.assign(x = [0, 1, 2], y = 2)
@@ -1482,6 +1489,24 @@ def test_assign(self):
         expected = expected.set_coords('c')
         self.assertDatasetIdentical(actual, expected)
 
+    def test_setitem_non_unique_index(self):
+        # regression test for GH943
+        original = Dataset({'data': ('x', np.arange(5))},
+                            coords={'x': [0, 1, 2, 0, 1]})
+        expected = Dataset({'data': ('x', np.arange(5))})
+
+        actual = original.copy()
+        actual['x'] = list(range(5))
+        self.assertDatasetIdentical(actual, expected)
+
+        actual = original.copy()
+        actual['x'] = ('x', list(range(5)))
+        self.assertDatasetIdentical(actual, expected)
+
+        actual = original.copy()
+        actual.coords['x'] = list(range(5))
+        self.assertDatasetIdentical(actual, expected)
+
     def test_delitem(self):
         data = create_test_data()
         all_items = set(data)