REF: move logic of 'block manager axis' into the BlockManager (#40075)

Author: jorisvandenbossche

pandas/core/frame.py

@@ -7843,12 +7843,11 @@ def diff(self, periods: int = 1, axis: Axis = 0) -> DataFrame:
                 raise ValueError("periods must be an integer")
             periods = int(periods)
 
-        bm_axis = self._get_block_manager_axis(axis)
-
-        if bm_axis == 0 and periods != 0:
+        axis = self._get_axis_number(axis)
+        if axis == 1 and periods != 0:
             return self - self.shift(periods, axis=axis)
 
-        new_data = self._mgr.diff(n=periods, axis=bm_axis)
+        new_data = self._mgr.diff(n=periods, axis=axis)
         return self._constructor(new_data).__finalize__(self, "diff")
 
     # ----------------------------------------------------------------------

pandas/core/generic.py

@@ -8968,8 +8968,6 @@ def _where(
                 self._info_axis, axis=self._info_axis_number, copy=False
             )
 
-        block_axis = self._get_block_manager_axis(axis)
-
         if inplace:
             # we may have different type blocks come out of putmask, so
             # reconstruct the block manager
@@ -8985,7 +8983,7 @@ def _where(
                 cond=cond,
                 align=align,
                 errors=errors,
-                axis=block_axis,
+                axis=axis,
             )
             result = self._constructor(new_data)
             return result.__finalize__(self)
@@ -9296,9 +9294,9 @@ def shift(
 
         if freq is None:
             # when freq is None, data is shifted, index is not
-            block_axis = self._get_block_manager_axis(axis)
+            axis = self._get_axis_number(axis)
             new_data = self._mgr.shift(
-                periods=periods, axis=block_axis, fill_value=fill_value
+                periods=periods, axis=axis, fill_value=fill_value
             )
             return self._constructor(new_data).__finalize__(self, method="shift")
 

pandas/core/internals/array_manager.py

@@ -406,12 +406,7 @@ def apply(
 
         return type(self)(result_arrays, new_axes)
 
-    def apply_2d(
-        self: T,
-        f,
-        ignore_failures: bool = False,
-        **kwargs,
-    ) -> T:
+    def apply_2d(self: T, f, ignore_failures: bool = False, **kwargs) -> T:
         """
         Variant of `apply`, but where the function should not be applied to
         each column independently, but to the full data as a 2D array.
@@ -430,7 +425,10 @@ def apply_2d(
 
         return type(self)(result_arrays, new_axes)
 
-    def apply_with_block(self: T, f, align_keys=None, **kwargs) -> T:
+    def apply_with_block(self: T, f, align_keys=None, swap_axis=True, **kwargs) -> T:
+        # switch axis to follow BlockManager logic
+        if swap_axis and "axis" in kwargs and self.ndim == 2:
+            kwargs["axis"] = 1 if kwargs["axis"] == 0 else 0
 
         align_keys = align_keys or []
         aligned_args = {k: kwargs[k] for k in align_keys}
@@ -542,7 +540,6 @@ def putmask(self, mask, new, align: bool = True):
         )
 
     def diff(self, n: int, axis: int) -> ArrayManager:
-        axis = self._normalize_axis(axis)
         if axis == 1:
             # DataFrame only calls this for n=0, in which case performing it
             # with axis=0 is equivalent
@@ -551,13 +548,13 @@ def diff(self, n: int, axis: int) -> ArrayManager:
         return self.apply(algos.diff, n=n, axis=axis)
 
     def interpolate(self, **kwargs) -> ArrayManager:
-        return self.apply_with_block("interpolate", **kwargs)
+        return self.apply_with_block("interpolate", swap_axis=False, **kwargs)
 
     def shift(self, periods: int, axis: int, fill_value) -> ArrayManager:
         if fill_value is lib.no_default:
             fill_value = None
 
-        if axis == 0 and self.ndim == 2:
+        if axis == 1 and self.ndim == 2:
             # TODO column-wise shift
             raise NotImplementedError
 

pandas/core/internals/managers.py

@@ -237,6 +237,12 @@ def shape(self) -> Shape:
     def ndim(self) -> int:
         return len(self.axes)
 
+    def _normalize_axis(self, axis):
+        # switch axis to follow BlockManager logic
+        if self.ndim == 2:
+            axis = 1 if axis == 0 else 0
+        return axis
+
     def set_axis(
         self, axis: int, new_labels: Index, verify_integrity: bool = True
     ) -> None:
@@ -560,6 +566,7 @@ def isna(self, func) -> BlockManager:
         return self.apply("apply", func=func)
 
     def where(self, other, cond, align: bool, errors: str, axis: int) -> BlockManager:
+        axis = self._normalize_axis(axis)
         if align:
             align_keys = ["other", "cond"]
         else:
@@ -594,12 +601,14 @@ def putmask(self, mask, new, align: bool = True):
         )
 
     def diff(self, n: int, axis: int) -> BlockManager:
+        axis = self._normalize_axis(axis)
         return self.apply("diff", n=n, axis=axis)
 
     def interpolate(self, **kwargs) -> BlockManager:
         return self.apply("interpolate", **kwargs)
 
     def shift(self, periods: int, axis: int, fill_value) -> BlockManager:
+        axis = self._normalize_axis(axis)
         if fill_value is lib.no_default:
             fill_value = None