ENH: 2D support for MaskedArray

pandas/core/array_algos/masked_reductions.py

@@ -3,7 +3,7 @@
 for missing values.
 """
 
-from typing import Callable
+from typing import Callable, Optional
 
 import numpy as np
 
@@ -20,6 +20,7 @@ def _sumprod(
     *,
     skipna: bool = True,
     min_count: int = 0,
+    axis: Optional[int] = None,
 ):
     """
     Sum or product for 1D masked array.
@@ -37,40 +38,58 @@ def _sumprod(
     min_count : int, default 0
         The required number of valid values to perform the operation. If fewer than
         ``min_count`` non-NA values are present the result will be NA.
+    axis : int, optional, default None
     """
     if not skipna:
-        if mask.any() or check_below_min_count(values.shape, None, min_count):
+        if mask.any(axis=axis) or check_below_min_count(values.shape, None, min_count):
             return libmissing.NA
         else:
-            return func(values)
+            return func(values, axis=axis)
     else:
-        if check_below_min_count(values.shape, mask, min_count):
+        if check_below_min_count(values.shape, mask, min_count) and (
+            axis is None or values.ndim == 1
+        ):
             return libmissing.NA
 
         if np_version_under1p17:
-            return func(values[~mask])
+            return func(values[~mask], axis=axis)
         else:
-            return func(values, where=~mask)
+            return func(values, where=~mask, axis=axis)
 
 
 def sum(
-    values: np.ndarray, mask: np.ndarray, *, skipna: bool = True, min_count: int = 0
+    values: np.ndarray,
+    mask: np.ndarray,
+    *,
+    skipna: bool = True,
+    min_count: int = 0,
+    axis: Optional[int] = None,
 ):
     return _sumprod(
-        np.sum, values=values, mask=mask, skipna=skipna, min_count=min_count
+        np.sum, values=values, mask=mask, skipna=skipna, min_count=min_count, axis=axis
     )
 
 
 def prod(
-    values: np.ndarray, mask: np.ndarray, *, skipna: bool = True, min_count: int = 0
+    values: np.ndarray,
+    mask: np.ndarray,
+    *,
+    skipna: bool = True,
+    min_count: int = 0,
+    axis: Optional[int] = None,
 ):
     return _sumprod(
-        np.prod, values=values, mask=mask, skipna=skipna, min_count=min_count
+        np.prod, values=values, mask=mask, skipna=skipna, min_count=min_count, axis=axis
     )
 
 
 def _minmax(
-    func: Callable, values: np.ndarray, mask: np.ndarray, *, skipna: bool = True
+    func: Callable,
+    values: np.ndarray,
+    mask: np.ndarray,
+    *,
+    skipna: bool = True,
+    axis: Optional[int] = None,
 ):
     """
     Reduction for 1D masked array.
@@ -85,6 +104,7 @@ def _minmax(
         Boolean numpy array (True values indicate missing values).
     skipna : bool, default True
         Whether to skip NA.
+    axis : int, optional, default None
     """
     if not skipna:
         if mask.any() or not values.size:
@@ -101,14 +121,27 @@ def _minmax(
             return libmissing.NA
 
 
-def min(values: np.ndarray, mask: np.ndarray, *, skipna: bool = True):
-    return _minmax(np.min, values=values, mask=mask, skipna=skipna)
+def min(
+    values: np.ndarray,
+    mask: np.ndarray,
+    *,
+    skipna: bool = True,
+    axis: Optional[int] = None,
+):
+    return _minmax(np.min, values=values, mask=mask, skipna=skipna, axis=axis)
 
 
-def max(values: np.ndarray, mask: np.ndarray, *, skipna: bool = True):
-    return _minmax(np.max, values=values, mask=mask, skipna=skipna)
+def max(
+    values: np.ndarray,
+    mask: np.ndarray,
+    *,
+    skipna: bool = True,
+    axis: Optional[int] = None,
+):
+    return _minmax(np.max, values=values, mask=mask, skipna=skipna, axis=axis)
 
 
+# TODO: axis kwarg
 def mean(values: np.ndarray, mask: np.ndarray, skipna: bool = True):
     if not values.size or mask.all():
         return libmissing.NA

pandas/core/arrays/_mixins.py

@@ -301,27 +301,6 @@ def _wrap_reduction_result(self, axis: Optional[int], result):
             return self._box_func(result)
         return self._from_backing_data(result)
 
-    # ------------------------------------------------------------------------
-
-    def __repr__(self) -> str:
-        if self.ndim == 1:
-            return super().__repr__()
-
-        from pandas.io.formats.printing import format_object_summary
-
-        # the short repr has no trailing newline, while the truncated
-        # repr does. So we include a newline in our template, and strip
-        # any trailing newlines from format_object_summary
-        lines = [
-            format_object_summary(x, self._formatter(), indent_for_name=False).rstrip(
-                ", \n"
-            )
-            for x in self
-        ]
-        data = ",\n".join(lines)
-        class_name = f"<{type(self).__name__}>"
-        return f"{class_name}\n[\n{data}\n]\nShape: {self.shape}, dtype: {self.dtype}"
-
     # ------------------------------------------------------------------------
     # __array_function__ methods
 

pandas/core/arrays/base.py

@@ -1097,6 +1097,9 @@ def view(self, dtype: Optional[Dtype] = None) -> ArrayLike:
     # ------------------------------------------------------------------------
 
     def __repr__(self) -> str:
+        if self.ndim > 1:
+            return self._repr_2d()
+
         from pandas.io.formats.printing import format_object_summary
 
         # the short repr has no trailing newline, while the truncated
@@ -1108,6 +1111,22 @@ def __repr__(self) -> str:
         class_name = f"<{type(self).__name__}>\n"
         return f"{class_name}{data}\nLength: {len(self)}, dtype: {self.dtype}"
 
+    def _repr_2d(self) -> str:
+        from pandas.io.formats.printing import format_object_summary
+
+        # the short repr has no trailing newline, while the truncated
+        # repr does. So we include a newline in our template, and strip
+        # any trailing newlines from format_object_summary
+        lines = [
+            format_object_summary(x, self._formatter(), indent_for_name=False).rstrip(
+                ", \n"
+            )
+            for x in self
+        ]
+        data = ",\n".join(lines)
+        class_name = f"<{type(self).__name__}>"
+        return f"{class_name}\n[\n{data}\n]\nShape: {self.shape}, dtype: {self.dtype}"
+
     def _formatter(self, boxed: bool = False) -> Callable[[Any], Optional[str]]:
         """
         Formatting function for scalar values.

pandas/core/arrays/boolean.py

@@ -193,10 +193,8 @@ def coerce_to_array(
             if mask_values is not None:
                 mask = mask | mask_values
 
-    if values.ndim != 1:
-        raise ValueError("values must be a 1D list-like")
-    if mask.ndim != 1:
-        raise ValueError("mask must be a 1D list-like")
+    if values.shape != mask.shape:
+        raise ValueError("values.shape and mask.shape must match")
 
     return values, mask
 
@@ -411,9 +409,9 @@ def _values_for_argsort(self) -> np.ndarray:
         """
         data = self._data.copy()
         data[self._mask] = -1
-        return data
+        return data.ravel("K")
 
-    def any(self, *, skipna: bool = True, **kwargs):
+    def any(self, *, skipna: bool = True, axis: Optional[int] = 0, **kwargs):
         """
         Return whether any element is True.
 
@@ -430,6 +428,7 @@ def any(self, *, skipna: bool = True, **kwargs):
             If `skipna` is False, the result will still be True if there is
             at least one element that is True, otherwise NA will be returned
             if there are NA's present.
+        axis : int or None, default 0
         **kwargs : any, default None
             Additional keywords have no effect but might be accepted for
             compatibility with NumPy.
@@ -472,16 +471,17 @@ def any(self, *, skipna: bool = True, **kwargs):
 
         values = self._data.copy()
         np.putmask(values, self._mask, False)
-        result = values.any()
+        result = values.any(axis=axis)
+
         if skipna:
             return result
         else:
-            if result or len(self) == 0 or not self._mask.any():
+            if result or self.size == 0 or not self._mask.any():
                 return result
             else:
                 return self.dtype.na_value
 
-    def all(self, *, skipna: bool = True, **kwargs):
+    def all(self, *, skipna: bool = True, axis: Optional[int] = 0, **kwargs):
         """
         Return whether all elements are True.
 
@@ -498,6 +498,7 @@ def all(self, *, skipna: bool = True, **kwargs):
             If `skipna` is False, the result will still be False if there is
             at least one element that is False, otherwise NA will be returned
             if there are NA's present.
+        axis : int or None, default 0
         **kwargs : any, default None
             Additional keywords have no effect but might be accepted for
             compatibility with NumPy.
@@ -538,12 +539,12 @@ def all(self, *, skipna: bool = True, **kwargs):
 
         values = self._data.copy()
         np.putmask(values, self._mask, True)
-        result = values.all()
+        result = values.all(axis=axis)
 
         if skipna:
             return result
         else:
-            if not result or len(self) == 0 or not self._mask.any():
+            if not result or self.size == 0 or not self._mask.any():
                 return result
             else:
                 return self.dtype.na_value

pandas/core/arrays/floating.py

@@ -394,21 +394,21 @@ def _cmp_method(self, other, op):
 
         return BooleanArray(result, mask)
 
-    def sum(self, *, skipna=True, min_count=0, **kwargs):
+    def sum(self, *, skipna=True, min_count=0, axis: Optional[int] = 0, **kwargs):
         nv.validate_sum((), kwargs)
-        return super()._reduce("sum", skipna=skipna, min_count=min_count)
+        return super()._reduce("sum", skipna=skipna, min_count=min_count, axis=axis)
 
-    def prod(self, *, skipna=True, min_count=0, **kwargs):
+    def prod(self, *, skipna=True, min_count=0, axis: Optional[int] = 0, **kwargs):
         nv.validate_prod((), kwargs)
-        return super()._reduce("prod", skipna=skipna, min_count=min_count)
+        return super()._reduce("prod", skipna=skipna, min_count=min_count, axis=axis)
 
-    def min(self, *, skipna=True, **kwargs):
+    def min(self, *, skipna=True, axis: Optional[int] = 0, **kwargs):
         nv.validate_min((), kwargs)
-        return super()._reduce("min", skipna=skipna)
+        return super()._reduce("min", skipna=skipna, axis=axis)
 
-    def max(self, *, skipna=True, **kwargs):
+    def max(self, *, skipna=True, axis: Optional[int] = 0, **kwargs):
         nv.validate_max((), kwargs)
-        return super()._reduce("max", skipna=skipna)
+        return super()._reduce("max", skipna=skipna, axis=axis)
 
     def _maybe_mask_result(self, result, mask, other, op_name: str):
         """

pandas/core/arrays/integer.py

@@ -423,7 +423,7 @@ def _values_for_argsort(self) -> np.ndarray:
         data = self._data.copy()
         if self._mask.any():
             data[self._mask] = data.min() - 1
-        return data
+        return data.ravel("K")
 
     def _cmp_method(self, other, op):
         from pandas.core.arrays import BooleanArray
@@ -470,21 +470,21 @@ def _cmp_method(self, other, op):
 
         return BooleanArray(result, mask)
 
-    def sum(self, *, skipna=True, min_count=0, **kwargs):
+    def sum(self, *, skipna=True, min_count=0, axis: Optional[int] = 0, **kwargs):
         nv.validate_sum((), kwargs)
-        return super()._reduce("sum", skipna=skipna, min_count=min_count)
+        return super()._reduce("sum", skipna=skipna, min_count=min_count, axis=axis)
 
-    def prod(self, *, skipna=True, min_count=0, **kwargs):
+    def prod(self, *, skipna=True, min_count=0, axis: Optional[int] = 0, **kwargs):
         nv.validate_prod((), kwargs)
-        return super()._reduce("prod", skipna=skipna, min_count=min_count)
+        return super()._reduce("prod", skipna=skipna, min_count=min_count, axis=axis)
 
-    def min(self, *, skipna=True, **kwargs):
+    def min(self, *, skipna=True, axis: Optional[int] = 0, **kwargs):
         nv.validate_min((), kwargs)
-        return super()._reduce("min", skipna=skipna)
+        return super()._reduce("min", skipna=skipna, axis=axis)
 
-    def max(self, *, skipna=True, **kwargs):
+    def max(self, *, skipna=True, axis: Optional[int] = 0, **kwargs):
         nv.validate_max((), kwargs)
-        return super()._reduce("max", skipna=skipna)
+        return super()._reduce("max", skipna=skipna, axis=axis)
 
     def _maybe_mask_result(self, result, mask, other, op_name: str):
         """