REF: simplify Block.diff (#48348)

* REF: simplify Block.diff

* fix ArrayManager.diff

Author: jbrockmendel

pandas/core/frame.py

@@ -9243,8 +9243,14 @@ def diff(self, periods: int = 1, axis: Axis = 0) -> DataFrame:
             periods = int(periods)
 
         axis = self._get_axis_number(axis)
-        if axis == 1 and periods != 0:
-            return self - self.shift(periods, axis=axis)
+        if axis == 1:
+            if periods != 0:
+                # in the periods == 0 case, this is equivalent diff of 0 periods
+                #  along axis=0, and the Manager method may be somewhat more
+                #  performant, so we dispatch in that case.
+                return self - self.shift(periods, axis=axis)
+            # With periods=0 this is equivalent to a diff with axis=0
+            axis = 0
 
         new_data = self._mgr.diff(n=periods, axis=axis)
         return self._constructor(new_data).__finalize__(self, "diff")

pandas/core/internals/array_manager.py

@@ -363,11 +363,7 @@ def putmask(self: T, mask, new, align: bool = True) -> T:
         )
 
     def diff(self: T, n: int, axis: int) -> T:
-        if axis == 1:
-            # DataFrame only calls this for n=0, in which case performing it
-            # with axis=0 is equivalent
-            assert n == 0
-            axis = 0
+        assert self.ndim == 2 and axis == 0  # caller ensures
         return self.apply(algos.diff, n=n, axis=axis)
 
     def interpolate(self: T, **kwargs) -> T:

pandas/core/internals/blocks.py

@@ -1271,6 +1271,7 @@ def interpolate(
 
     def diff(self, n: int, axis: int = 1) -> list[Block]:
         """return block for the diff of the values"""
+        # only reached with ndim == 2 and axis == 1
         new_values = algos.diff(self.values, n, axis=axis)
         return [self.make_block(values=new_values)]
 
@@ -1830,17 +1831,10 @@ def getitem_block_index(self, slicer: slice) -> ExtensionBlock:
         return type(self)(new_values, self._mgr_locs, ndim=self.ndim)
 
     def diff(self, n: int, axis: int = 1) -> list[Block]:
-        if axis == 0 and n != 0:
-            # n==0 case will be a no-op so let is fall through
-            # Since we only have one column, the result will be all-NA.
-            #  Create this result by shifting along axis=0 past the length of
-            #  our values.
-            return super().diff(len(self.values), axis=0)
-        if axis == 1:
-            # TODO(EA2D): unnecessary with 2D EAs
-            # we are by definition 1D.
-            axis = 0
-        return super().diff(n, axis)
+        # only reached with ndim == 2 and axis == 1
+        # TODO(EA2D): Can share with NDArrayBackedExtensionBlock
+        new_values = algos.diff(self.values, n, axis=0)
+        return [self.make_block(values=new_values)]
 
     def shift(self, periods: int, axis: int = 0, fill_value: Any = None) -> list[Block]:
         """
@@ -1964,6 +1958,7 @@ def diff(self, n: int, axis: int = 0) -> list[Block]:
         The arguments here are mimicking shift so they are called correctly
         by apply.
         """
+        # only reached with ndim == 2 and axis == 1
         values = self.values
 
         new_values = values - values.shift(n, axis=axis)

pandas/core/internals/managers.py

@@ -413,6 +413,7 @@ def putmask(self, mask, new, align: bool = True):
         )
 
     def diff(self: T, n: int, axis: int) -> T:
+        # only reached with self.ndim == 2 and axis == 1
         axis = self._normalize_axis(axis)
         return self.apply("diff", n=n, axis=axis)