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

Author: jorisvandenbossche

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,13 @@ def _get_data(axis_matters):
                 raise NotImplementedError(msg)
             return data
 
+        if axis == 0 and column_wise:
+            # column-wise reduction
+            df = self
+            if numeric_only is True:
+                df = _get_data(axis_matters=True)
+            return DataFrame._reduce_columns(df, op, name, skipna=skipna, **kwds)
+
         if numeric_only is not None and axis in [0, 1]:
             df = self
             if numeric_only is True:
@@ -7994,6 +8018,47 @@ def blk_func(values):
             result = self._constructor_sliced(result, index=labels)
         return result
 
+    def _reduce_columns(self, op, name, skipna=True, **kwds):
+        """
+        Reduce DataFrame column-wise.
+
+        Parameters
+        ----------
+        op : func
+            The reducing function to be called on the values. Only used
+            for columns backed by a numpy ndarray.
+        name : str
+            The name of the reduction.
+        skipna, **kwds : keywords to pass to the `op` function
+
+        Returns
+        -------
+        Series
+        """
+        result = []
+
+        for arr in self._iter_arrays():
+            if isinstance(arr, ExtensionArray):
+                # dispatch to ExtensionArray interface
+                val = arr._reduce(name, skipna=skipna, **kwds)
+            else:
+                # dispatch to numpy arrays
+                with np.errstate(all="ignore"):
+                    val = op(arr, skipna=skipna, **kwds)
+
+            result.append(val)
+
+        return self._constructor_sliced(result, index=self.columns)
+
+    def _iter_arrays(self):
+        """
+        Iterate over the arrays of all columns in order.
+
+        This returns the values as stored in the Block (ndarray or ExtensionArray).
+        """
+        for i in range(len(self.columns)):
+            yield self._data.iget_values(i)
+
     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

@@ -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: