diff --git a/doc/source/getting_started/basics.rst b/doc/source/getting_started/basics.rst
index 125990f7cadcd..6301fee7775cf 100644
--- a/doc/source/getting_started/basics.rst
+++ b/doc/source/getting_started/basics.rst
@@ -1950,6 +1950,7 @@ sparse :class:`SparseDtype` (none) :class:`arrays.
intervals :class:`IntervalDtype` :class:`Interval` :class:`arrays.IntervalArray` :ref:`advanced.intervalindex`
nullable integer :class:`Int64Dtype`, ... (none) :class:`arrays.IntegerArray` :ref:`integer_na`
Strings :class:`StringDtype` :class:`str` :class:`arrays.StringArray` :ref:`text`
+Boolean (with NA) :class:`BooleanDtype` :class:`bool` :class:`arrays.BooleanArray` :ref:`api.arrays.bool`
=================== ========================= ================== ============================= =============================
Pandas has two ways to store strings.
diff --git a/doc/source/reference/arrays.rst b/doc/source/reference/arrays.rst
index 0c435e06ac57f..cf14d28772f4c 100644
--- a/doc/source/reference/arrays.rst
+++ b/doc/source/reference/arrays.rst
@@ -25,6 +25,7 @@ Nullable Integer :class:`Int64Dtype`, ... (none) :ref:`api.array
Categorical :class:`CategoricalDtype` (none) :ref:`api.arrays.categorical`
Sparse :class:`SparseDtype` (none) :ref:`api.arrays.sparse`
Strings :class:`StringDtype` :class:`str` :ref:`api.arrays.string`
+Boolean (with NA) :class:`BooleanDtype` :class:`bool` :ref:`api.arrays.bool`
=================== ========================= ================== =============================
Pandas and third-party libraries can extend NumPy's type system (see :ref:`extending.extension-types`).
@@ -485,6 +486,28 @@ The ``Series.str`` accessor is available for ``Series`` backed by a :class:`arra
See :ref:`api.series.str` for more.
+.. _api.arrays.bool:
+
+Boolean data with missing values
+--------------------------------
+
+The boolean dtype (with the alias ``"boolean"``) provides support for storing
+boolean data (True, False values) with missing values, which is not possible
+with a bool :class:`numpy.ndarray`.
+
+.. autosummary::
+ :toctree: api/
+ :template: autosummary/class_without_autosummary.rst
+
+ arrays.BooleanArray
+
+.. autosummary::
+ :toctree: api/
+ :template: autosummary/class_without_autosummary.rst
+
+ BooleanDtype
+
+
.. Dtype attributes which are manually listed in their docstrings: including
.. it here to make sure a docstring page is built for them
diff --git a/doc/source/whatsnew/v1.0.0.rst b/doc/source/whatsnew/v1.0.0.rst
index 3990eec2435d9..7d11d90eeb670 100644
--- a/doc/source/whatsnew/v1.0.0.rst
+++ b/doc/source/whatsnew/v1.0.0.rst
@@ -102,6 +102,30 @@ String accessor methods returning integers will return a value with :class:`Int6
We recommend explicitly using the ``string`` data type when working with strings.
See :ref:`text.types` for more.
+.. _whatsnew_100.boolean:
+
+Boolean data type with missing values support
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+We've added :class:`BooleanDtype` / :class:`~arrays.BooleanArray`, an extension
+type dedicated to boolean data that can hold missing values. With the default
+``'bool`` data type based on a numpy bool array, the column can only hold
+True or False values and not missing values. This new :class:`BooleanDtype`
+can store missing values as well by keeping track of this in a separate mask.
+(:issue:`29555`)
+
+.. ipython:: python
+
+ pd.Series([True, False, None], dtype=pd.BooleanDtype())
+
+You can use the alias ``"boolean"`` as well.
+
+.. ipython:: python
+
+ s = pd.Series([True, False, None], dtype="boolean")
+ s
+
+
.. _whatsnew_1000.enhancements.other:
Other enhancements
diff --git a/pandas/__init__.py b/pandas/__init__.py
index 5d163e411c0ac..cd697b757a26a 100644
--- a/pandas/__init__.py
+++ b/pandas/__init__.py
@@ -67,6 +67,7 @@
IntervalDtype,
DatetimeTZDtype,
StringDtype,
+ BooleanDtype,
# missing
isna,
isnull,
diff --git a/pandas/arrays/__init__.py b/pandas/arrays/__init__.py
index 9870b5bed076d..61832a8b6d621 100644
--- a/pandas/arrays/__init__.py
+++ b/pandas/arrays/__init__.py
@@ -4,6 +4,7 @@
See :ref:`extending.extension-types` for more.
"""
from pandas.core.arrays import (
+ BooleanArray,
Categorical,
DatetimeArray,
IntegerArray,
@@ -16,6 +17,7 @@
)
__all__ = [
+ "BooleanArray",
"Categorical",
"DatetimeArray",
"IntegerArray",
diff --git a/pandas/conftest.py b/pandas/conftest.py
index b032e14d8f7e1..78e5b5e12b7e9 100644
--- a/pandas/conftest.py
+++ b/pandas/conftest.py
@@ -293,6 +293,20 @@ def compare_operators_no_eq_ne(request):
return request.param
+@pytest.fixture(
+ params=["__and__", "__rand__", "__or__", "__ror__", "__xor__", "__rxor__"]
+)
+def all_logical_operators(request):
+ """
+ Fixture for dunder names for common logical operations
+
+ * |
+ * &
+ * ^
+ """
+ return request.param
+
+
@pytest.fixture(params=[None, "gzip", "bz2", "zip", "xz"])
def compression(request):
"""
diff --git a/pandas/core/api.py b/pandas/core/api.py
index 7df2165201a99..65f0178b19187 100644
--- a/pandas/core/api.py
+++ b/pandas/core/api.py
@@ -12,6 +12,7 @@
from pandas.core.algorithms import factorize, unique, value_counts
from pandas.core.arrays import Categorical
+from pandas.core.arrays.boolean import BooleanDtype
from pandas.core.arrays.integer import (
Int8Dtype,
Int16Dtype,
diff --git a/pandas/core/arrays/__init__.py b/pandas/core/arrays/__init__.py
index 03d998707c26b..df26cd94b5ed9 100644
--- a/pandas/core/arrays/__init__.py
+++ b/pandas/core/arrays/__init__.py
@@ -4,6 +4,7 @@
ExtensionScalarOpsMixin,
try_cast_to_ea,
)
+from .boolean import BooleanArray # noqa: F401
from .categorical import Categorical # noqa: F401
from .datetimes import DatetimeArray # noqa: F401
from .integer import IntegerArray, integer_array # noqa: F401
diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py
index fa0e025c22c88..a444a4e46d0d7 100644
--- a/pandas/core/arrays/base.py
+++ b/pandas/core/arrays/base.py
@@ -1088,6 +1088,15 @@ def _add_comparison_ops(cls):
cls.__le__ = cls._create_comparison_method(operator.le)
cls.__ge__ = cls._create_comparison_method(operator.ge)
+ @classmethod
+ def _add_logical_ops(cls):
+ cls.__and__ = cls._create_logical_method(operator.and_)
+ cls.__rand__ = cls._create_logical_method(ops.rand_)
+ cls.__or__ = cls._create_logical_method(operator.or_)
+ cls.__ror__ = cls._create_logical_method(ops.ror_)
+ cls.__xor__ = cls._create_logical_method(operator.xor)
+ cls.__rxor__ = cls._create_logical_method(ops.rxor)
+
class ExtensionScalarOpsMixin(ExtensionOpsMixin):
"""
diff --git a/pandas/core/arrays/boolean.py b/pandas/core/arrays/boolean.py
new file mode 100644
index 0000000000000..c118b6fe26549
--- /dev/null
+++ b/pandas/core/arrays/boolean.py
@@ -0,0 +1,745 @@
+import numbers
+from typing import TYPE_CHECKING, Type
+import warnings
+
+import numpy as np
+
+from pandas._libs import lib
+from pandas.compat import set_function_name
+
+from pandas.core.dtypes.base import ExtensionDtype
+from pandas.core.dtypes.cast import astype_nansafe
+from pandas.core.dtypes.common import (
+ is_bool_dtype,
+ is_extension_array_dtype,
+ is_float,
+ is_float_dtype,
+ is_integer,
+ is_integer_dtype,
+ is_list_like,
+ is_scalar,
+ 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
+
+if TYPE_CHECKING:
+ from pandas._typing import Scalar
+
+
+@register_extension_dtype
+class BooleanDtype(ExtensionDtype):
+ """
+ Extension dtype for boolean data.
+
+ .. versionadded:: 1.0.0
+
+ .. warning::
+
+ BooleanDtype is considered experimental. The implementation and
+ parts of the API may change without warning.
+
+ Attributes
+ ----------
+ None
+
+ Methods
+ -------
+ None
+
+ Examples
+ --------
+ >>> pd.BooleanDtype()
+ BooleanDtype
+ """
+
+ @property
+ def na_value(self) -> "Scalar":
+ """
+ BooleanDtype uses :attr:`numpy.nan` as the missing NA value.
+
+ .. warning::
+
+ `na_value` may change in a future release.
+ """
+ return np.nan
+
+ @property
+ def type(self) -> Type:
+ return np.bool_
+
+ @property
+ def kind(self) -> str:
+ return "b"
+
+ @property
+ def name(self) -> str:
+ """
+ The alias for BooleanDtype is ``'boolean'``.
+ """
+ return "boolean"
+
+ @classmethod
+ def construct_from_string(cls, string: str) -> ExtensionDtype:
+ if string == "boolean":
+ return cls()
+ return super().construct_from_string(string)
+
+ @classmethod
+ def construct_array_type(cls) -> "Type[BooleanArray]":
+ return BooleanArray
+
+ def __repr__(self) -> str:
+ return "BooleanDtype"
+
+ @property
+ def _is_boolean(self) -> bool:
+ return True
+
+
+def coerce_to_array(values, mask=None, copy: bool = False):
+ """
+ Coerce the input values array to numpy arrays with a mask.
+
+ Parameters
+ ----------
+ values : 1D list-like
+ mask : bool 1D array, optional
+ copy : bool, default False
+ if True, copy the input
+
+ Returns
+ -------
+ tuple of (values, mask)
+ """
+ if isinstance(values, BooleanArray):
+ if mask is not None:
+ raise ValueError("cannot pass mask for BooleanArray input")
+ values, mask = values._data, values._mask
+ if copy:
+ values = values.copy()
+ mask = mask.copy()
+ return values, mask
+
+ mask_values = None
+ if isinstance(values, np.ndarray) and values.dtype == np.bool_:
+ if copy:
+ values = values.copy()
+ else:
+ # TODO conversion from integer/float ndarray can be done more efficiently
+ # (avoid roundtrip through object)
+ values_object = np.asarray(values, dtype=object)
+
+ inferred_dtype = lib.infer_dtype(values_object, skipna=True)
+ integer_like = ("floating", "integer", "mixed-integer-float")
+ if inferred_dtype not in ("boolean", "empty") + integer_like:
+ raise TypeError("Need to pass bool-like values")
+
+ mask_values = isna(values_object)
+ values = np.zeros(len(values), dtype=bool)
+ values[~mask_values] = values_object[~mask_values].astype(bool)
+
+ # if the values were integer-like, validate it were actually 0/1's
+ if inferred_dtype in integer_like:
+ if not np.all(
+ values[~mask_values].astype(float)
+ == values_object[~mask_values].astype(float)
+ ):
+ raise TypeError("Need to pass bool-like values")
+
+ if mask is None and mask_values is None:
+ mask = np.zeros(len(values), dtype=bool)
+ elif mask is None:
+ mask = mask_values
+ else:
+ if isinstance(mask, np.ndarray) and mask.dtype == np.bool_:
+ if mask_values is not None:
+ mask = mask | mask_values
+ else:
+ if copy:
+ mask = mask.copy()
+ else:
+ mask = np.array(mask, dtype=bool)
+ if mask_values is not None:
+ mask = mask | mask_values
+
+ if not values.ndim == 1:
+ raise ValueError("values must be a 1D list-like")
+ if not mask.ndim == 1:
+ raise ValueError("mask must be a 1D list-like")
+
+ return values, mask
+
+
+class BooleanArray(ExtensionArray, ExtensionOpsMixin):
+ """
+ Array of boolean (True/False) data with missing values.
+
+ This is a pandas Extension array for boolean data, under the hood
+ represented by 2 numpy arrays: a boolean array with the data and
+ a boolean array with the mask (True indicating missing).
+
+ To construct an BooleanArray from generic array-like input, use
+ :func:`pandas.array` specifying ``dtype="boolean"`` (see examples
+ below).
+
+ .. versionadded:: 1.0.0
+
+ .. warning::
+
+ BooleanArray is considered experimental. The implementation and
+ parts of the API may change without warning.
+
+ Parameters
+ ----------
+ values : numpy.ndarray
+ A 1-d boolean-dtype array with the data.
+ mask : numpy.ndarray
+ A 1-d boolean-dtype array indicating missing values (True
+ indicates missing).
+ copy : bool, default False
+ Whether to copy the `values` and `mask` arrays.
+
+ Attributes
+ ----------
+ None
+
+ Methods
+ -------
+ None
+
+ Returns
+ -------
+ BooleanArray
+
+ Examples
+ --------
+ Create an BooleanArray with :func:`pandas.array`:
+
+ >>> pd.array([True, False, None], dtype="boolean")
+ <BooleanArray>
+ [True, False, NaN]
+ Length: 3, dtype: boolean
+ """
+
+ 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(
+ "values should be boolean numpy array. Use "
+ "the 'array' function instead"
+ )
+ if not (isinstance(mask, np.ndarray) and mask.dtype == np.bool_):
+ raise TypeError(
+ "mask should be boolean numpy array. Use "
+ "the 'array' function instead"
+ )
+ if not values.ndim == 1:
+ raise ValueError("values must be a 1D array")
+ if not mask.ndim == 1:
+ raise ValueError("mask must be a 1D array")
+
+ if copy:
+ values = values.copy()
+ mask = mask.copy()
+
+ self._data = values
+ self._mask = mask
+ self._dtype = BooleanDtype()
+
+ @property
+ def dtype(self):
+ return self._dtype
+
+ @classmethod
+ def _from_sequence(cls, scalars, dtype=None, copy: bool = False):
+ if dtype:
+ assert dtype == "boolean"
+ values, mask = coerce_to_array(scalars, copy=copy)
+ return BooleanArray(values, mask)
+
+ @classmethod
+ def _from_factorized(cls, values, original: "BooleanArray"):
+ return cls._from_sequence(values, dtype=original.dtype)
+
+ def _formatter(self, boxed=False):
+ def fmt(x):
+ if isna(x):
+ return "NaN"
+ return str(x)
+
+ return fmt
+
+ def __getitem__(self, item):
+ if is_integer(item):
+ if self._mask[item]:
+ return self.dtype.na_value
+ return self._data[item]
+ return type(self)(self._data[item], self._mask[item])
+
+ def _coerce_to_ndarray(self, force_bool: bool = False):
+ """
+ Coerce to an ndarary of object dtype or bool dtype (if force_bool=True).
+
+ Parameters
+ ----------
+ force_bool : bool, default False
+ If True, return bool array or raise error if not possible (in
+ presence of missing values)
+ """
+ if force_bool:
+ if not self.isna().any():
+ return self._data
+ else:
+ raise ValueError(
+ "cannot convert to bool numpy array in presence of missing values"
+ )
+ data = self._data.astype(object)
+ data[self._mask] = self._na_value
+ 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
+ """
+ if dtype is not None:
+ if is_bool_dtype(dtype):
+ return self._coerce_to_ndarray(force_bool=True)
+ # TODO can optimize this to not go through object dtype for
+ # numeric dtypes
+ arr = self._coerce_to_ndarray()
+ return arr.astype(dtype, copy=False)
+ # by default (no dtype specified), return an object array
+ return self._coerce_to_ndarray()
+
+ 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):
+ # For BooleanArray inputs, we apply the ufunc to ._data
+ # and mask the result.
+ if method == "reduce":
+ # Not clear how to handle missing values in reductions. Raise.
+ raise NotImplementedError("The 'reduce' method is not supported.")
+ out = kwargs.get("out", ())
+
+ for x in inputs + out:
+ if not isinstance(x, self._HANDLED_TYPES + (BooleanArray,)):
+ return NotImplemented
+
+ # for binary ops, use our custom dunder methods
+ result = ops.maybe_dispatch_ufunc_to_dunder_op(
+ self, ufunc, method, *inputs, **kwargs
+ )
+ if result is not NotImplemented:
+ return result
+
+ mask = np.zeros(len(self), dtype=bool)
+ inputs2 = []
+ for x in inputs:
+ if isinstance(x, BooleanArray):
+ mask |= x._mask
+ inputs2.append(x._data)
+ else:
+ inputs2.append(x)
+
+ def reconstruct(x):
+ # we don't worry about scalar `x` here, since we
+ # raise for reduce up above.
+
+ if is_bool_dtype(x.dtype):
+ m = mask.copy()
+ return BooleanArray(x, m)
+ else:
+ x[mask] = np.nan
+ return x
+
+ result = getattr(ufunc, method)(*inputs2, **kwargs)
+ if isinstance(result, tuple):
+ tuple(reconstruct(x) for x in result)
+ 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:
+ value = [value]
+ value, mask = coerce_to_array(value)
+
+ if _is_scalar:
+ value = value[0]
+ mask = mask[0]
+
+ 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'.
+
+ Parameters
+ ----------
+ dtype : str or dtype
+ Typecode or data-type to which the array is cast.
+ copy : bool, default True
+ Whether to copy the data, even if not necessary. If False,
+ a copy is made only if the old dtype does not match the
+ new dtype.
+
+ Returns
+ -------
+ array : ndarray or ExtensionArray
+ NumPy ndarray, BooleanArray or IntergerArray with 'dtype' for its dtype.
+
+ Raises
+ ------
+ TypeError
+ if incompatible type with an BooleanDtype, equivalent of same_kind
+ casting
+ """
+ dtype = pandas_dtype(dtype)
+
+ if isinstance(dtype, BooleanDtype):
+ values, mask = coerce_to_array(self, copy=copy)
+ return BooleanArray(values, mask, copy=False)
+
+ if is_bool_dtype(dtype):
+ # astype_nansafe converts np.nan to True
+ if self.isna().any():
+ raise ValueError("cannot convert float NaN to bool")
+ else:
+ return self._data.astype(dtype, copy=copy)
+ if is_extension_array_dtype(dtype) and is_integer_dtype(dtype):
+ from pandas.core.arrays import IntegerArray
+
+ return IntegerArray(
+ self._data.astype(dtype.numpy_dtype), self._mask.copy(), copy=False
+ )
+ # coerce
+ data = self._coerce_to_ndarray()
+ return astype_nansafe(data, dtype, copy=None)
+
+ def value_counts(self, dropna=True):
+ """
+ Returns a Series containing counts of each category.
+
+ Every category will have an entry, even those with a count of 0.
+
+ Parameters
+ ----------
+ dropna : bool, default True
+ Don't include counts of NaN.
+
+ Returns
+ -------
+ counts : Series
+
+ See Also
+ --------
+ Series.value_counts
+
+ """
+
+ from pandas import Index, Series
+
+ # compute counts on the data with no nans
+ data = self._data[~self._mask]
+ value_counts = Index(data).value_counts()
+ array = value_counts.values
+
+ # TODO(extension)
+ # if we have allow Index to hold an ExtensionArray
+ # this is easier
+ index = value_counts.index.values.astype(bool).astype(object)
+
+ # if we want nans, count the mask
+ if not dropna:
+
+ # TODO(extension)
+ # appending to an Index *always* infers
+ # w/o passing the dtype
+ array = np.append(array, [self._mask.sum()])
+ index = Index(
+ np.concatenate([index, np.array([np.nan], dtype=object)]), dtype=object
+ )
+
+ return Series(array, index=index)
+
+ def _values_for_argsort(self) -> np.ndarray:
+ """
+ Return values for sorting.
+
+ Returns
+ -------
+ ndarray
+ The transformed values should maintain the ordering between values
+ within the array.
+
+ See Also
+ --------
+ ExtensionArray.argsort
+ """
+ data = self._data.copy()
+ data[self._mask] = -1
+ return data
+
+ @classmethod
+ def _create_logical_method(cls, op):
+ def logical_method(self, other):
+
+ if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndexClass)):
+ # Rely on pandas to unbox and dispatch to us.
+ return NotImplemented
+
+ other = lib.item_from_zerodim(other)
+ mask = None
+
+ if isinstance(other, BooleanArray):
+ other, mask = other._data, other._mask
+ elif is_list_like(other):
+ other = np.asarray(other, dtype="bool")
+ if other.ndim > 1:
+ raise NotImplementedError(
+ "can only perform ops with 1-d structures"
+ )
+ if len(self) != len(other):
+ raise ValueError("Lengths must match to compare")
+ other, mask = coerce_to_array(other, copy=False)
+
+ # numpy will show a DeprecationWarning on invalid elementwise
+ # comparisons, this will raise in the future
+ with warnings.catch_warnings():
+ warnings.filterwarnings("ignore", "elementwise", FutureWarning)
+ with np.errstate(all="ignore"):
+ result = op(self._data, other)
+
+ # nans propagate
+ if mask is None:
+ mask = self._mask
+ else:
+ mask = self._mask | mask
+
+ return BooleanArray(result, mask)
+
+ name = "__{name}__".format(name=op.__name__)
+ return set_function_name(logical_method, name, cls)
+
+ @classmethod
+ def _create_comparison_method(cls, op):
+ op_name = op.__name__
+
+ def cmp_method(self, other):
+
+ if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndexClass)):
+ # Rely on pandas to unbox and dispatch to us.
+ return NotImplemented
+
+ other = lib.item_from_zerodim(other)
+ mask = None
+
+ if isinstance(other, BooleanArray):
+ other, mask = other._data, other._mask
+
+ elif is_list_like(other):
+ other = np.asarray(other)
+ if other.ndim > 1:
+ raise NotImplementedError(
+ "can only perform ops with 1-d structures"
+ )
+ if len(self) != len(other):
+ raise ValueError("Lengths must match to compare")
+
+ # numpy will show a DeprecationWarning on invalid elementwise
+ # comparisons, this will raise in the future
+ with warnings.catch_warnings():
+ warnings.filterwarnings("ignore", "elementwise", FutureWarning)
+ with np.errstate(all="ignore"):
+ result = op(self._data, other)
+
+ # nans propagate
+ if mask is None:
+ mask = self._mask
+ else:
+ mask = self._mask | mask
+
+ result[mask] = op_name == "ne"
+ return BooleanArray(result, np.zeros(len(result), dtype=bool), copy=False)
+
+ name = "__{name}__".format(name=op.__name__)
+ return set_function_name(cmp_method, name, cls)
+
+ def _reduce(self, name, skipna=True, **kwargs):
+ data = self._data
+ mask = self._mask
+
+ # coerce to a nan-aware float if needed
+ if mask.any():
+ data = self._data.astype("float64")
+ data[mask] = self._na_value
+
+ op = getattr(nanops, "nan" + name)
+ result = op(data, axis=0, skipna=skipna, mask=mask, **kwargs)
+
+ # if we have a boolean op, don't coerce
+ if name in ["any", "all"]:
+ pass
+
+ # if we have numeric op that would result in an int, coerce to int if possible
+ elif name in ["sum", "prod"] and notna(result):
+ int_result = np.int64(result)
+ if int_result == result:
+ result = int_result
+
+ elif name in ["min", "max"] and notna(result):
+ result = np.bool_(result)
+
+ return result
+
+ def _maybe_mask_result(self, result, mask, other, op_name):
+ """
+ Parameters
+ ----------
+ result : array-like
+ mask : array-like bool
+ other : scalar or array-like
+ op_name : str
+ """
+ # if we have a float operand we are by-definition
+ # a float result
+ # or our op is a divide
+ if (is_float_dtype(other) or is_float(other)) or (
+ op_name in ["rtruediv", "truediv"]
+ ):
+ result[mask] = np.nan
+ return result
+
+ if is_bool_dtype(result):
+ return BooleanArray(result, mask, copy=False)
+
+ elif is_integer_dtype(result):
+ from pandas.core.arrays import IntegerArray
+
+ return IntegerArray(result, mask, copy=False)
+ else:
+ result[mask] = np.nan
+ return result
+
+ @classmethod
+ def _create_arithmetic_method(cls, op):
+ op_name = op.__name__
+
+ def boolean_arithmetic_method(self, other):
+
+ if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndexClass)):
+ # Rely on pandas to unbox and dispatch to us.
+ return NotImplemented
+
+ other = lib.item_from_zerodim(other)
+ mask = None
+
+ if isinstance(other, BooleanArray):
+ other, mask = other._data, other._mask
+
+ elif is_list_like(other):
+ other = np.asarray(other)
+ if other.ndim > 1:
+ raise NotImplementedError(
+ "can only perform ops with 1-d structures"
+ )
+ if len(self) != len(other):
+ raise ValueError("Lengths must match")
+
+ # nans propagate
+ if mask is None:
+ mask = self._mask
+ else:
+ mask = self._mask | mask
+
+ with np.errstate(all="ignore"):
+ result = op(self._data, other)
+
+ # divmod returns a tuple
+ if op_name == "divmod":
+ div, mod = result
+ return (
+ self._maybe_mask_result(div, mask, other, "floordiv"),
+ self._maybe_mask_result(mod, mask, other, "mod"),
+ )
+
+ return self._maybe_mask_result(result, mask, other, op_name)
+
+ name = "__{name}__".format(name=op_name)
+ return set_function_name(boolean_arithmetic_method, name, cls)
+
+
+BooleanArray._add_logical_ops()
+BooleanArray._add_comparison_ops()
+BooleanArray._add_arithmetic_ops()
diff --git a/pandas/core/dtypes/missing.py b/pandas/core/dtypes/missing.py
index aeba4eebc498e..25d6f87143d72 100644
--- a/pandas/core/dtypes/missing.py
+++ b/pandas/core/dtypes/missing.py
@@ -448,7 +448,7 @@ def array_equivalent(left, right, strict_nan: bool = False) -> bool:
return False
else:
try:
- if np.any(left_value != right_value):
+ if np.any(np.asarray(left_value != right_value)):
return False
except TypeError as err:
if "Cannot compare tz-naive" in str(err):
diff --git a/pandas/tests/api/test_api.py b/pandas/tests/api/test_api.py
index 5d11e160bbd71..1282aa6edd538 100644
--- a/pandas/tests/api/test_api.py
+++ b/pandas/tests/api/test_api.py
@@ -80,6 +80,7 @@ class TestPDApi(Base):
"PeriodDtype",
"IntervalDtype",
"DatetimeTZDtype",
+ "BooleanDtype",
"Int8Dtype",
"Int16Dtype",
"Int32Dtype",
diff --git a/pandas/tests/arrays/test_boolean.py b/pandas/tests/arrays/test_boolean.py
new file mode 100644
index 0000000000000..5cfc7c3837875
--- /dev/null
+++ b/pandas/tests/arrays/test_boolean.py
@@ -0,0 +1,509 @@
+import operator
+
+import numpy as np
+import pytest
+
+import pandas.util._test_decorators as td
+
+import pandas as pd
+from pandas.arrays import BooleanArray
+from pandas.core.arrays.boolean import coerce_to_array
+from pandas.tests.extension.base import BaseOpsUtil
+import pandas.util.testing as tm
+
+
+def make_data():
+ return [True, False] * 4 + [np.nan] + [True, False] * 44 + [np.nan] + [True, False]
+
+
+@pytest.fixture
+def dtype():
+ return pd.BooleanDtype()
+
+
+@pytest.fixture
+def data(dtype):
+ return pd.array(make_data(), dtype=dtype)
+
+
+def test_boolean_array_constructor():
+ values = np.array([True, False, True, False], dtype="bool")
+ mask = np.array([False, False, False, True], dtype="bool")
+
+ result = BooleanArray(values, mask)
+ expected = pd.array([True, False, True, None], dtype="boolean")
+ tm.assert_extension_array_equal(result, expected)
+
+ with pytest.raises(TypeError, match="values should be boolean numpy array"):
+ BooleanArray(values.tolist(), mask)
+
+ with pytest.raises(TypeError, match="mask should be boolean numpy array"):
+ BooleanArray(values, mask.tolist())
+
+ with pytest.raises(TypeError, match="values should be boolean numpy array"):
+ BooleanArray(values.astype(int), mask)
+
+ with pytest.raises(TypeError, match="mask should be boolean numpy array"):
+ BooleanArray(values, None)
+
+ with pytest.raises(ValueError, match="values must be a 1D array"):
+ BooleanArray(values.reshape(1, -1), mask)
+
+ with pytest.raises(ValueError, match="mask must be a 1D array"):
+ BooleanArray(values, mask.reshape(1, -1))
+
+
+def test_boolean_array_constructor_copy():
+ values = np.array([True, False, True, False], dtype="bool")
+ mask = np.array([False, False, False, True], dtype="bool")
+
+ result = BooleanArray(values, mask)
+ assert result._data is values
+ assert result._mask is mask
+
+ result = BooleanArray(values, mask, copy=True)
+ assert result._data is not values
+ assert result._mask is not mask
+
+
+def test_to_boolean_array():
+ expected = BooleanArray(
+ np.array([True, False, True]), np.array([False, False, False])
+ )
+
+ result = pd.array([True, False, True], dtype="boolean")
+ tm.assert_extension_array_equal(result, expected)
+ result = pd.array(np.array([True, False, True]), dtype="boolean")
+ tm.assert_extension_array_equal(result, expected)
+ result = pd.array(np.array([True, False, True], dtype=object), dtype="boolean")
+ tm.assert_extension_array_equal(result, expected)
+
+ # with missing values
+ expected = BooleanArray(
+ np.array([True, False, True]), np.array([False, False, True])
+ )
+
+ result = pd.array([True, False, None], dtype="boolean")
+ tm.assert_extension_array_equal(result, expected)
+ result = pd.array(np.array([True, False, None], dtype=object), dtype="boolean")
+ tm.assert_extension_array_equal(result, expected)
+
+
+def test_to_boolean_array_all_none():
+ expected = BooleanArray(np.array([True, True, True]), np.array([True, True, True]))
+
+ result = pd.array([None, None, None], dtype="boolean")
+ tm.assert_extension_array_equal(result, expected)
+ result = pd.array(np.array([None, None, None], dtype=object), dtype="boolean")
+ tm.assert_extension_array_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+ "a, b",
+ [
+ ([True, None], [True, np.nan]),
+ ([None], [np.nan]),
+ ([None, np.nan], [np.nan, np.nan]),
+ ([np.nan, np.nan], [np.nan, np.nan]),
+ ],
+)
+def test_to_boolean_array_none_is_nan(a, b):
+ result = pd.array(a, dtype="boolean")
+ expected = pd.array(b, dtype="boolean")
+ tm.assert_extension_array_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+ "values",
+ [
+ ["foo", "bar"],
+ ["1", "2"],
+ # "foo",
+ [1, 2],
+ [1.0, 2.0],
+ pd.date_range("20130101", periods=2),
+ np.array(["foo"]),
+ [np.nan, {"a": 1}],
+ ],
+)
+def test_to_boolean_array_error(values):
+ # error in converting existing arrays to BooleanArray
+ with pytest.raises(TypeError):
+ pd.array(values, dtype="boolean")
+
+
+def test_to_boolean_array_integer_like():
+ # integers of 0's and 1's
+ result = pd.array([1, 0, 1, 0], dtype="boolean")
+ expected = pd.array([True, False, True, False], dtype="boolean")
+ tm.assert_extension_array_equal(result, expected)
+
+ result = pd.array(np.array([1, 0, 1, 0]), dtype="boolean")
+ tm.assert_extension_array_equal(result, expected)
+
+ result = pd.array(np.array([1.0, 0.0, 1.0, 0.0]), dtype="boolean")
+ tm.assert_extension_array_equal(result, expected)
+
+ # with missing values
+ result = pd.array([1, 0, 1, None], dtype="boolean")
+ expected = pd.array([True, False, True, None], dtype="boolean")
+ tm.assert_extension_array_equal(result, expected)
+
+ result = pd.array(np.array([1.0, 0.0, 1.0, np.nan]), dtype="boolean")
+ tm.assert_extension_array_equal(result, expected)
+
+
+def test_coerce_to_array():
+ # TODO this is currently not public API
+ values = np.array([True, False, True, False], dtype="bool")
+ mask = np.array([False, False, False, True], dtype="bool")
+ result = BooleanArray(*coerce_to_array(values, mask=mask))
+ expected = BooleanArray(values, mask)
+ tm.assert_extension_array_equal(result, expected)
+ assert result._data is values
+ assert result._mask is mask
+ result = BooleanArray(*coerce_to_array(values, mask=mask, copy=True))
+ expected = BooleanArray(values, mask)
+ tm.assert_extension_array_equal(result, expected)
+ assert result._data is not values
+ assert result._mask is not mask
+
+ # mixed missing from values and mask
+ values = [True, False, None, False]
+ mask = np.array([False, False, False, True], dtype="bool")
+ result = BooleanArray(*coerce_to_array(values, mask=mask))
+ expected = BooleanArray(
+ np.array([True, False, True, True]), np.array([False, False, True, True])
+ )
+ tm.assert_extension_array_equal(result, expected)
+ result = BooleanArray(*coerce_to_array(np.array(values, dtype=object), mask=mask))
+ tm.assert_extension_array_equal(result, expected)
+ result = BooleanArray(*coerce_to_array(values, mask=mask.tolist()))
+ tm.assert_extension_array_equal(result, expected)
+
+ # raise errors for wrong dimension
+ values = np.array([True, False, True, False], dtype="bool")
+ mask = np.array([False, False, False, True], dtype="bool")
+
+ with pytest.raises(ValueError, match="values must be a 1D list-like"):
+ coerce_to_array(values.reshape(1, -1))
+
+ with pytest.raises(ValueError, match="mask must be a 1D list-like"):
+ coerce_to_array(values, mask=mask.reshape(1, -1))
+
+
+def test_coerce_to_array_from_boolean_array():
+ # passing BooleanArray to coerce_to_array
+ values = np.array([True, False, True, False], dtype="bool")
+ mask = np.array([False, False, False, True], dtype="bool")
+ arr = BooleanArray(values, mask)
+ result = BooleanArray(*coerce_to_array(arr))
+ tm.assert_extension_array_equal(result, arr)
+ # no copy
+ assert result._data is arr._data
+ assert result._mask is arr._mask
+
+ result = BooleanArray(*coerce_to_array(arr), copy=True)
+ tm.assert_extension_array_equal(result, arr)
+ assert result._data is not arr._data
+ assert result._mask is not arr._mask
+
+ with pytest.raises(ValueError, match="cannot pass mask for BooleanArray input"):
+ coerce_to_array(arr, mask=mask)
+
+
+def test_coerce_to_numpy_array():
+ # with missing values -> object dtype
+ arr = pd.array([True, False, None], dtype="boolean")
+ result = np.array(arr)
+ expected = np.array([True, False, None], dtype="object")
+ tm.assert_numpy_array_equal(result, expected)
+
+ # also with no missing values -> object dtype
+ arr = pd.array([True, False, True], dtype="boolean")
+ result = np.array(arr)
+ expected = np.array([True, False, True], dtype="object")
+ tm.assert_numpy_array_equal(result, expected)
+
+ # force bool dtype
+ result = np.array(arr, dtype="bool")
+ expected = np.array([True, False, True], dtype="bool")
+ tm.assert_numpy_array_equal(result, expected)
+ # with missing values will raise error
+ arr = pd.array([True, False, None], dtype="boolean")
+ with pytest.raises(ValueError):
+ np.array(arr, dtype="bool")
+
+
+def test_astype():
+ # with missing values
+ arr = pd.array([True, False, None], dtype="boolean")
+ msg = "cannot convert float NaN to"
+
+ with pytest.raises(ValueError, match=msg):
+ arr.astype("int64")
+
+ with pytest.raises(ValueError, match=msg):
+ arr.astype("bool")
+
+ result = arr.astype("float64")
+ expected = np.array([1, 0, np.nan], dtype="float64")
+ tm.assert_numpy_array_equal(result, expected)
+
+ # no missing values
+ arr = pd.array([True, False, True], dtype="boolean")
+ result = arr.astype("int64")
+ expected = np.array([1, 0, 1], dtype="int64")
+ tm.assert_numpy_array_equal(result, expected)
+
+ result = arr.astype("bool")
+ expected = np.array([True, False, True], dtype="bool")
+ tm.assert_numpy_array_equal(result, expected)
+
+
+def test_astype_to_boolean_array():
+ # astype to BooleanArray
+ arr = pd.array([True, False, None], dtype="boolean")
+
+ result = arr.astype("boolean")
+ tm.assert_extension_array_equal(result, arr)
+ result = arr.astype(pd.BooleanDtype())
+ tm.assert_extension_array_equal(result, arr)
+
+
+def test_astype_to_integer_array():
+ # astype to IntegerArray
+ arr = pd.array([True, False, None], dtype="boolean")
+
+ result = arr.astype("Int64")
+ expected = pd.array([1, 0, None], dtype="Int64")
+ tm.assert_extension_array_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+ "ufunc", [np.add, np.logical_or, np.logical_and, np.logical_xor]
+)
+def test_ufuncs_binary(ufunc):
+ # two BooleanArrays
+ a = pd.array([True, False, None], dtype="boolean")
+ result = ufunc(a, a)
+ expected = pd.array(ufunc(a._data, a._data), dtype="boolean")
+ expected[a._mask] = np.nan
+ tm.assert_extension_array_equal(result, expected)
+
+ s = pd.Series(a)
+ result = ufunc(s, a)
+ expected = pd.Series(ufunc(a._data, a._data), dtype="boolean")
+ expected[a._mask] = np.nan
+ tm.assert_series_equal(result, expected)
+
+ # Boolean with numpy array
+ arr = np.array([True, True, False])
+ result = ufunc(a, arr)
+ expected = pd.array(ufunc(a._data, arr), dtype="boolean")
+ expected[a._mask] = np.nan
+ tm.assert_extension_array_equal(result, expected)
+
+ result = ufunc(arr, a)
+ expected = pd.array(ufunc(arr, a._data), dtype="boolean")
+ expected[a._mask] = np.nan
+ tm.assert_extension_array_equal(result, expected)
+
+ # BooleanArray with scalar
+ result = ufunc(a, True)
+ expected = pd.array(ufunc(a._data, True), dtype="boolean")
+ expected[a._mask] = np.nan
+ tm.assert_extension_array_equal(result, expected)
+
+ result = ufunc(True, a)
+ expected = pd.array(ufunc(True, a._data), dtype="boolean")
+ expected[a._mask] = np.nan
+ tm.assert_extension_array_equal(result, expected)
+
+ # not handled types
+ with pytest.raises(TypeError):
+ ufunc(a, "test")
+
+
+@pytest.mark.parametrize("ufunc", [np.logical_not])
+def test_ufuncs_unary(ufunc):
+ a = pd.array([True, False, None], dtype="boolean")
+ result = ufunc(a)
+ expected = pd.array(ufunc(a._data), dtype="boolean")
+ expected[a._mask] = np.nan
+ tm.assert_extension_array_equal(result, expected)
+
+ s = pd.Series(a)
+ result = ufunc(s)
+ expected = pd.Series(ufunc(a._data), dtype="boolean")
+ expected[a._mask] = np.nan
+ tm.assert_series_equal(result, expected)
+
+
+@pytest.mark.parametrize("values", [[True, False], [True, None]])
+def test_ufunc_reduce_raises(values):
+ a = pd.array(values, dtype="boolean")
+ with pytest.raises(NotImplementedError):
+ np.add.reduce(a)
+
+
+class TestLogicalOps(BaseOpsUtil):
+ def get_op_from_name(self, op_name):
+ short_opname = op_name.strip("_")
+ short_opname = short_opname if "xor" in short_opname else short_opname + "_"
+ try:
+ op = getattr(operator, short_opname)
+ except AttributeError:
+ # Assume it is the reverse operator
+ rop = getattr(operator, short_opname[1:])
+ op = lambda x, y: rop(y, x)
+
+ return op
+
+ def _compare_other(self, data, op_name, other):
+ op = self.get_op_from_name(op_name)
+
+ # array
+ result = pd.Series(op(data, other))
+ expected = pd.Series(op(data._data, other), dtype="boolean")
+
+ # fill the nan locations
+ expected[data._mask] = np.nan
+
+ tm.assert_series_equal(result, expected)
+
+ # series
+ s = pd.Series(data)
+ result = op(s, other)
+
+ expected = pd.Series(data._data)
+ expected = op(expected, other)
+ expected = pd.Series(expected, dtype="boolean")
+
+ # fill the nan locations
+ expected[data._mask] = np.nan
+
+ tm.assert_series_equal(result, expected)
+
+ def test_scalar(self, data, all_logical_operators):
+ op_name = all_logical_operators
+ self._compare_other(data, op_name, True)
+
+ def test_array(self, data, all_logical_operators):
+ op_name = all_logical_operators
+ other = pd.array([True] * len(data), dtype="boolean")
+ self._compare_other(data, op_name, other)
+ other = np.array([True] * len(data))
+ self._compare_other(data, op_name, other)
+ other = pd.Series([True] * len(data), dtype="boolean")
+ self._compare_other(data, op_name, other)
+
+
+class TestComparisonOps(BaseOpsUtil):
+ def _compare_other(self, data, op_name, other):
+ op = self.get_op_from_name(op_name)
+
+ # array
+ result = pd.Series(op(data, other))
+ expected = pd.Series(op(data._data, other), dtype="boolean")
+
+ # fill the nan locations
+ expected[data._mask] = op_name == "__ne__"
+
+ tm.assert_series_equal(result, expected)
+
+ # series
+ s = pd.Series(data)
+ result = op(s, other)
+
+ expected = pd.Series(data._data)
+ expected = op(expected, other)
+ expected = expected.astype("boolean")
+
+ # fill the nan locations
+ expected[data._mask] = op_name == "__ne__"
+
+ tm.assert_series_equal(result, expected)
+
+ def test_compare_scalar(self, data, all_compare_operators):
+ op_name = all_compare_operators
+ self._compare_other(data, op_name, True)
+
+ def test_compare_array(self, data, all_compare_operators):
+ op_name = all_compare_operators
+ other = pd.array([True] * len(data), dtype="boolean")
+ self._compare_other(data, op_name, other)
+ other = np.array([True] * len(data))
+ self._compare_other(data, op_name, other)
+ other = pd.Series([True] * len(data))
+ self._compare_other(data, op_name, other)
+
+
+class TestArithmeticOps(BaseOpsUtil):
+ def test_error(self, data, all_arithmetic_operators):
+ # invalid ops
+
+ op = all_arithmetic_operators
+ s = pd.Series(data)
+ ops = getattr(s, op)
+ opa = getattr(data, op)
+
+ # invalid scalars
+ with pytest.raises(TypeError):
+ ops("foo")
+ with pytest.raises(TypeError):
+ ops(pd.Timestamp("20180101"))
+
+ # invalid array-likes
+ if op not in ("__mul__", "__rmul__"):
+ # TODO(extension) numpy's mul with object array sees booleans as numbers
+ with pytest.raises(TypeError):
+ ops(pd.Series("foo", index=s.index))
+
+ # 2d
+ result = opa(pd.DataFrame({"A": s}))
+ assert result is NotImplemented
+
+ with pytest.raises(NotImplementedError):
+ opa(np.arange(len(s)).reshape(-1, len(s)))
+
+
+@pytest.mark.parametrize("dropna", [True, False])
+def test_reductions_return_types(dropna, data, all_numeric_reductions):
+ op = all_numeric_reductions
+ s = pd.Series(data)
+ if dropna:
+ s = s.dropna()
+
+ if op in ("sum", "prod"):
+ assert isinstance(getattr(s, op)(), np.int64)
+ elif op in ("min", "max"):
+ assert isinstance(getattr(s, op)(), np.bool_)
+ else:
+ # "mean", "std", "var", "median", "kurt", "skew"
+ assert isinstance(getattr(s, op)(), np.float64)
+
+
+# TODO when BooleanArray coerces to object dtype numpy array, need to do conversion
+# manually in the indexing code
+# def test_indexing_boolean_mask():
+# arr = pd.array([1, 2, 3, 4], dtype="Int64")
+# mask = pd.array([True, False, True, False], dtype="boolean")
+# result = arr[mask]
+# expected = pd.array([1, 3], dtype="Int64")
+# tm.assert_extension_array_equal(result, expected)
+
+# # missing values -> error
+# mask = pd.array([True, False, True, None], dtype="boolean")
+# with pytest.raises(IndexError):
+# result = arr[mask]
+
+
+@td.skip_if_no("pyarrow", min_version="0.15.0")
+def test_arrow_array(data):
+ # protocol added in 0.15.0
+ import pyarrow as pa
+
+ arr = pa.array(data)
+ expected = pa.array(np.array(data, dtype=object), type=pa.bool_(), from_pandas=True)
+ assert arr.equals(expected)
diff --git a/pandas/tests/dtypes/test_common.py b/pandas/tests/dtypes/test_common.py
index 6d91d13027f69..912fce6339716 100644
--- a/pandas/tests/dtypes/test_common.py
+++ b/pandas/tests/dtypes/test_common.py
@@ -529,6 +529,9 @@ def test_is_bool_dtype():
assert com.is_bool_dtype(np.array([True, False]))
assert com.is_bool_dtype(pd.Index([True, False]))
+ assert com.is_bool_dtype(pd.BooleanDtype())
+ assert com.is_bool_dtype(pd.array([True, False, None], dtype="boolean"))
+
@pytest.mark.filterwarnings("ignore:'is_extension_type' is deprecated:FutureWarning")
@pytest.mark.parametrize(
diff --git a/pandas/tests/extension/test_boolean.py b/pandas/tests/extension/test_boolean.py
new file mode 100644
index 0000000000000..089dd798b2512
--- /dev/null
+++ b/pandas/tests/extension/test_boolean.py
@@ -0,0 +1,333 @@
+"""
+This file contains a minimal set of tests for compliance with the extension
+array interface test suite, and should contain no other tests.
+The test suite for the full functionality of the array is located in
+`pandas/tests/arrays/`.
+
+The tests in this file are inherited from the BaseExtensionTests, and only
+minimal tweaks should be applied to get the tests passing (by overwriting a
+parent method).
+
+Additional tests should either be added to one of the BaseExtensionTests
+classes (if they are relevant for the extension interface for all dtypes), or
+be added to the array-specific tests in `pandas/tests/arrays/`.
+
+"""
+import numpy as np
+import pytest
+
+from pandas.compat.numpy import _np_version_under1p14
+
+import pandas as pd
+from pandas.core.arrays.boolean import BooleanDtype
+from pandas.tests.extension import base
+import pandas.util.testing as tm
+
+
+def make_data():
+ return [True, False] * 4 + [np.nan] + [True, False] * 44 + [np.nan] + [True, False]
+
+
+@pytest.fixture
+def dtype():
+ return BooleanDtype()
+
+
+@pytest.fixture
+def data(dtype):
+ return pd.array(make_data(), dtype=dtype)
+
+
+@pytest.fixture
+def data_for_twos(dtype):
+ return pd.array(np.ones(100), dtype=dtype)
+
+
+@pytest.fixture
+def data_missing(dtype):
+ return pd.array([np.nan, True], dtype=dtype)
+
+
+@pytest.fixture
+def data_for_sorting(dtype):
+ return pd.array([True, True, False], dtype=dtype)
+
+
+@pytest.fixture
+def data_missing_for_sorting(dtype):
+ return pd.array([True, np.nan, False], dtype=dtype)
+
+
+@pytest.fixture
+def na_cmp():
+ # we are np.nan
+ return lambda x, y: np.isnan(x) and np.isnan(y)
+
+
+@pytest.fixture
+def na_value():
+ return np.nan
+
+
+@pytest.fixture
+def data_for_grouping(dtype):
+ b = True
+ a = False
+ na = np.nan
+ return pd.array([b, b, na, na, a, a, b], dtype=dtype)
+
+
+class TestDtype(base.BaseDtypeTests):
+ pass
+
+
+class TestInterface(base.BaseInterfaceTests):
+ pass
+
+
+class TestConstructors(base.BaseConstructorsTests):
+ pass
+
+
+class TestGetitem(base.BaseGetitemTests):
+ pass
+
+
+class TestSetitem(base.BaseSetitemTests):
+ pass
+
+
+class TestMissing(base.BaseMissingTests):
+ pass
+
+
+class TestArithmeticOps(base.BaseArithmeticOpsTests):
+ def check_opname(self, s, op_name, other, exc=None):
+ # overwriting to indicate ops don't raise an error
+ super().check_opname(s, op_name, other, exc=None)
+
+ def _check_op(self, s, op, other, op_name, exc=NotImplementedError):
+ if exc is None:
+ if op_name in ("__sub__", "__rsub__"):
+ # subtraction for bools raises TypeError (but not yet in 1.13)
+ if _np_version_under1p14:
+ pytest.skip("__sub__ does not yet raise in numpy 1.13")
+ with pytest.raises(TypeError):
+ op(s, other)
+
+ return
+
+ result = op(s, other)
+ expected = s.combine(other, op)
+
+ if op_name in (
+ "__floordiv__",
+ "__rfloordiv__",
+ "__pow__",
+ "__rpow__",
+ "__mod__",
+ "__rmod__",
+ ):
+ # combine keeps boolean type
+ expected = expected.astype("Int8")
+ elif op_name in ("__truediv__", "__rtruediv__"):
+ # combine with bools does not generate the correct result
+ # (numpy behaviour for div is to regard the bools as numeric)
+ expected = s.astype(float).combine(other, op)
+ if op_name == "__rpow__":
+ # for rpow, combine does not propagate NaN
+ expected[result.isna()] = np.nan
+ self.assert_series_equal(result, expected)
+ else:
+ with pytest.raises(exc):
+ op(s, other)
+
+ def _check_divmod_op(self, s, op, other, exc=None):
+ # override to not raise an error
+ super()._check_divmod_op(s, op, other, None)
+
+ @pytest.mark.skip(reason="BooleanArray does not error on ops")
+ def test_error(self, data, all_arithmetic_operators):
+ # other specific errors tested in the boolean array specific tests
+ pass
+
+
+class TestComparisonOps(base.BaseComparisonOpsTests):
+ def check_opname(self, s, op_name, other, exc=None):
+ # overwriting to indicate ops don't raise an error
+ super().check_opname(s, op_name, other, exc=None)
+
+ def _compare_other(self, s, data, op_name, other):
+ self.check_opname(s, op_name, other)
+
+
+class TestReshaping(base.BaseReshapingTests):
+ pass
+
+
+class TestMethods(base.BaseMethodsTests):
+ @pytest.mark.parametrize("na_sentinel", [-1, -2])
+ def test_factorize(self, data_for_grouping, na_sentinel):
+ # override because we only have 2 unique values
+ labels, uniques = pd.factorize(data_for_grouping, na_sentinel=na_sentinel)
+ expected_labels = np.array(
+ [0, 0, na_sentinel, na_sentinel, 1, 1, 0], dtype=np.intp
+ )
+ expected_uniques = data_for_grouping.take([0, 4])
+
+ tm.assert_numpy_array_equal(labels, expected_labels)
+ self.assert_extension_array_equal(uniques, expected_uniques)
+
+ def test_combine_le(self, data_repeated):
+ # override because expected needs to be boolean instead of bool dtype
+ orig_data1, orig_data2 = data_repeated(2)
+ s1 = pd.Series(orig_data1)
+ s2 = pd.Series(orig_data2)
+ result = s1.combine(s2, lambda x1, x2: x1 <= x2)
+ expected = pd.Series(
+ [a <= b for (a, b) in zip(list(orig_data1), list(orig_data2))],
+ dtype="boolean",
+ )
+ self.assert_series_equal(result, expected)
+
+ val = s1.iloc[0]
+ result = s1.combine(val, lambda x1, x2: x1 <= x2)
+ expected = pd.Series([a <= val for a in list(orig_data1)], dtype="boolean")
+ self.assert_series_equal(result, expected)
+
+ def test_searchsorted(self, data_for_sorting, as_series):
+ # override because we only have 2 unique values
+ data_for_sorting = pd.array([True, False], dtype="boolean")
+ b, a = data_for_sorting
+ arr = type(data_for_sorting)._from_sequence([a, b])
+
+ if as_series:
+ arr = pd.Series(arr)
+ assert arr.searchsorted(a) == 0
+ assert arr.searchsorted(a, side="right") == 1
+
+ assert arr.searchsorted(b) == 1
+ assert arr.searchsorted(b, side="right") == 2
+
+ result = arr.searchsorted(arr.take([0, 1]))
+ expected = np.array([0, 1], dtype=np.intp)
+
+ tm.assert_numpy_array_equal(result, expected)
+
+ # sorter
+ sorter = np.array([1, 0])
+ assert data_for_sorting.searchsorted(a, sorter=sorter) == 0
+
+
+class TestCasting(base.BaseCastingTests):
+ pass
+
+
+class TestGroupby(base.BaseGroupbyTests):
+ """
+ Groupby-specific tests are overridden because boolean only has 2
+ unique values, base tests uses 3 groups.
+ """
+
+ def test_grouping_grouper(self, data_for_grouping):
+ df = pd.DataFrame(
+ {"A": ["B", "B", None, None, "A", "A", "B"], "B": data_for_grouping}
+ )
+ gr1 = df.groupby("A").grouper.groupings[0]
+ gr2 = df.groupby("B").grouper.groupings[0]
+
+ tm.assert_numpy_array_equal(gr1.grouper, df.A.values)
+ tm.assert_extension_array_equal(gr2.grouper, data_for_grouping)
+
+ @pytest.mark.parametrize("as_index", [True, False])
+ def test_groupby_extension_agg(self, as_index, data_for_grouping):
+ df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1], "B": data_for_grouping})
+ result = df.groupby("B", as_index=as_index).A.mean()
+ _, index = pd.factorize(data_for_grouping, sort=True)
+
+ index = pd.Index(index, name="B")
+ expected = pd.Series([3, 1], index=index, name="A")
+ if as_index:
+ self.assert_series_equal(result, expected)
+ else:
+ expected = expected.reset_index()
+ self.assert_frame_equal(result, expected)
+
+ def test_groupby_extension_no_sort(self, data_for_grouping):
+ df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1], "B": data_for_grouping})
+ result = df.groupby("B", sort=False).A.mean()
+ _, index = pd.factorize(data_for_grouping, sort=False)
+
+ index = pd.Index(index, name="B")
+ expected = pd.Series([1, 3], index=index, name="A")
+ self.assert_series_equal(result, expected)
+
+ def test_groupby_extension_transform(self, data_for_grouping):
+ valid = data_for_grouping[~data_for_grouping.isna()]
+ df = pd.DataFrame({"A": [1, 1, 3, 3, 1], "B": valid})
+
+ result = df.groupby("B").A.transform(len)
+ expected = pd.Series([3, 3, 2, 2, 3], name="A")
+
+ self.assert_series_equal(result, expected)
+
+ def test_groupby_extension_apply(self, data_for_grouping, groupby_apply_op):
+ df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1], "B": data_for_grouping})
+ df.groupby("B").apply(groupby_apply_op)
+ df.groupby("B").A.apply(groupby_apply_op)
+ df.groupby("A").apply(groupby_apply_op)
+ df.groupby("A").B.apply(groupby_apply_op)
+
+ def test_groupby_apply_identity(self, data_for_grouping):
+ df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1], "B": data_for_grouping})
+ result = df.groupby("A").B.apply(lambda x: x.array)
+ expected = pd.Series(
+ [
+ df.B.iloc[[0, 1, 6]].array,
+ df.B.iloc[[2, 3]].array,
+ df.B.iloc[[4, 5]].array,
+ ],
+ index=pd.Index([1, 2, 3], name="A"),
+ name="B",
+ )
+ self.assert_series_equal(result, expected)
+
+ def test_in_numeric_groupby(self, data_for_grouping):
+ df = pd.DataFrame(
+ {
+ "A": [1, 1, 2, 2, 3, 3, 1],
+ "B": data_for_grouping,
+ "C": [1, 1, 1, 1, 1, 1, 1],
+ }
+ )
+ result = df.groupby("A").sum().columns
+
+ if data_for_grouping.dtype._is_numeric:
+ expected = pd.Index(["B", "C"])
+ else:
+ expected = pd.Index(["C"])
+
+ tm.assert_index_equal(result, expected)
+
+
+class TestNumericReduce(base.BaseNumericReduceTests):
+ def check_reduce(self, s, op_name, skipna):
+ result = getattr(s, op_name)(skipna=skipna)
+ expected = getattr(s.astype("float64"), op_name)(skipna=skipna)
+ # override parent function to cast to bool for min/max
+ if op_name in ("min", "max") and not pd.isna(expected):
+ expected = bool(expected)
+ tm.assert_almost_equal(result, expected)
+
+
+class TestBooleanReduce(base.BaseBooleanReduceTests):
+ pass
+
+
+class TestPrinting(base.BasePrintingTests):
+ pass
+
+
+# TODO parsing not yet supported
+# class TestParsing(base.BaseParsingTests):
+# pass