ENH: allow storing ExtensionArrays in Index
State: closed

- [x] closes #22861
- [x] whatsnew

Still TODO
- [x] 
288
15 not-yet-passing tests (with --skip-slow --skip-db)
- [ ] flesh out tests/extension/base/ea_index.py
- [x] 
deprecation cycle
 (decided we can probably just change)
- [ ] optimize EAIndexEngine/NullableIndexEngine once #43870 is resolved
- [x] EA.putmask implementation ATM is just a stub (and needs tests)

cc @jorisvandenbossche @TomAugspurger 

xref https://github.com/pandas-dev/pandas/issues/39133

Labels:
name: ExtensionArray
description: Extending pandas with custom dtypes or arrays.
name: Index
description: Related to the Index class or subclasses

Comments:
Woops, meant to make this a Draft PR.  Is there a way to convert it?
> Woops, meant to make this a Draft PR. Is there a way to convert it?

Under the reviewers section, there should be a `Convert to draft` option (have gone ahead and clicked it :)
Hello @jbrockmendel! Thanks for updating this PR. We checked the lines you've touched for [PEP 8](https://www.python.org/dev/peps/pep-0008) issues, and found:











































































































There are currently no PEP 8 issues detected in this Pull Request. Cheers! :beers: 

##### Comment last updated at 2021-12-29 22:43:08 UTC
> The general ExtensionEngine seems to work with an actual ExtensionArray. An alternative could be to have it work an an ndarray that the EA could provide? Although a potential disadvantage of that approach is then that such an ndarray needs to be materialized always (in case this EA -> ndarray conversion is costly), while the the current way doesn't need that (but also cannot make use of existing optimized engines).

I chose to keep the EA intact instead of casting to ndarray bc most cases (for get_loc, the main concern) can use EA methods (searchsorted, `__eq__`) directly.

I suspect many cases will be able to use NDArrayBackedExtensionArray, in which case using one of the non-EA engines would be nice.  I haven't implemented a way of doing that.

> For example, for the NullableEngine, you are currently not using any hash table. Did you look at that / decide that is not possible or desirable? Or is that a potential future improvement, and you are focusing first on getting it working with a base implementation?

Right, first trying to get everything working, then will look at optimizations.  (Also for NullableEngine.get_loc at least I have a different optimization in mind I want to try first).

> For reviewability: I suppose that in theory some of the changes in Index class are mostly for allowing to store an EA in the Index, somewhat independent of the engine changes, and thus could be done separately?

Yep, a bunch of my recent PRs have been exactly that.  More coming up, e.g. `eq_NA_compat` fixes problems with Index[object] containing pd.NA (though the function needs to be re-written) so i'll break that off before long.  Also the float16 check in FloatingArray and the isna check in testing.pyx.

> the nullable dtypes could use an object dtype array for the engine (if that's not buggy with NA), and that could then also work for starting the Index implementation and tests).

ATM the NullableEngine isn't a pain point.  The remaining test failures are mostly in setops (xref #44000) and value_counts ordering. 

Giving some thought to how to implement tests for this.  Suppose hypothetically that we decided that the extant tests with the `index` and `index_with_missing` fixtures are sufficient, so we just need to give EA authors a way to shim their indexes into those fixtures.  The options that come to mind are 1) tell authors to directly edit pandas.conftest.indices_dict and 2) collect tests using the relevant fixtures into classes that the authors can import into their test files and override the fixtures there.  2) seems like The Right Way To Do it, but also really tedious.  I'm open to other ideas.
> Giving some thought to how to implement tests for this.

Just a quick comment without looking in much detail. For EAs in general (and thus for _external_ EAs), I think we can just use the existing "base" extension tests and expand those with some basic indexing tests (I see you added `/tests/extension/base/ea_index.py` with a start for this). This won't be as comprehensive as our own indexing tests (using the `indices_dict` from conftest), but I would personally not start recommending external authors to use this as well, and keep it on the base extension tests. 
And for all our internal EAs, we of course have the more comprehensive testing with our own fixtures. 
>> the nullable dtypes could use an object dtype array for the engine (if that's not buggy with NA), and that could then also work for starting the Index implementation and tests).
>
> ATM the NullableEngine isn't a pain point. The remaining test failures are mostly in setops (xref #44000) and value_counts ordering.

Coming back to this, it might still be useful to do I think (if not complicating things for testing). While the NullableEngine might not be a pain point for getting things to work here, it still _is_ 500 lines of new code that needs to be reviwed (so half of the diff in the current PR, although it's probably the easier half to review). Doing this separate will still make the review process easier I think? (and thus faster to get the first part merged) What do you think?

I did a quick local check with the following diff (on top on this PR), basically disabling the use of the two new engines:

```patch
--- a/pandas/core/indexes/base.py
+++ b/pandas/core/indexes/base.py
@@ -841,12 +841,12 @@ class Index(IndexOpsMixin, PandasObject):
         # For base class (object dtype) we get ObjectEngine
 
         if isinstance(self._values, BaseMaskedArray):
-            return libindex.NullableEngine(self._values)
+            return libindex.ObjectEngine(self._values.astype(object))
         elif (
             isinstance(self._values, ExtensionArray)
             and self._engine_type is libindex.ObjectEngine
         ):
-            return libindex.ExtensionEngine(self._values)
+            return libindex.ObjectEngine(self._values.astype(object))
 
         engine_type = cast("type[libindex.IndexEngine]", self._engine_type)
         # to avoid a reference cycle, bind `target_values` to a local variable, so
@@ -3197,10 +3197,10 @@ class Index(IndexOpsMixin, PandasObject):
         name = get_op_result_name(self, other)
         if isinstance(result, Index):
             if result.name != name:
-                return result.rename(name)
-            return result
+                result = result.rename(name)
         else:
-            return self._shallow_copy(result, name=name)
+            result = self._shallow_copy(result, name=name)
+        return result.astype(self.dtype, copy=False)
 
     # TODO: standardize return type of non-union setops type(self vs other)
     @final
@@ -4764,6 +4764,8 @@ class Index(IndexOpsMixin, PandasObject):
         """
         # error: Incompatible return value type (got "Union[ExtensionArray,
         # ndarray]", expected "ndarray")
+        if type(self) is Index and isinstance(self._values, ExtensionArray):
+            return self._values.astype(object)
         return self._values  # type: ignore[return-value]
 
     def _from_join_target(self, result: np.ndarray) -> ArrayLike:

```

and running the tests locally they all pass (the change in the middle is `_wrap_setop_result` to get the set op tests passing).
> Coming back to this, it might still be useful to do I think (if not complicating things for testing). While the NullableEngine might not be a pain point for getting things to work here, it still is 500 lines of new code that needs to be reviwed (so half of the diff in the current PR, although it's probably the easier half to review). Doing this separate will still make the review process easier I think? (and thus faster to get the first part merged) What do you think?

I'm a bit concerned about subtle behavior changes that may come from casting to ndarray[object], but I think we can definitely trim some of the diff by using ExtensionEngine for now and leaving NullableEngine to a follow-on PR.
Thanks for the updates. I hope you don't mind, but I pushed a small commit (but touching many files) to run the initial base extension test you add for the EAs (so subclassing that base class everywhere).
Updated, whatsnew added, SparseArray behavior deprecated.
gentle ping; the Index[SparseArray] deprecation isn't the most important, but it'd be nice to do it right.
prob can move the todo's to another issue and close the original #39133 
thanks @jbrockmendel glad to get this in
@jbrockmendel did you already open a PR with the NullableEngine / ExtensionEngine to review that part? 
> did you already open a PR with the NullableEngine / ExtensionEngine to review that part?

No.  I've got a branch near-ready, but there is a "values_for_argsort" usage that i want to get rid of before pushing.