rework tests & implementation in concat of tz-aware dataframes

Author: jreback

pandas/core/dtypes/concat.py

@@ -416,6 +416,13 @@ def _maybe_unwrap(x):
                        fastpath=True)
 
 
+def _concatenate_2d(to_concat, axis):
+    # coerce to 2d if needed & concatenate
+    if axis == 1:
+        to_concat = [np.atleast_2d(x) for x in to_concat]
+    return np.concatenate(to_concat, axis=axis)
+
+
 def _concat_datetime(to_concat, axis=0, typs=None):
     """
     provide concatenation of an datetimelike array of arrays each of which is a
@@ -432,61 +439,57 @@ def _concat_datetime(to_concat, axis=0, typs=None):
     a single array, preserving the combined dtypes
     """
 
-    def convert_to_pydatetime(x, axis):
-        # coerce to an object dtype
+    if typs is None:
+        typs = get_dtype_kinds(to_concat)
 
-        # if dtype is of datetimetz or timezone
-        if x.dtype.kind == _NS_DTYPE.kind:
-            if getattr(x, 'tz', None) is not None:
-                x = x.astype(object).values
-            else:
-                shape = x.shape
-                x = tslib.ints_to_pydatetime(x.view(np.int64).ravel(),
-                                             box="timestamp")
-                x = x.reshape(shape)
+    # multiple types, need to coerce to object
+    if len(typs) != 1:
+        return _concatenate_2d([_convert_datetimelike_to_object(x)
+                                for x in to_concat],
+                               axis=axis)
 
-        elif x.dtype == _TD_DTYPE:
-            shape = x.shape
-            x = tslib.ints_to_pytimedelta(x.view(np.int64).ravel(), box=True)
-            x = x.reshape(shape)
+    # must be single dtype
+    if any(typ.startswith('datetime') for typ in typs):
 
-        if axis == 1:
-            x = np.atleast_2d(x)
-        return x
+        if 'datetime' in typs:
+            to_concat = [np.array(x, copy=False).view(np.int64)
+                         for x in to_concat]
+            return _concatenate_2d(to_concat, axis=axis).view(_NS_DTYPE)
+        else:
+            # when to_concat has different tz, len(typs) > 1.
+            # thus no need to care
+            return _concat_datetimetz(to_concat)
 
-    if typs is None:
-        typs = get_dtype_kinds(to_concat)
+    elif 'timedelta' in typs:
+        return _concatenate_2d([x.view(np.int64) for x in to_concat],
+                               axis=axis).view(_TD_DTYPE)
 
-    # must be single dtype
-    if len(typs) == 1:
-        _contains_datetime = any(typ.startswith('datetime') for typ in typs)
-        _contains_period = any(typ.startswith('period') for typ in typs)
+    elif any(typ.startswith('period') for typ in typs):
+        # PeriodIndex must be handled by PeriodIndex,
+        # Thus can't meet this condition ATM
+        # Must be changed when we adding PeriodDtype
+        raise NotImplementedError("unable to concat PeriodDtype")
 
-        if _contains_datetime:
 
-            if 'datetime' in typs:
-                new_values = np.concatenate([x.view(np.int64) for x in
-                                             to_concat], axis=axis)
-                return new_values.view(_NS_DTYPE)
-            else:
-                # when to_concat has different tz, len(typs) > 1.
-                # thus no need to care
-                return _concat_datetimetz(to_concat)
-
-        elif 'timedelta' in typs:
-            new_values = np.concatenate([x.view(np.int64) for x in to_concat],
-                                        axis=axis)
-            return new_values.view(_TD_DTYPE)
-
-        elif _contains_period:
-            # PeriodIndex must be handled by PeriodIndex,
-            # Thus can't meet this condition ATM
-            # Must be changed when we adding PeriodDtype
-            raise NotImplementedError
-
-    # need to coerce to object
-    to_concat = [convert_to_pydatetime(x, axis) for x in to_concat]
-    return np.concatenate(to_concat, axis=axis)
+def _convert_datetimelike_to_object(x):
+    # coerce datetimelike array to object dtype
+
+    # if dtype is of datetimetz or timezone
+    if x.dtype.kind == _NS_DTYPE.kind:
+        if getattr(x, 'tz', None) is not None:
+            x = x.astype(object).values
+        else:
+            shape = x.shape
+            x = tslib.ints_to_pydatetime(x.view(np.int64).ravel(),
+                                         box="timestamp")
+            x = x.reshape(shape)
+
+    elif x.dtype == _TD_DTYPE:
+        shape = x.shape
+        x = tslib.ints_to_pytimedelta(x.view(np.int64).ravel(), box=True)
+        x = x.reshape(shape)
+
+    return x
 
 
 def _concat_datetimetz(to_concat, name=None):

pandas/core/indexes/base.py

@@ -2183,19 +2183,17 @@ def _assert_take_fillable(self, values, indices, allow_fill=True,
                               fill_value=None, na_value=np.nan):
         """ Internal method to handle NA filling of take """
         indices = _ensure_platform_int(indices)
+
         # only fill if we are passing a non-None fill_value
         if allow_fill and fill_value is not None:
             if (indices < -1).any():
                 msg = ('When allow_fill=True and fill_value is not None, '
                        'all indices must be >= -1')
                 raise ValueError(msg)
-            mask = indices == -1
-            if mask.all():
-                taken = np.full(indices.shape, fill_value=na_value)
-            else:
-                taken = values.take(indices)
-                if mask.any():
-                    taken[mask] = na_value
+            taken = algos.take(values,
+                               indices,
+                               allow_fill=allow_fill,
+                               fill_value=na_value)
         else:
             taken = values.take(indices)
         return taken

pandas/core/internals.py

@@ -5837,8 +5837,7 @@ def get_reindexed_values(self, empty_dtype, upcasted_na):
 
                 if getattr(self.block, 'is_datetimetz', False) or \
                         is_datetimetz(empty_dtype):
-                    missing_arr = np.full(np.prod(self.shape), fill_value)
-                    return DatetimeIndex(missing_arr, dtype=empty_dtype)
+                    pass
                 elif getattr(self.block, 'is_categorical', False):
                     pass
                 elif getattr(self.block, 'is_sparse', False):

pandas/tests/reshape/test_concat.py

@@ -1917,131 +1917,73 @@ def test_concat_tz_series_tzlocal(self):
         tm.assert_series_equal(result, pd.Series(x + y))
         assert result.dtype == 'datetime64[ns, tzlocal()]'
 
-    def test_concat_NaT_dataframes_all_NaT_axis_0(self):
+    @pytest.mark.parametrize('tz1', [None, 'UTC'])
+    @pytest.mark.parametrize('tz2', [None, 'UTC'])
+    @pytest.mark.parametrize('s', [pd.NaT, pd.Timestamp('20150101')])
+    def test_concat_NaT_dataframes_all_NaT_axis_0(self, tz1, tz2, s):
         # GH 12396
 
         # tz-naive
-        first = pd.DataFrame([[pd.NaT], [pd.NaT]])
-        second = pd.DataFrame([[pd.NaT]])
+        first = pd.DataFrame([[pd.NaT], [pd.NaT]]).apply(
+            lambda x: x.dt.tz_localize(tz1))
+        second = pd.DataFrame([s]).apply(lambda x: x.dt.tz_localize(tz2))
 
         result = pd.concat([first, second], axis=0)
-        expected = pd.DataFrame([pd.NaT, pd.NaT, pd.NaT], index=[0, 1, 0])
-        assert_frame_equal(result, expected)
+        expected = pd.DataFrame(pd.Series(
+            [pd.NaT, pd.NaT, s], index=[0, 1, 0]))
+        expected = expected.apply(lambda x: x.dt.tz_localize(tz2))
+        if tz1 != tz2:
+            expected = expected.astype(object)
 
-        # one side timezone-aware
-        # upcasts for mixed case
-        first = pd.DataFrame(pd.Series([pd.NaT, pd.NaT]).dt.tz_localize('UTC'))
-        result = pd.concat([first, second], axis=0)
-        expected = pd.DataFrame(
-            pd.Series([pd.NaT, pd.NaT, pd.NaT]).dt.tz_localize('UTC'),
-            index=[0, 1, 0]
-        )
         assert_frame_equal(result, expected)
 
