ExtensionArray.argsort and nargsort

[TomAugspurger]

47ffcd6 changes how we handle sorting EAs. We've made nargsort compatible with extension arrays. It now calls ExtensionArray._values_for_argsort and then applies pandas nargsort sorting algorithm to the values.

This may have broken the (implied) ability for ExtensionArray's to "bring their own" sorting algorithm. For example, you can imagine a GPU-backed EA that has a very fast sorting algorithm available. While ExtensionArray.argsort can use that custom sorting algo, pandas may not. Anywhere that uses nargsort(obj) will now essentially do nargsort(thing._values_for_argsort()) for EAs.

[jorisvandenbossche]

My current thinking is that nargsort should dispatch to EA.argsort() instead of using np.argsort if an EA is passed in. And then nargsort can further deal with ascending / descending and position of NaNs (something that EA.argsort does not handle).

Of course, that means that EA.argsort will call nargsort on its "values to be sorted".
So something like nargsort(EA) -> EA.argsort() -> nargsort(EA._values_for_argsort) with default settings -> the result of this is further processed in nargsort for final result of nargsort(EA)

So we should maybe see how this works for the missing values (eg for the EA, only for the non-missing values EA.argsort is called? or do we skip the separation of the values in missing / non-missing as we do for np.argsort for EAs?)

[TomAugspurger]

Is this still relevant now that NumPy has the array_function protocol? If an array implements that get to use their own sorting algo.

[TomAugspurger]

Pushing off 1.0, though this issue is a bit unclear to me now.

[jbrockmendel]

This may have broken the (implied) ability for ExtensionArray's to "bring their own" sorting algorithm

So something like nargsort(EA) -> EA.argsort() -> nargsort(EA._values_for_argsort) with default settings -> the result of this is further processed in nargsort for final result of nargsort(EA)

Agreed on both counts. A recent attempt to change nargsort/EA.argsort to work this way broke the world. Have either of you tried doing this?

[jbrockmendel]

@jorisvandenbossche have you looked at this recently? IIUC this is a blocker for #43840. You seem to have the strongest opinions w/r/t to argsort (xref #33941, #33942, #37815) and i don't fully grok them, would probably be made clearer with an implementation.

[jbrockmendel]

nargsort now dispatches to an EA's argsort method, so I think this is closable?

[mroeschke]

Yeah it appears this is covered in the EA interface so closing