ENH: allow storing ExtensionArrays in Index

[jbrockmendel]

• closes Add interface for defining an ExtensionIndex #22861
• whatsnew

Still TODO

• 28815 not-yet-passing tests (with --skip-slow --skip-db)
• flesh out tests/extension/base/ea_index.py
• deprecation cycle (decided we can probably just change)
• optimize EAIndexEngine/NullableIndexEngine once BUG: get_indexer_non_unique with np.datetime64("NaT") and np.timedelta64("NaT")  #43870 is resolved
• EA.putmask implementation ATM is just a stub (and needs tests)

cc @jorisvandenbossche @TomAugspurger

xref #39133

[jbrockmendel]

Woops, meant to make this a Draft PR. Is there a way to convert it?

[mzeitlin11]

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 :)

[pep8speaks]

Hello @jbrockmendel! Thanks for updating this PR. We checked the lines you've touched for PEP 8 issues, and found:

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

Comment last updated at 2021-12-29 22:43:08 UTC

[jorisvandenbossche]

Nice work!

Can you give a bit more high-level context on what is now implemented and how you approached it?
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?
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).

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? But I don't know if that would work in practice of course (eg for an initial PR to store EA in Index, 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).

[jorisvandenbossche]

Might be best to extract data/mask from the MaskedArray values and index those inside the loop?

[jbrockmendel]

yah, i want to see if i can do this in a way that allows for sharing code with ExtensionEngine without a big perf hit.

[jbrockmendel]

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.

[jbrockmendel]

Updated, whatsnew added, SparseArray behavior deprecated.

[jbrockmendel]

gentle ping; the Index[SparseArray] deprecation isn't the most important, but it'd be nice to do it right.

[jreback]

looks good, just a small question.

[jreback]

prob can move the todo's to another issue and close the original #39133

[jreback]

thanks @jbrockmendel glad to get this in

[jorisvandenbossche]

@jbrockmendel did you already open a PR with the NullableEngine / ExtensionEngine to review that part?

[jbrockmendel]

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.