-        # both sides timezone-aware
-        # upcasts to tz-aware
-        second = pd.DataFrame(pd.Series([pd.NaT]).dt.tz_localize('UTC'))
-        result = pd.concat([first, second], axis=0)
-        assert_frame_equal(result, expected)
-
-    def test_concat_NaT_dataframes_all_NaT_axis_1(self):
+    @pytest.mark.parametrize('tz1', [None, 'UTC'])
+    @pytest.mark.parametrize('tz2', [None, 'UTC'])
+    def test_concat_NaT_dataframes_all_NaT_axis_1(self, tz1, tz2):
         # GH 12396
 
-        # tz-naive
-        first = pd.DataFrame([[pd.NaT], [pd.NaT]])
-        second = pd.DataFrame([[pd.NaT]], columns=[1])
-        expected = pd.DataFrame([[pd.NaT, pd.NaT], [pd.NaT, pd.NaT]],
-                                columns=[0, 1])
-        result = pd.concat([first, second], axis=1)
-        assert_frame_equal(result, expected)
-
-        # one side timezone-aware
-        # upcasts result to tz-aware
-        first = pd.DataFrame(pd.Series([pd.NaT, pd.NaT]).dt.tz_localize('UTC'))
+        first = pd.DataFrame(pd.Series([pd.NaT, pd.NaT]).dt.tz_localize(tz1))
+        second = pd.DataFrame(pd.Series(
+            [pd.NaT]).dt.tz_localize(tz2), columns=[1])
         expected = pd.DataFrame(
-            {0: pd.Series([pd.NaT, pd.NaT]).dt.tz_localize('UTC'),
-             1: pd.Series([pd.NaT, pd.NaT])}
+            {0: pd.Series([pd.NaT, pd.NaT]).dt.tz_localize(tz1),
+             1: pd.Series([pd.NaT, pd.NaT]).dt.tz_localize(tz2)}
         )
         result = pd.concat([first, second], axis=1)
         assert_frame_equal(result, expected)
 
-        # both sides timezone-aware
-        # upcasts result to tz-aware
-        second[1] = second[1].dt.tz_localize('UTC')
-        expected = pd.DataFrame(
-            {0: pd.Series([pd.NaT, pd.NaT]).dt.tz_localize('UTC'),
-             1: pd.Series([pd.NaT, pd.NaT]).dt.tz_localize('UTC')}
-        )
-        result = pd.concat([first, second], axis=1)
-        assert_frame_equal(result, expected)
-
-    def test_concat_NaT_dataframes_mixed_timestamps_and_NaT(self):
+    @pytest.mark.parametrize('tz1', [None, 'UTC'])
+    @pytest.mark.parametrize('tz2', [None, 'UTC'])
+    def test_concat_NaT_series_dataframe_all_NaT(self, tz1, tz2):
         # GH 12396
 
         # tz-naive
-        first = pd.DataFrame([[pd.NaT], [pd.NaT]])
-        second = pd.DataFrame([[pd.Timestamp('2015/01/01')],
-                               [pd.Timestamp('2016/01/01')]],
-                              index=[2, 3])
-        expected = pd.DataFrame([pd.NaT, pd.NaT,
-                                 pd.Timestamp('2015/01/01'),
-                                 pd.Timestamp('2016/01/01')])
-
-        result = pd.concat([first, second], axis=0)
-        assert_frame_equal(result, expected)
-
-        # one side timezone-aware
-        second = second[0].dt.tz_localize('UTC')
-        expected = pd.DataFrame(
-            pd.Series([pd.NaT, pd.NaT,
-                       pd.Timestamp('2015/01/01'),
-                       pd.Timestamp('2016/01/01')]).dt.tz_localize('UTC')
-        )
-        result = pd.concat([first, second], axis=0)
-        assert_frame_equal(result, expected)
-
-    def test_concat_NaT_series_dataframe_all_NaT(self):
-        # GH 12396
-
-        # tz-naive
-        first = pd.Series([pd.NaT, pd.NaT])
-        second = pd.DataFrame([[pd.Timestamp('2015/01/01')],
-                               [pd.Timestamp('2016/01/01')]],
+        first = pd.Series([pd.NaT, pd.NaT]).dt.tz_localize(tz1)
+        second = pd.DataFrame([[pd.Timestamp('2015/01/01', tz=tz2)],
+                               [pd.Timestamp('2016/01/01', tz=tz2)]],
                               index=[2, 3])
 
         expected = pd.DataFrame([pd.NaT, pd.NaT,
-                                 pd.Timestamp('2015/01/01'),
-                                 pd.Timestamp('2016/01/01')])
+                                 pd.Timestamp('2015/01/01', tz=tz2),
+                                 pd.Timestamp('2016/01/01', tz=tz2)])
+        if tz1 != tz2:
+            expected = expected.astype(object)
 
         result = pd.concat([first, second])
         assert_frame_equal(result, expected)
 
