ENH: add masked algorithm for mean() function (#34814)

Author: Akshatt

doc/source/whatsnew/v1.3.0.rst

@@ -167,7 +167,7 @@ Deprecations
 Performance improvements
 ~~~~~~~~~~~~~~~~~~~~~~~~
 - Performance improvement in :meth:`IntervalIndex.isin` (:issue:`38353`)
--
+- Performance improvement in :meth:`Series.mean` for nullable data types (:issue:`34814`)
 -
 
 .. ---------------------------------------------------------------------------

pandas/core/array_algos/masked_reductions.py

@@ -107,3 +107,12 @@ def min(values: np.ndarray, mask: np.ndarray, *, skipna: bool = True):
 
 def max(values: np.ndarray, mask: np.ndarray, *, skipna: bool = True):
     return _minmax(np.max, values=values, mask=mask, skipna=skipna)
+
+
+def mean(values: np.ndarray, mask: np.ndarray, skipna: bool = True):
+    if not values.size or mask.all():
+        return libmissing.NA
+    _sum = _sumprod(np.sum, values=values, mask=mask, skipna=skipna)
+    count = np.count_nonzero(~mask)
+    mean_value = _sum / count
+    return mean_value

pandas/core/arrays/masked.py

@@ -394,7 +394,7 @@ def _reduce(self, name: str, *, skipna: bool = True, **kwargs):
         data = self._data
         mask = self._mask
 
-        if name in {"sum", "prod", "min", "max"}:
+        if name in {"sum", "prod", "min", "max", "mean"}:
             op = getattr(masked_reductions, name)
             return op(data, mask, skipna=skipna, **kwargs)
 

pandas/tests/reductions/test_reductions.py

@@ -688,6 +688,23 @@ def test_empty_multi(self, method, unit):
         expected = Series([1, np.nan], index=["a", "b"])
         tm.assert_series_equal(result, expected)
 
+    @pytest.mark.parametrize("method", ["mean"])
+    @pytest.mark.parametrize("dtype", ["Float64", "Int64", "boolean"])
+    def test_ops_consistency_on_empty_nullable(self, method, dtype):
+
+        # GH#34814
+        # consistency for nullable dtypes on empty or ALL-NA mean
+
+        # empty series
+        eser = Series([], dtype=dtype)
+        result = getattr(eser, method)()
+        assert result is pd.NA
+
+        # ALL-NA series
+        nser = Series([np.nan], dtype=dtype)
+        result = getattr(nser, method)()
+        assert result is pd.NA
+
     @pytest.mark.parametrize("method", ["mean", "median", "std", "var"])
     def test_ops_consistency_on_empty(self, method):