BUG: __setitem__ with a tuple induces NaN with a tz-aware DatetimeInd…

…ex (pandas-dev#16889)
hhuuggoo · Jul 20, 2017 · fc20f32 · fc20f32
1 parent 8d0c025
commit fc20f32
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 3 deletions.
diff --git a/doc/source/whatsnew/v0.21.0.txt b/doc/source/whatsnew/v0.21.0.txt
@@ -196,7 +196,8 @@ Indexing
 - Fix :func:`MultiIndex.sort_index` ordering when ``ascending`` argument is a list, but not all levels are specified, or are in a different order (:issue:`16934`).
 - Fixes bug where indexing with ``np.inf`` caused an ``OverflowError`` to be raised (:issue:`16957`)
 - Bug in reindexing on an empty ``CategoricalIndex`` (:issue:`16770`)
-
+- Fixes ``DataFrame.loc`` for setting with alignment and tz-aware ``DatetimeIndex`` (:issue:`16889`)
+
 I/O
 ^^^
 

diff --git a/pandas/core/indexing.py b/pandas/core/indexing.py
@@ -760,10 +760,12 @@ def _align_frame(self, indexer, df):
             for i, ix in enumerate(indexer):
                 ax = self.obj.axes[i]
                 if is_sequence(ix) or isinstance(ix, slice):
+                    if isinstance(ix, np.ndarray):
+                        ix = ix.ravel()
                     if idx is None:
-                        idx = ax[ix].ravel()
+                        idx = ax[ix]
                     elif cols is None:
-                        cols = ax[ix].ravel()
+                        cols = ax[ix]
                     else:
                         break
                 else:

diff --git a/pandas/tests/indexing/test_datetime.py b/pandas/tests/indexing/test_datetime.py
@@ -8,6 +8,33 @@
 
 class TestDatetimeIndex(object):
 
+    def test_setitem_with_datetime_tz(self):
+        # 16889
+        # support .loc with alignment and tz-aware DatetimeIndex
+        mask = np.array([True, False, True, False])
+
+        idx = pd.date_range('20010101', periods=4, tz='UTC')
+        df = pd.DataFrame({'a': np.arange(4)}, index=idx).astype('float64')
+
+        result = df.copy()
+        result.loc[mask, :] = df.loc[mask, :]
+        tm.assert_frame_equal(result, df)
+
+        result = df.copy()
+        result.loc[mask] = df.loc[mask]
+        tm.assert_frame_equal(result, df)
+
+        idx = pd.date_range('20010101', periods=4)
+        df = pd.DataFrame({'a': np.arange(4)}, index=idx).astype('float64')
+
+        result = df.copy()
+        result.loc[mask, :] = df.loc[mask, :]
+        tm.assert_frame_equal(result, df)
+
+        result = df.copy()
+        result.loc[mask] = df.loc[mask]
+        tm.assert_frame_equal(result, df)
+
     def test_indexing_with_datetime_tz(self):
 
         # 8260