PERF: efficient argmax/argmin for SparseArray

[GYHHAHA]

• partially closes ENH: implement efficient sorting methods for SparseArray (argsort/argmin/argmax) #34197
• Tests added and passed if fixing a bug or adding a new feature
• All code checks passed.
• Added type annotations to new arguments/methods/functions.
• Added an entry in the latest doc/source/whatsnew/vX.X.X.rst file if fixing a bug or adding a new feature.

Thanks for the review. Currently, only argmax/argmin are implemented since argsort has many annoying corner cases and I have to spent more time on the correctness and take the follow-up in a separate PR. Simple benchmark:

>>> val = np.random.rand(1000000)
>>> mask = val < 0.99
>>> val[mask] = np.nan
>>> arr = SparseArray(val)
>>> %timeit arr.argmax()
8.6 ms ± 343 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) <- master
44.3 µs ± 934 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each) <- this pr

[pep8speaks]

Hello @GYHHAHA! 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 2022-07-27 16:40:06 UTC

[mzeitlin11]

Thanks for the pr @GYHHAHA! Some comments below, general question is why can't we just dispatch to a SparseArray.argmin/argmax which directly uses np.argmin/max instead of having to modify the general ExtensionArray path (for example, how SparseArray handles min/max. This could also avoid materializing the data.

[GYHHAHA]

@mzeitlin11 I think you are right. Originally I want to reuse some general codes. It seems to be unnecessary. I will handle this in SparseArray.

[GYHHAHA]

I find the current implementation for _first_fill_value_loc is wrong. A simple bad case:

>>> arr = SparseArray([np.nan, 1, 0, 0, np.nan, 2], fill_value=1)
>>> arr._first_fill_value_loc()
5 # should be 1

I will commit a fix after merging this one.

[mzeitlin11]

I find the current implementation for _first_fill_value_loc is wrong. A simple bad case:

>>> arr = SparseArray([np.nan, 1, 0, 0, np.nan, 2], fill_value=1)
>>> arr._first_fill_value_loc()
5 # should be 1

I will commit a fix after merging this one.

I'm not entirely sure what that function is meant to do, but the name might also just be misleading. It looks to only be used in unique as some kind of mechanism for figuring out where to insert the fill value (probably to keep sorted order in the result?).

[GYHHAHA]

The reason why unique need to find the loc is to keep the insertion place right, which aligns the normal array.unique() result. @mzeitlin11

[mzeitlin11]

@GYHHAHA what's the status here? This generally lgtm - but does the _first_fill_value_loc issues you discovered cause any regressions in cases where _first_fill_value_loc is wrong? Also just confirming benchmarks still look good with new impl?

[GYHHAHA]

How about I submit the fix PR for unique and _first_fill_value_loc (and add tests) now? After the fix merged, then we can deal with the current one. @mzeitlin11

[mzeitlin11]

How about I submit the fix PR for unique and _first_fill_value_loc (and add tests) now? After the fix merged, then we can deal with the current one. @mzeitlin11

Sounds great, thanks!

[mzeitlin11]

Thanks @GYHHAHA! LGTM pending one small comment. Can you please also show latest benchmarks for the cases you showed in the pr description?

[GYHHAHA]

@mzeitlin11 New benchmark is added in the description.

[mroeschke]

Thanks @GYHHAHA