Bug: Save original index and remap after function completes

[Jeffrharr]

• closes PERF: Surprisingly slow nlargest with duplicates in the index #55767
• 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.

Note: I'm new to this project, so this is my first PR.

Saves the index for SeriesNLargest at algorithm start and resets it before returning. This fixes performance issues when the index has many duplicate values.

Results:

• The original statistics can be viewed in the original ticket, but slow_df was several ms. Note that the results in the ticket were not from my development machine and specific timings differ.

In [4]: import pandas as pd
   ...: import numpy as np
   ...: 
   ...: N = 1500
   ...: N_HALF = 750
   ...: 
   ...: slow_df = pd.DataFrame({'a':  np.random.rand(N)}, index=np.concatenate([[1] * N_HALF, np.arange(N_HALF)]))
   ...: print("slow_df")
   ...: %timeit slow_df['a'].nlargest()
   ...: 
   ...: fast_df = pd.DataFrame({'a': np.random.rand(N)})
   ...: print("fast_df")
   ...: %timeit fast_df['a'].nlargest()

slow_df
427 μs ± 11.4 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
fast_df
420 μs ± 5.4 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Tests

The existing tests should cover this unless we want to add specific tests via the asv_bench.

Addendum

I also modified the call to sort to use sort(kind="stable") to get consistent ordering which is what is currently happening in the equivalent Frame method (it was using kind=mergesort which is equivalent to kind=stable, but kept for portability). I can remove this -- it may be better in another PR.
https://numpy.org/doc/stable/reference/generated/numpy.sort.html#numpy.sort

[rhshadrach]

Thanks for the PR! I'm seeing a performance regression when the index does not contain duplicates. Can we do this conditionally.

[Jeffrharr]

Thanks for the PR! I'm seeing a performance regression when the index does not contain duplicates. Can we do this conditionally.

I appreciate the review!

Can you confirm the performance regression numbers that you are seeing? Note that the timings in the original ticket were from a different machine than the one I'm working on.

ASV Benchmarks appear similar enough to be reasonable in the standard case.

-- Original Function --

[37.66%] ··· series_methods.NSort.time_nlargest                              ok
[37.66%] ··· ======= ==========
               keep            
             ------- ----------
              first   3.59±0ms 
               last   3.70±0ms 
               all    3.84±0ms 
             ======= ==========

[37.71%] ··· series_methods.NSort.time_nsmallest                             ok
[37.71%] ··· ======= ==========
               keep            
             ------- ----------
              first   3.26±0ms 
               last   2.94±0ms 
               all    3.31±0ms 
             ======= ==========

-- Function with my changes --

[37.66%] ··· series_methods.NSort.time_nlargest                              ok
[37.66%] ··· ======= ==========
               keep            
             ------- ----------
              first   3.71±0ms 
               last   3.41±0ms 
               all    3.60±0ms 
             ======= ==========

[37.71%] ··· series_methods.NSort.time_nsmallest                             ok
[37.71%] ··· ======= ==========
               keep            
             ------- ----------
              first   2.86±0ms 
               last   2.58±0ms 
               all    3.20±0ms 
             ======= ==========

[Jeffrharr]

Thanks for the PR! I'm seeing a performance regression when the index does not contain duplicates. Can we do this conditionally.

Hey, can I get a re-review on this?

[rhshadrach]

Thanks for the ping, will get to it in the next day or two.

[rhshadrach]

lgtm

[mroeschke]

Suggested change

- :meth:`Series.nlargest` uses a 'stable' sort internally and will preserve original ordering.

[rhshadrach]

@mroeschke - there is a change in Series.nlargest here where previously we weren't using a stable sorting algorithm in certain cases. Now the only dtype I think this can possibly effect is object (as otherwise equal implies identical), and even there it is quite uncommon to have an impact. Perhaps the line needs to make this more clear that it is unlikely to have any impact, or are you thinking it isn't even worth mentioning?

[mroeschke]

Ah gotcha. OK doesnt hurt to include it in the whatsnew notes

[mroeschke]

Thanks @Jeffrharr