ENH: Index[bool] (#45061)

jbrockmendel · web-flow · commit ee6b0a09fff7 · 2022-02-05T14:01:21.000-05:00
diff --git a/doc/source/whatsnew/v1.5.0.rst b/doc/source/whatsnew/v1.5.0.rst
@@ -37,6 +37,7 @@ Other enhancements
 - :meth:`to_numeric` now preserves float64 arrays when downcasting would generate values not representable in float32 (:issue:`43693`)
 - :meth:`Series.reset_index` and :meth:`DataFrame.reset_index` now support the argument ``allow_duplicates`` (:issue:`44410`)
 - :meth:`.GroupBy.min` and :meth:`.GroupBy.max` now supports `Numba <https://numba.pydata.org/>`_ execution with the ``engine`` keyword (:issue:`45428`)
+- Implemented a ``bool``-dtype :class:`Index`, passing a bool-dtype arraylike to ``pd.Index`` will now retain ``bool`` dtype instead of casting to ``object`` (:issue:`45061`)
 -
 
 .. ---------------------------------------------------------------------------
diff --git a/pandas/_libs/index.pyi b/pandas/_libs/index.pyi
@@ -42,6 +42,7 @@ class ObjectEngine(IndexEngine): ...
 class DatetimeEngine(Int64Engine): ...
 class TimedeltaEngine(DatetimeEngine): ...
 class PeriodEngine(Int64Engine): ...
+class BoolEngine(UInt8Engine): ...
 
 class BaseMultiIndexCodesEngine:
     levels: list[np.ndarray]
diff --git a/pandas/_libs/index.pyx b/pandas/_libs/index.pyx
@@ -802,6 +802,13 @@ cdef class BaseMultiIndexCodesEngine:
 include "index_class_helper.pxi"
 
 
+cdef class BoolEngine(UInt8Engine):
+    cdef _check_type(self, object val):
+        if not util.is_bool_object(val):
+            raise KeyError(val)
+        return <uint8_t>val
+
+
 @cython.internal
 @cython.freelist(32)
 cdef class SharedEngine:
diff --git a/pandas/conftest.py b/pandas/conftest.py
@@ -555,7 +555,8 @@ def _create_mi_with_dt64tz_level():
     "num_uint8": tm.makeNumericIndex(100, dtype="uint8"),
     "num_float64": tm.makeNumericIndex(100, dtype="float64"),
     "num_float32": tm.makeNumericIndex(100, dtype="float32"),
-    "bool": tm.makeBoolIndex(10),
+    "bool-object": tm.makeBoolIndex(10).astype(object),
+    "bool-dtype": Index(np.random.randn(10) < 0),
     "categorical": tm.makeCategoricalIndex(100),
     "interval": tm.makeIntervalIndex(100),
     "empty": Index([]),
@@ -630,7 +631,7 @@ def index_flat_unique(request):
         key
         for key in indices_dict
         if not (
-            key in ["int", "uint", "range", "empty", "repeats"]
+            key in ["int", "uint", "range", "empty", "repeats", "bool-dtype"]
             or key.startswith("num_")
         )
         and not isinstance(indices_dict[key], MultiIndex)
diff --git a/pandas/core/algorithms.py b/pandas/core/algorithms.py
@@ -220,9 +220,6 @@ def _reconstruct_data(
     elif is_bool_dtype(dtype):
         values = values.astype(dtype, copy=False)
 
-        # we only support object dtypes bool Index
-        if isinstance(original, ABCIndex):
-            values = values.astype(object, copy=False)
     elif dtype is not None:
         if is_datetime64_dtype(dtype):
             dtype = np.dtype("datetime64[ns]")
@@ -830,7 +827,10 @@ def value_counts(
     -------
     Series
     """
-    from pandas.core.series import Series
+    from pandas import (
+        Index,
+        Series,
+    )
 
     name = getattr(values, "name", None)
 
@@ -868,7 +868,13 @@ def value_counts(
         else:
             keys, counts = value_counts_arraylike(values, dropna)
 
-            result = Series(counts, index=keys, name=name)
+            # For backwards compatibility, we let Index do its normal type
+            #  inference, _except_ for if if infers from object to bool.
+            idx = Index._with_infer(keys)
+            if idx.dtype == bool and keys.dtype == object:
+                idx = idx.astype(object)
+
+            result = Series(counts, index=idx, name=name)
 
     if sort:
         result = result.sort_values(ascending=ascending)
diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py
@@ -505,6 +505,10 @@ def __new__(
             if data.dtype.kind in ["i", "u", "f"]:
                 # maybe coerce to a sub-class
                 arr = data
+            elif data.dtype.kind == "b":
+                # No special subclass, and Index._ensure_array won't do this
+                #  for us.
+                arr = np.asarray(data)
             else:
                 arr = com.asarray_tuplesafe(data, dtype=_dtype_obj)
 
@@ -702,7 +706,7 @@ def _with_infer(cls, *args, **kwargs):
             # "Union[ExtensionArray, ndarray[Any, Any]]"; expected
             # "ndarray[Any, Any]"
             values = lib.maybe_convert_objects(result._values)  # type: ignore[arg-type]
-            if values.dtype.kind in ["i", "u", "f"]:
+            if values.dtype.kind in ["i", "u", "f", "b"]:
                 return Index(values, name=result.name)
 
         return result
@@ -872,9 +876,12 @@ def _engine(
         ):
             return libindex.ExtensionEngine(target_values)
 
+        target_values = cast(np.ndarray, target_values)
         # to avoid a reference cycle, bind `target_values` to a local variable, so
         # `self` is not passed into the lambda.
-        target_values = cast(np.ndarray, target_values)
+        if target_values.dtype == bool:
+            return libindex.BoolEngine(target_values)
+
         # error: Argument 1 to "ExtensionEngine" has incompatible type
         # "ndarray[Any, Any]"; expected "ExtensionArray"
         return self._engine_type(target_values)  # type:ignore[arg-type]
@@ -2680,7 +2687,6 @@ def _is_all_dates(self) -> bool:
         """
         Whether or not the index values only consist of dates.
         """
-
         if needs_i8_conversion(self.dtype):
             return True
         elif self.dtype != _dtype_obj:
@@ -7302,7 +7308,7 @@ def _maybe_cast_data_without_dtype(
             FutureWarning,
             stacklevel=3,
         )
-    if result.dtype.kind in ["b", "c"]:
+    if result.dtype.kind in ["c"]:
         return subarr
     result = ensure_wrapped_if_datetimelike(result)
     return result
diff --git a/pandas/core/tools/datetimes.py b/pandas/core/tools/datetimes.py
@@ -1076,6 +1076,8 @@ def to_datetime(
             result = convert_listlike(arg, format)
     else:
         result = convert_listlike(np.array([arg]), format)[0]
+        if isinstance(arg, bool) and isinstance(result, np.bool_):
+            result = bool(result)  # TODO: avoid this kludge.
 
     #  error: Incompatible return value type (got "Union[Timestamp, NaTType,
     # Series, Index]", expected "Union[DatetimeIndex, Series, float, str,
diff --git a/pandas/core/util/hashing.py b/pandas/core/util/hashing.py
@@ -319,7 +319,7 @@ def _hash_ndarray(
 
     # First, turn whatever array this is into unsigned 64-bit ints, if we can
     # manage it.
-    elif isinstance(dtype, bool):
+    elif dtype == bool:
         vals = vals.astype("u8")
     elif issubclass(dtype.type, (np.datetime64, np.timedelta64)):
         vals = vals.view("i8").astype("u8", copy=False)
diff --git a/pandas/tests/base/test_value_counts.py b/pandas/tests/base/test_value_counts.py
@@ -284,7 +284,7 @@ def test_value_counts_with_nan(dropna, index_or_series):
     obj = klass(values)
     res = obj.value_counts(dropna=dropna)
     if dropna is True:
-        expected = Series([1], index=[True])
+        expected = Series([1], index=Index([True], dtype=obj.dtype))
     else:
         expected = Series([1, 1, 1], index=[True, pd.NA, np.nan])
     tm.assert_series_equal(res, expected)
diff --git a/pandas/tests/indexes/common.py b/pandas/tests/indexes/common.py
@@ -216,6 +216,8 @@ def test_ensure_copied_data(self, index):
             # RangeIndex cannot be initialized from data
             # MultiIndex and CategoricalIndex are tested separately
             return
+        elif index.dtype == object and index.inferred_type == "boolean":
+            init_kwargs["dtype"] = index.dtype
 
         index_type = type(index)
         result = index_type(index.values, copy=True, **init_kwargs)
@@ -522,6 +524,9 @@ def test_fillna(self, index):
         # GH 11343
         if len(index) == 0:
             return
+        elif index.dtype == bool:
+            # can't hold NAs
+            return
         elif isinstance(index, NumericIndex) and is_integer_dtype(index.dtype):
             return
         elif isinstance(index, MultiIndex):
diff --git a/pandas/tests/indexes/multi/test_indexing.py b/pandas/tests/indexes/multi/test_indexing.py
@@ -621,13 +621,22 @@ def test_get_loc_implicit_cast(self, level, dtypes):
         idx = MultiIndex.from_product(levels)
         assert idx.get_loc(tuple(key)) == 3
 
-    def test_get_loc_cast_bool(self):
-        # GH 19086 : int is casted to bool, but not vice-versa
-        levels = [[False, True], np.arange(2, dtype="int64")]
+    @pytest.mark.parametrize("dtype", [bool, object])
+    def test_get_loc_cast_bool(self, dtype):
+        # GH 19086 : int is casted to bool, but not vice-versa (for object dtype)
+        #  With bool dtype, we don't cast in either direction.
+        levels = [Index([False, True], dtype=dtype), np.arange(2, dtype="int64")]
         idx = MultiIndex.from_product(levels)
 
-        assert idx.get_loc((0, 1)) == 1
-        assert idx.get_loc((1, 0)) == 2
+        if dtype is bool:
+            with pytest.raises(KeyError, match=r"^\(0, 1\)$"):
+                assert idx.get_loc((0, 1)) == 1
+            with pytest.raises(KeyError, match=r"^\(1, 0\)$"):
+                assert idx.get_loc((1, 0)) == 2
+        else:
+            # We use python object comparisons, which treat 0 == False and 1 == True
+            assert idx.get_loc((0, 1)) == 1
+            assert idx.get_loc((1, 0)) == 2
 
         with pytest.raises(KeyError, match=r"^\(False, True\)$"):
             idx.get_loc((False, True))
diff --git a/pandas/tests/indexes/test_any_index.py b/pandas/tests/indexes/test_any_index.py
@@ -49,6 +49,10 @@ def test_mutability(index):
 def test_map_identity_mapping(index):
     # GH#12766
     result = index.map(lambda x: x)
+    if index.dtype == object and result.dtype == bool:
+        assert (index == result).all()
+        # TODO: could work that into the 'exact="equiv"'?
+        return  # FIXME: doesn't belong in this file anymore!
     tm.assert_index_equal(result, index, exact="equiv")
 
 
diff --git a/pandas/tests/indexes/test_base.py b/pandas/tests/indexes/test_base.py
@@ -321,15 +321,21 @@ def test_view_with_args(self, index):
             "unicode",
             "string",
             pytest.param("categorical", marks=pytest.mark.xfail(reason="gh-25464")),
-            "bool",
+            "bool-object",
+            "bool-dtype",
             "empty",
         ],
         indirect=True,
     )
     def test_view_with_args_object_array_raises(self, index):
-        msg = "Cannot change data-type for object array"
-        with pytest.raises(TypeError, match=msg):
-            index.view("i8")
+        if index.dtype == bool:
+            msg = "When changing to a larger dtype"
+            with pytest.raises(ValueError, match=msg):
+                index.view("i8")
+        else:
+            msg = "Cannot change data-type for object array"
+            with pytest.raises(TypeError, match=msg):
+                index.view("i8")
 
     @pytest.mark.parametrize("index", ["int", "range"], indirect=True)
     def test_astype(self, index):
@@ -397,9 +403,9 @@ def test_is_(self):
 
     def test_asof_numeric_vs_bool_raises(self):
         left = Index([1, 2, 3])
-        right = Index([True, False])
+        right = Index([True, False], dtype=object)
 
-        msg = "Cannot compare dtypes int64 and object"
+        msg = "Cannot compare dtypes int64 and bool"
         with pytest.raises(TypeError, match=msg):
             left.asof(right[0])
         # TODO: should right.asof(left[0]) also raise?
@@ -591,7 +597,8 @@ def test_append_empty_preserve_name(self, name, expected):
         "index, expected",
         [
             ("string", False),
-            ("bool", False),
+            ("bool-object", False),
+            ("bool-dtype", False),
             ("categorical", False),
             ("int", True),
             ("datetime", False),
@@ -606,7 +613,8 @@ def test_is_numeric(self, index, expected):
         "index, expected",
         [
             ("string", True),
-            ("bool", True),
+            ("bool-object", True),
+            ("bool-dtype", False),
             ("categorical", False),
             ("int", False),
             ("datetime", False),
@@ -621,7 +629,8 @@ def test_is_object(self, index, expected):
         "index, expected",
         [
             ("string", False),
-            ("bool", False),
+            ("bool-object", False),
+            ("bool-dtype", False),
             ("categorical", False),
             ("int", False),
             ("datetime", True),
diff --git a/pandas/tests/indexes/test_common.py b/pandas/tests/indexes/test_common.py
@@ -88,7 +88,9 @@ def test_constructor_non_hashable_name(self, index_flat):
 
     def test_constructor_unwraps_index(self, index_flat):
         a = index_flat
-        b = type(a)(a)
+        # Passing dtype is necessary for Index([True, False], dtype=object)
+        #  case.
+        b = type(a)(a, dtype=a.dtype)
         tm.assert_equal(a._data, b._data)
 
     def test_to_flat_index(self, index_flat):
@@ -426,6 +428,9 @@ def test_hasnans_isnans(self, index_flat):
             return
         elif isinstance(index, NumericIndex) and is_integer_dtype(index.dtype):
             return
+        elif index.dtype == bool:
+            # values[1] = np.nan below casts to True!
+            return
 
         values[1] = np.nan
 
diff --git a/pandas/tests/indexes/test_index_new.py b/pandas/tests/indexes/test_index_new.py
@@ -74,7 +74,7 @@ def test_constructor_dtypes_to_object(self, cast_index, vals):
             index = Index(vals)
 
         assert type(index) is Index
-        assert index.dtype == object
+        assert index.dtype == bool
 
     def test_constructor_categorical_to_object(self):
         # GH#32167 Categorical data and dtype=object should return object-dtype
diff --git a/pandas/tests/indexes/test_numpy_compat.py b/pandas/tests/indexes/test_numpy_compat.py
@@ -68,16 +68,18 @@ def test_numpy_ufuncs_basic(index, func):
         with tm.external_error_raised((TypeError, AttributeError)):
             with np.errstate(all="ignore"):
                 func(index)
-    elif isinstance(index, NumericIndex) or (
-        not isinstance(index.dtype, np.dtype) and index.dtype._is_numeric
+    elif (
+        isinstance(index, NumericIndex)
+        or (not isinstance(index.dtype, np.dtype) and index.dtype._is_numeric)
+        or index.dtype == bool
     ):
         # coerces to float (e.g. np.sin)
         with np.errstate(all="ignore"):
             result = func(index)
             exp = Index(func(index.values), name=index.name)
 
         tm.assert_index_equal(result, exp)
-        if type(index) is not Index:
+        if type(index) is not Index or index.dtype == bool:
             # i.e NumericIndex
             assert isinstance(result, Float64Index)
         else:
@@ -117,8 +119,10 @@ def test_numpy_ufuncs_other(index, func):
         with tm.external_error_raised(TypeError):
             func(index)
 
-    elif isinstance(index, NumericIndex) or (
-        not isinstance(index.dtype, np.dtype) and index.dtype._is_numeric
+    elif (
+        isinstance(index, NumericIndex)
+        or (not isinstance(index.dtype, np.dtype) and index.dtype._is_numeric)
+        or index.dtype == bool
     ):
         # Results in bool array
         result = func(index)
diff --git a/pandas/tests/indexes/test_setops.py b/pandas/tests/indexes/test_setops.py
@@ -9,7 +9,6 @@
 import pytest
 
 from pandas.core.dtypes.cast import find_common_type
-from pandas.core.dtypes.common import is_dtype_equal
 
 from pandas import (
     CategoricalIndex,
@@ -55,14 +54,20 @@ def test_union_different_types(index_flat, index_flat2, request):
 
     if (
         not idx1.is_unique
+        and not idx2.is_unique
+        and not idx2.is_monotonic_decreasing
         and idx1.dtype.kind == "i"
-        and is_dtype_equal(idx2.dtype, "boolean")
+        and idx2.dtype.kind == "b"
     ) or (
         not idx2.is_unique
+        and not idx1.is_unique
+        and not idx1.is_monotonic_decreasing
         and idx2.dtype.kind == "i"
-        and is_dtype_equal(idx1.dtype, "boolean")
+        and idx1.dtype.kind == "b"
     ):
-        mark = pytest.mark.xfail(reason="GH#44000 True==1", raises=ValueError)
+        mark = pytest.mark.xfail(
+            reason="GH#44000 True==1", raises=ValueError, strict=False
+        )
         request.node.add_marker(mark)
 
     common_dtype = find_common_type([idx1.dtype, idx2.dtype])
@@ -231,7 +236,11 @@ def test_union_base(self, index):
     def test_difference_base(self, sort, index):
         first = index[2:]
         second = index[:4]
-        if isinstance(index, CategoricalIndex) or index.is_boolean():
+        if index.is_boolean():
+            # i think (TODO: be sure) there assumptions baked in about
+            #  the index fixture that don't hold here?
+            answer = set(first).difference(set(second))
+        elif isinstance(index, CategoricalIndex):
             answer = []
         else:
             answer = index[4:]
diff --git a/pandas/tests/reshape/concat/test_append_common.py b/pandas/tests/reshape/concat/test_append_common.py
diff --git a/pandas/tests/series/indexing/test_setitem.py b/pandas/tests/series/indexing/test_setitem.py
diff --git a/pandas/tests/series/methods/test_drop.py b/pandas/tests/series/methods/test_drop.py
diff --git a/pandas/tests/series/methods/test_value_counts.py b/pandas/tests/series/methods/test_value_counts.py
diff --git a/pandas/tests/series/test_logical_ops.py b/pandas/tests/series/test_logical_ops.py
diff --git a/pandas/tests/test_algos.py b/pandas/tests/test_algos.py

Original file line number	Diff line number	Diff line change
`@@ -37,6 +37,7 @@ Other enhancements`
`37`	`37`	- :meth:`to_numeric` now preserves float64 arrays when downcasting would generate values not representable in float32 (:issue:`43693`)
`38`	`38`	- :meth:`Series.reset_index` and :meth:`DataFrame.reset_index` now support the argument ``allow_duplicates`` (:issue:`44410`)
`39`	`39`	- :meth:`.GroupBy.min` and :meth:`.GroupBy.max` now supports `Numba <https://numba.pydata.org/>`_ execution with the ``engine`` keyword (:issue:`45428`)
	`40`	+- Implemented a ``bool``-dtype :class:`Index`, passing a bool-dtype arraylike to ``pd.Index`` will now retain ``bool`` dtype instead of casting to ``object`` (:issue:`45061`)
`40`	`41`	`-`
`41`	`42`
`42`	`43`	`.. ---------------------------------------------------------------------------`