ENH/PERF: enable column-wise reductions for EA-backed columns

pandas/core/frame.py

@@ -7852,6 +7852,23 @@ def _count_level(self, level, axis=0, numeric_only=False):
     def _reduce(
         self, op, name, axis=0, skipna=True, numeric_only=None, filter_type=None, **kwds
     ):
+        """
+        Reduce DataFrame over axis with given operation.
+
+        Parameters
+        ----------
+        op : func
+            The reducing function to be called on the values.
+        name : str
+            The name of the reduction.
+        axis : int
+        numeric_only : bool, optional
+        filter_type : None or "bool"
+            Set to "bool" for ops that give boolean results.
+        skipna, **kwds : keywords to pass to the `op` function
+
+        """
+        column_wise = kwds.pop("column_wise", False)
 
         assert filter_type is None or filter_type == "bool", filter_type
 
@@ -7898,6 +7915,19 @@ def _get_data(axis_matters):
                 raise NotImplementedError(msg)
             return data
 
+        def blk_func(values):
+            if isinstance(values, ExtensionArray):
+                return values._reduce(name, skipna=skipna, **kwds)
+            else:
+                return op(values, axis=1, skipna=skipna, **kwds)
+
+        if axis == 0 and column_wise:
+            # column-wise reduction
+            df = self
+            if numeric_only is True:
+                df = _get_data(axis_matters=True)
+            return df._reduce_columns(blk_func)
+
         if numeric_only is not None and axis in [0, 1]:
             df = self
             if numeric_only is True:
@@ -7908,12 +7938,6 @@ def _get_data(axis_matters):
 
             out_dtype = "bool" if filter_type == "bool" else None
 
-            def blk_func(values):
-                if values.ndim == 1 and not isinstance(values, np.ndarray):
-                    # we can't pass axis=1
-                    return op(values, axis=0, skipna=skipna, **kwds)
-                return op(values, axis=1, skipna=skipna, **kwds)
-
             # After possibly _get_data and transposing, we are now in the
             #  simple case where we can use BlockManager._reduce
             res = df._data.reduce(blk_func)
@@ -7994,6 +8018,27 @@ def blk_func(values):
             result = self._constructor_sliced(result, index=labels)
         return result
 
+    def _reduce_columns(self, op):
+        """
+        Reduce DataFrame column-wise.
+
+        Parameters
+        ----------
+        op : func
+            The reducing function to be called on the values.
+
+        Returns
+        -------
+        Series
+        """
+        result = []
+
+        for i in range(len(self.columns)):
+            val = op(self._data.iget_values(i))
+            result.append(val)
+
+        return self._constructor_sliced(result, index=self.columns)
+
     def nunique(self, axis=0, dropna=True) -> Series:
         """
         Count distinct observations over requested axis.

pandas/core/generic.py

@@ -11067,6 +11067,7 @@ def stat_func(
         min_count=0,
         **kwargs,
     ):
+        column_wise = kwargs.pop("column_wise", False)
         if name == "sum":
             nv.validate_sum(tuple(), kwargs)
         elif name == "prod":
@@ -11088,6 +11089,7 @@ def stat_func(
             skipna=skipna,
             numeric_only=numeric_only,
             min_count=min_count,
+            column_wise=column_wise,
         )
 
     return set_function_name(stat_func, name, cls)
@@ -11117,6 +11119,7 @@ def _make_stat_function(
     def stat_func(
         self, axis=None, skipna=None, level=None, numeric_only=None, **kwargs
     ):
+        column_wise = kwargs.pop("column_wise", False)
         if name == "median":
             nv.validate_median(tuple(), kwargs)
         else:
@@ -11128,7 +11131,12 @@ def stat_func(
         if level is not None:
             return self._agg_by_level(name, axis=axis, level=level, skipna=skipna)
         return self._reduce(
-            func, name=name, axis=axis, skipna=skipna, numeric_only=numeric_only
+            func,
+            name=name,
+            axis=axis,
+            skipna=skipna,
+            numeric_only=numeric_only,
+            column_wise=column_wise,
         )
 
     return set_function_name(stat_func, name, cls)

pandas/core/internals/managers.py

@@ -349,24 +349,24 @@ def _verify_integrity(self) -> None:
                 f"tot_items: {tot_items}"
             )
 
-    def reduce(self, func, *args, **kwargs):
+    def reduce(self, func):
         # If 2D, we assume that we're operating column-wise
         if self.ndim == 1:
             # we'll be returning a scalar
             blk = self.blocks[0]
-            return func(blk.values, *args, **kwargs)
+            return func(blk.values)
 
         res = {}
         for blk in self.blocks:
-            bres = func(blk.values, *args, **kwargs)
+            bres = func(blk.values)
 
             if np.ndim(bres) == 0:
                 # EA
                 assert blk.shape[0] == 1
                 new_res = zip(blk.mgr_locs.as_array, [bres])
             else:
                 assert bres.ndim == 1, bres.shape
-                assert blk.shape[0] == len(bres), (blk.shape, bres.shape, args, kwargs)
+                assert blk.shape[0] == len(bres), (blk.shape, bres.shape)
                 new_res = zip(blk.mgr_locs.as_array, bres)
 
             nr = dict(new_res)
@@ -998,6 +998,14 @@ def iget(self, i: int) -> "SingleBlockManager":
             fastpath=True,
         )
 
+    def iget_values(self, i: int):
+        """
+        Return the data for column i as the values (ndarray or ExtensionArray).
+        """
+        block = self.blocks[self.blknos[i]]
+        values = block.iget(self.blklocs[i])
+        return values
+
     def delete(self, item):
         """
         Delete selected item (items if non-unique) in-place.

pandas/core/series.py

@@ -3871,6 +3871,7 @@ def _reduce(
         If we have an ndarray as a value, then simply perform the operation,
         otherwise delegate to the object.
         """
+        kwds.pop("column_wise", None)
         delegate = self._values
 
         if axis is not None: