REF: Implement BaseMaskedArray class for integer/boolean ExtensionArrays (#30789)

Author: jorisvandenbossche

pandas/core/arrays/boolean.py

@@ -15,24 +15,19 @@
  is_extension_array_dtype,
  is_float,
  is_float_dtype,
- is_integer,
  is_integer_dtype,
  is_list_like,
  is_numeric_dtype,
- is_object_dtype,
  is_scalar,
- is_string_dtype,
  pandas_dtype,
 )
 from pandas.core.dtypes.dtypes import register_extension_dtype
 from pandas.core.dtypes.generic import ABCDataFrame, ABCIndexClass, ABCSeries
 from pandas.core.dtypes.missing import isna, notna
 
 from pandas.core import nanops, ops
-from pandas.core.algorithms import take
-from pandas.core.arrays import ExtensionArray, ExtensionOpsMixin
-import pandas.core.common as com
-from pandas.core.indexers import check_bool_array_indexer
+
+from .masked import BaseMaskedArray
 
 if TYPE_CHECKING:
  from pandas._typing import Scalar
@@ -199,7 +194,7 @@ def coerce_to_array(values, mask=None, copy: bool = False):
  return values, mask
 
 
-class BooleanArray(ExtensionArray, ExtensionOpsMixin):
+class BooleanArray(BaseMaskedArray):
  """
  Array of boolean (True/False) data with missing values.
 
@@ -253,6 +248,9 @@ class BooleanArray(ExtensionArray, ExtensionOpsMixin):
  Length: 3, dtype: boolean
  """
 
+ # The value used to fill '_data' to avoid upcasting
+ _internal_fill_value = False
+
  def __init__(self, values: np.ndarray, mask: np.ndarray, copy: bool = False):
  if not (isinstance(values, np.ndarray) and values.dtype == np.bool_):
  raise TypeError(
@@ -297,127 +295,6 @@ def _values_for_factorize(self) -> Tuple[np.ndarray, Any]:
  def _from_factorized(cls, values, original: "BooleanArray"):
  return cls._from_sequence(values, dtype=original.dtype)
 
- def _formatter(self, boxed=False):
- return str
-
- @property
- def _hasna(self) -> bool:
- # Note: this is expensive right now! The hope is that we can
- # make this faster by having an optional mask, but not have to change
- # source code using it..
- return self._mask.any()
-
- def __getitem__(self, item):
- if is_integer(item):
- if self._mask[item]:
- return self.dtype.na_value
- return self._data[item]
-
- elif com.is_bool_indexer(item):
- item = check_bool_array_indexer(self, item)
-
- return type(self)(self._data[item], self._mask[item])
-
- def to_numpy(
- self, dtype=None, copy=False, na_value: "Scalar" = lib.no_default,
- ):
- """
- Convert to a NumPy Array.
-
- By default converts to an object-dtype NumPy array. Specify the `dtype` and
- `na_value` keywords to customize the conversion.
-
- Parameters
- ----------
- dtype : dtype, default object
- The numpy dtype to convert to.
- copy : bool, default False
- Whether to ensure that the returned value is a not a view on
- the array. Note that ``copy=False`` does not *ensure* that
- ``to_numpy()`` is no-copy. Rather, ``copy=True`` ensure that
- a copy is made, even if not strictly necessary. This is typically
- only possible when no missing values are present and `dtype`
- is a boolean dtype.
- na_value : scalar, optional
- Scalar missing value indicator to use in numpy array. Defaults
- to the native missing value indicator of this array (pd.NA).
-
- Returns
- -------
- numpy.ndarray
-
- Examples
- --------
- An object-dtype is the default result
-
- >>> a = pd.array([True, False], dtype="boolean")
- >>> a.to_numpy()
- array([True, False], dtype=object)
-
- When no missing values are present, a boolean dtype can be used.
-
- >>> a.to_numpy(dtype="bool")
- array([ True, False])
-
- However, requesting a bool dtype will raise a ValueError if
- missing values are present and the default missing value :attr:`NA`
- is used.
-
- >>> a = pd.array([True, False, pd.NA], dtype="boolean")
- >>> a
- <BooleanArray>
- [True, False, NA]
- Length: 3, dtype: boolean
-
- >>> a.to_numpy(dtype="bool")
- Traceback (most recent call last):
- ...
- ValueError: cannot convert to bool numpy array in presence of missing values
-
- Specify a valid `na_value` instead
-
- >>> a.to_numpy(dtype="bool", na_value=False)
- array([ True, False, False])
- """
- if na_value is lib.no_default:
- na_value = libmissing.NA
- if dtype is None:
- dtype = object
- if self._hasna:
- if (
- not (is_object_dtype(dtype) or is_string_dtype(dtype))
- and na_value is libmissing.NA
- ):
- raise ValueError(
- f"cannot convert to '{dtype}'-dtype NumPy array "
- "with missing values. Specify an appropriate 'na_value' "
- "for this dtype."
- )
- # don't pass copy to astype -> always need a copy since we are mutating
- data = self._data.astype(dtype)
- data[self._mask] = na_value
- else:
- data = self._data.astype(dtype, copy=copy)
- return data
-
- __array_priority__ = 1000 # higher than ndarray so ops dispatch to us
-
- def __array__(self, dtype=None):
- """
- the array interface, return my values
- We return an object array here to preserve our scalar values
- """
- # by default (no dtype specified), return an object array
- return self.to_numpy(dtype=dtype)
-
- def __arrow_array__(self, type=None):
- """
- Convert myself into a pyarrow Array.
- """
- import pyarrow as pa
-
- return pa.array(self._data, mask=self._mask, type=type)
-
  _HANDLED_TYPES = (np.ndarray, numbers.Number, bool, np.bool_)
 
  def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
@@ -465,40 +342,6 @@ def reconstruct(x):
  else:
  return reconstruct(result)
 
- def __iter__(self):
- for i in range(len(self)):
- if self._mask[i]:
- yield self.dtype.na_value
- else:
- yield self._data[i]
-
- def take(self, indexer, allow_fill=False, fill_value=None):
- # we always fill with False internally
- # to avoid upcasting
- data_fill_value = False if isna(fill_value) else fill_value
- result = take(
- self._data, indexer, fill_value=data_fill_value, allow_fill=allow_fill
- )
-
- mask = take(self._mask, indexer, fill_value=True, allow_fill=allow_fill)
-
- # if we are filling
- # we only fill where the indexer is null
- # not existing missing values
- # TODO(jreback) what if we have a non-na float as a fill value?
- if allow_fill and notna(fill_value):
- fill_mask = np.asarray(indexer) == -1
- result[fill_mask] = fill_value
- mask = mask ^ fill_mask
-
- return type(self)(result, mask, copy=False)
-
- def copy(self):
- data, mask = self._data, self._mask
- data = data.copy()
- mask = mask.copy()
- return type(self)(data, mask, copy=False)
-
  def __setitem__(self, key, value):
  _is_scalar = is_scalar(value)
  if _is_scalar:
@@ -512,26 +355,6 @@ def __setitem__(self, key, value):
  self._data[key] = value
  self._mask[key] = mask
 
- def __len__(self):
- return len(self._data)
-
- @property
- def nbytes(self):
- return self._data.nbytes + self._mask.nbytes
-
- def isna(self):
- return self._mask
-
- @property
- def _na_value(self):
- return self._dtype.na_value
-
- @classmethod
- def _concat_same_type(cls, to_concat):
- data = np.concatenate([x._data for x in to_concat])
- mask = np.concatenate([x._mask for x in to_concat])
- return cls(data, mask)
-
  def astype(self, dtype, copy=True):
  """
  Cast to a NumPy array or ExtensionArray with 'dtype'.