perf: Add fast paths for series.arg_sort and dataframe.sort

[siddharth-vi]

Add fast paths for series.arg_sort and dataframe.sort(single column) in the following cases -

1. Key column has no nulls and is sorted in the required order
2. Key column has no nulls and is sorted in the opposite order
3. Key column has nulls ,is sorted in the required order and has nulls in the correct place

For the second case, if maintain_order is false, technically we can simply return the reverse of 0..len. However current behavior is that we still do a stable sort on it, maintaining order of elements. This is also the expectation in some unrelated test cases. To maintain this behavior, I have implemented a linear algorithm to do a stable reverse sort. This is faster than the current implementation which is also linear for sorting already sorted data, but takes more time due to creating a copy of the data.

Closes #19364

Benchmarks

dataframe.sort

Code

import numpy as np
import polars as pl
n = 1_000_000
df = pl.DataFrame({
    "a": np.random.rand(n),
    "b": np.random.rand(n),
})

# -- One column
print("A unsorted")
df2 = df.select("a")
%timeit df2.sort("a")

print("A sorted")
df2 = df.select("a").sort("a")
%timeit df2.sort("a")

# -- Two columns
print("A unsorted, B unsorted")
df2 = df.clone()
%timeit df2.sort("a")

print("A sorted, B unsorted")
df2 = df.sort("a")
%timeit df2.sort("a")

print("A sorted reverse, B unsorted")
df2 = df.sort("a",descending=True)
%timeit df2.sort("a")


df_nulls = pl.DataFrame({
    "a": np.random.rand(n).tolist()+[None],
    "b": np.random.rand(n).tolist()+[None],
})
print("A sorted, with null")
df2_nulls = df_nulls.sort("a")
%timeit df2_nulls.sort("a")

case	Time taken Before	Time taken After
A unsorted	32.7 ms ± 651 μs	32 ms ± 1.33 ms
A sorted	11.6 ms ± 52.6 μs	6.08 μs ± 24.2 ns
A sorted reverse	17.6 ms ± 1.43 ms	4.75 ms ± 244 μs
A sorted, with null	14.1 ms ± 560 μs	6.14 μs ± 43.5 ns

series.arg_sort

Code

import numpy as np
#import polars as pl
n = 1_000_000
df = pl.DataFrame({
    "a": np.random.rand(n),
    "b": np.random.rand(n),
})
s=df["a"]
%timeit s.arg_sort()

s=df.sort("a")["a"]
%timeit s.arg_sort()
%timeit s.arg_sort(descending=True)
df_nulls = pl.DataFrame({
    "a": np.random.rand(n).tolist()+[None],
    "b": np.random.rand(n).tolist()+[None],
})
s_nulls=df_nulls.sort("a")["a"]
%timeit s_nulls.arg_sort()

case	Time taken Before	Time taken After
A unsorted	27.2 ms ± 715 μs	27.5 ms ± 472 μs
A sorted	9.85 ms ± 51.2 μs	71.2 μs ± 362 ns
A sorted reverse	12.7 ms ± 298 μs	1.99 ms ± 9.96 μs
A sorted, with null	11 ms ± 67.7 μs	69.8 μs ± 744 ns

Testing

We need to ensure that even if we take fast path we do not change the final array. I have modified some pre existing tests to also test for fast path. I have added an additional test which checks for correctness of sorting.

[codecov]

Codecov Report

Attention: Patch coverage is 89.78495% with 19 lines in your changes missing coverage. Please review.

Project coverage is 79.61%. Comparing base (c92612a) to head (c38636c).
Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
...polars-core/src/chunked_array/ops/sort/arg_sort.rs	86.52%	19 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #19872      +/-   ##
==========================================
- Coverage   79.62%   79.61%   -0.01%     
==========================================
  Files        1564     1564              
  Lines      217989   218171     +182     
  Branches     2477     2477              
==========================================
+ Hits       173564   173704     +140     
- Misses      43857    43899      +42     
  Partials      568      568              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[siddharth-vi]

Disabled fast path for categorical columns due to #19900

[ritchie46]

Thanks, I have left some comments.

We should split the bug fix PR from the optimizations, they should not be in a single PR.

[ritchie46]

I think we can simplify here somewhat after we replace with std::cmp::min

[siddharth-vi]

Have rewritten this, let me know if it looks fine now.

[ritchie46]

We need to use iter as into_iter boxes.

I also don't like flattened iterators. If we can write this with explicitly looping over the chunks we should prefer that.

[siddharth-vi]

Have rewritten this, let me know if it looks fine now.

[ritchie46]

What does this return? An empty Vec?

[siddharth-vi]

Yes, for arrays of length zero we return an empty vec (From Vec::with_capacity(0)).

[siddharth-vi]

Have created a separate PR for the bug fix- #20004

[siddharth-vi]

@ritchie46 I have made the requested changes, please have a look .

[ritchie46]

Thanks a lot @siddharth-vi. Great improvements. Left one comment, but I think we can do that in a later PR.

[ritchie46]

No blocker for this PR, but I rather see this in a generic function. Could be a follow up.