move diff into its own function in DatetimeTZBlock

Author: mroeschke

pandas/core/algorithms.py

@@ -1501,12 +1501,7 @@ def diff(arr, n, axis=0):
     is_timedelta = False
     if needs_i8_conversion(arr):
         dtype = np.float64
-        if is_datetime64tz_dtype(arr):
-            # Block data is usually a 2D array
-            # except DatetimeTZBlock which is a 1D array (DatetimeIndex)
-            arr = arr.view('i8').reshape(1, arr.shape[0])
-        else:
-            arr = arr.view('i8')
+        arr = arr.view('i8')
         na = iNaT
         is_timedelta = True
 

pandas/core/internals.py

@@ -1269,11 +1269,6 @@ def take_nd(self, indexer, axis, new_mgr_locs=None, fill_tuple=None):
 
     def diff(self, n, axis=1, mgr=None):
         """ return block for the diff of the values """
-        if isinstance(self, DatetimeTZBlock) and axis == 0:
-            # This method will iterate (self.apply) over each DateTimeTZBlock
-            # but axis=0 will need to gather data from all DateTimeTZBlocks
-            # in the manager in order for diff to work correctly.
-            raise NotImplementedError
         new_values = algos.diff(self.values, n, axis=axis)
         return [self.make_block(values=new_values)]
 
@@ -2857,7 +2852,8 @@ def shift(self, periods, axis=0, mgr=None):
 
         # think about moving this to the DatetimeIndex. This is a non-freq
         # (number of periods) shift ###
-
+        if axis == 0:
+            raise NotImplementedError
         N = len(self)
         indexer = np.zeros(N, dtype=int)
         if periods > 0:
@@ -2876,6 +2872,11 @@ def shift(self, periods, axis=0, mgr=None):
         return [self.make_block_same_class(new_values,
                                            placement=self.mgr_locs)]
 
+    def diff(self, n, axis=0, mgr=None):
+        """1st discrete difference"""
+        new_values = self.values - self.shift(n, axis=axis)[0].values
+        return [TimeDeltaBlock(new_values, placement=self.mgr_locs)]
+
     def concat_same_type(self, to_concat, placement=None):
         """
         Concatenate list of single blocks of the same type.

pandas/tests/frame/test_timeseries.py

@@ -59,19 +59,26 @@ def test_diff(self):
 
     @pytest.mark.parametrize('axis', [0, 1])
     @pytest.mark.parametrize('tz', [None, 'UTC'])
-    @pytest.mark.xfail(raises=NotImplementedError)
     def test_diff_datetime(self, axis, tz):
         # GH 18578
         df = DataFrame({0: date_range('2010', freq='D', periods=2, tz=tz),
                         1: date_range('2010', freq='D', periods=2, tz=tz)})
-        result = df.diff(axis=axis)  # xfails for axis=1 and tz='UTC'
         if axis == 1:
-            expected = DataFrame({0: pd.TimedeltaIndex(['NaT', 'NaT']),
-                                  1: pd.TimedeltaIndex(['0 days', '0 days'])})
+            if tz is None:
+                result = df.diff(axis=axis)
+                expected = DataFrame({0: pd.TimedeltaIndex(['NaT', 'NaT']),
+                                      1: pd.TimedeltaIndex(['0 days',
+                                                            '0 days'])})
+                assert_frame_equal(result, expected)
+            else:
+                with pytest.raises(NotImplementedError):
+                    result = df.diff(axis=axis)
+
         else:
+            result = df.diff(axis=axis)
             expected = DataFrame({0: pd.TimedeltaIndex(['NaT', '1 days']),
                                   1: pd.TimedeltaIndex(['NaT', '1 days'])})
-        assert_frame_equal(result, expected)
+            assert_frame_equal(result, expected)
 
     def test_diff_timedelta(self):
         # GH 4533
@@ -294,6 +301,28 @@ def test_shift(self):
         result = df.shift(1, axis='columns')
         assert_frame_equal(result, expected)
 
+    @pytest.mark.parametrize('axis', [0, 1])
+    @pytest.mark.parametrize('tz', [None, 'UTC'])
+    def test_shift_datetime(self, axis, tz):
+        # GH 18578
+        df = DataFrame(date_range('2010', freq='D', periods=2, tz=tz))
+        df[1] = df[0]
+        if axis == 1:
+            if tz is None:
+                result = df.shift(axis=axis)
+                expected = DataFrame(DatetimeIndex([pd.NaT] * 2))
+                expected[1] = df[1]
+                assert_frame_equal(result, expected)
+            else:
+                with pytest.raises(NotImplementedError):
+                    result = df.shift(axis=axis)
+
+        else:
+            result = df.shift(axis=axis)
+            expected = DataFrame([pd.NaT, Timestamp('2010', tz=tz)])
+            expected[1] = expected[0]
+            assert_frame_equal(result, expected)
+
     def test_shift_bool(self):
         df = DataFrame({'high': [True, False],
                         'low': [False, False]})