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

…ex (pandas-dev#16889)
hhuuggoo · Jul 13, 2017 · 36a7b63 · 36a7b63
1 parent 9d13227
commit 36a7b63
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 2 deletions.
diff --git a/doc/source/whatsnew/v0.21.0.txt b/doc/source/whatsnew/v0.21.0.txt
@@ -144,6 +144,7 @@ Bug Fixes
 ~~~~~~~~~
 
 - Fixes regression in 0.20, :func:`Series.aggregate` and :func:`DataFrame.aggregate` allow dictionaries as return values again (:issue:`16741`)
+- Fixes `__setitem__` with a tuple for dataframes with a tz-aware `DatetimeIndex` (:issue: `16889`)
 
 Conversion
 ^^^^^^^^^^

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 setitem with datetimeindex with tz
+        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