-        # one side timezone-aware
-        second[0] = second[0].dt.tz_localize('UTC')
-        result = pd.concat([first, second])
-
-        expected = pd.DataFrame(
-            pd.Series([pd.NaT, pd.NaT,
-                       pd.Timestamp('2015/01/01'),
-                       pd.Timestamp('2016/01/01')]).dt.tz_localize('UTC')
-        )
-        assert_frame_equal(result, expected)
-
-        # both sides timezone-aware
-        first = first.dt.tz_localize('UTC')
-        result = pd.concat([first, second])
-        assert_frame_equal(result, expected)
+    @pytest.mark.parametrize('tz', [None, 'UTC'])
+    def test_concat_NaT_dataframes(self, tz):
+        # GH 12396
 
-        # mixed tz
         first = pd.DataFrame([[pd.NaT], [pd.NaT]])
-        second = pd.DataFrame([[pd.Timestamp('2015/01/01', tz='UTC')],
-                               [pd.Timestamp('2016/01/01', tz='US/Eastern')]],
+        first = first.apply(lambda x: x.dt.tz_localize(tz))
+        second = pd.DataFrame([[pd.Timestamp('2015/01/01', tz=tz)],
+                               [pd.Timestamp('2016/01/01', tz=tz)]],
                               index=[2, 3])
-
-        expected = pd.DataFrame([pd.NaT,
-                                 pd.NaT,
-                                 pd.Timestamp('2015/01/01', tz='UTC'),
-                                 pd.Timestamp('2016/01/01', tz='US/Eastern')])
+        expected = pd.DataFrame([pd.NaT, pd.NaT,
+                                 pd.Timestamp('2015/01/01', tz=tz),
+                                 pd.Timestamp('2016/01/01', tz=tz)])
 
         result = pd.concat([first, second], axis=0)
         assert_frame_equal(result, expected)
@@ -2107,32 +2049,21 @@ def test_concat_empty_series(self):
                            columns=['x', 0])
         tm.assert_frame_equal(res, exp)
 
+    @pytest.mark.parametrize('tz', [None, 'UTC'])
+    @pytest.mark.parametrize('values', [[], [1, 2, 3]])
+    def test_concat_empty_series_timelike(self, tz, values):
         # GH 18447
-        # tz-naive
-        first = Series(pd.to_datetime([], utc=False))
-        second = Series([1, 2, 3])
-        expected = DataFrame([[pd.NaT, 1], [pd.NaT, 2], [pd.NaT, 3]])
-        result = concat([first, second], axis=1)
-        assert_frame_equal(result, expected)
 
-        # timezone-aware
-        first = Series(pd.to_datetime([], utc=True))
-        second = Series([1, 2, 3])
+        first = Series([], dtype='M8[ns]').dt.tz_localize(tz)
+        second = Series(values)
         expected = DataFrame(
-            {0: pd.Series([pd.NaT, pd.NaT, pd.NaT]).dt.tz_localize('UTC'),
-             1: pd.Series([1, 2, 3])}
-        )
+            {0: pd.Series([pd.NaT] * len(values),
+                          dtype='M8[ns]'
+                          ).dt.tz_localize(tz),
+             1: values})
         result = concat([first, second], axis=1)
         assert_frame_equal(result, expected)
 
-        # both empty
-        first = Series(pd.to_datetime([], utc=True))
-        second = Series([])
-        result = concat([first, second], axis=1)
-        assert result.size == 0
-        assert result.dtypes[0] == 'datetime64[ns, UTC]'
-        assert result.dtypes[1] == 'float64'
-
     def test_default_index(self):
         # is_series and ignore_index
         s1 = pd.Series([1, 2, 3], name='x')