Add trailing_nans parameter to control NaN placement in sort results by r-devulap · Pull Request #235 · numpy/x86-simd-sort

r-devulap · 2026-05-11T09:14:03Z

Summary of changes in `nan` branch

`trailing_nans` parameter

Added bool trailing_nans = true to all sort/select routines (qsort, qselect, partial_qsort, argsort, argselect). When hasnan=true, this controls where NaN values land in the result:

trailing_nans=true (default): NaNs at the end, regardless of sort direction
trailing_nans=false: NaNs at the beginning, regardless of sort direction

Previously the library replaced NaNs with ±inf as a workaround; now it partitions them directly to the correct end. Bit-exact NaN payloads are preserved.

`descending` parameter for `argselect`

Added bool descending = false to argselect, completing its API to match argsort and qselect. When descending=true:

The k-th element is the k-th largest
All elements before index k are ≥ it
SIMD path partitions at mirror position arrsize-1-k then reverses; scalar path selects a std::greater comparator

Tests

Added trailing/leading NaN coverage for qsort, argsort, and qselect
Replaced the single test_argselect with four tests: ascending, descending, trailing NaNs, leading NaNs
New utils/custom-compare.h with NaN-aware comparators shared across tests

All sorting and selection routines (qsort, qselect, partial_qsort, argsort, argselect) now accept an optional `bool trailing_nans` parameter (default: true). When hasnan=true: - trailing_nans=true (default): NaN values appear at the end of the result, independent of sort direction. - trailing_nans=false: NaN values appear at the beginning of the result, independent of sort direction. Previously, NaN placement was coupled to the sort direction: ascending always put NaN at the end, descending always put NaN at the beginning. This change decouples the two. Implementation notes: - qsort path: replace_inf_with_nan now takes both `descending` and `trailing_nans`; when they differ, the real-value portion is rotated before NaN slots are written back. - qselect path: NaN is moved to the desired end before partitioning, keeping the existing move_nans_to_{end,start}_of_array helpers. - argsort path: after std_argsort_withnan + optional reverse, a std::rotate moves NaN indices to the desired end if needed. - argselect path: comparator lambda respects trailing_nans. - keyvalue_qsort/select path: descending reverse now only reverses the real portion [0, index_last_elem], leaving trailing NaN in place. - Scalar fallback: four NaN-aware comparators (compare_nan_end/begin, compare_arg_nan_end/begin) added to utils/custom-compare.h. Tests updated to use compare_nan_end for descending reference sorts and to skip value-comparison checks when arr[k] is NaN.

seberg · 2026-05-11T11:15:15Z

Nice, I have to take a slightly closer look. But should it be nans_largest=true or so (I dunno what the best name is). Because I think with the current choice passing descending=true, has_nan=true would default to trailing_nans=true and thus change the result?
(Just a very minor translation of where you use the direct value vs. the "if they differ".)

MaanasArora · 2026-05-11T12:34:52Z

Thanks, this is nice! Going to take a deeper look soon - but, I was wondering if it's actually easier to just change replace_nan_with_inf when possible, to use negative (versus positive) infinity based on trailing_nans? Rotation seems a bit roundabout and could have some overhead, I suppose. Though not entirely familiar with how possible that is (though I did try it locally it and it seemed to work).

r-devulap · 2026-05-25T03:55:34Z

Thanks, this is nice! Going to take a deeper look soon - but, I was wondering if it's actually easier to just change replace_nan_with_inf when possible, to use negative (versus positive) infinity based on trailing_nans? Rotation seems a bit roundabout and could have some overhead, I suppose. Though not entirely familiar with how possible that is (though I did try it locally it and it seemed to work).

It might easier to just reuse existing helper functions: move_nans_to_end_of_array and move_nans_to_start_of_array to move the nan's at the start and then only sort the remaining elements. I will make that change.

- Add bool descending = false parameter to all argselect APIs and internal implementations (avx512, avx2, scalar, static, dispatch) - In xss_argselect: reverse the index array after partitioning when descending is true; pass descending to std_argselect_withnan for NaN-array fast path - std_argselect_withnan: honour descending flag in the comparator - scalar argselect: pick comparator once based on descending and trailing_nans, then call nth_element once - Tests: replace test_argselect with four typed tests covering ascending, descending, trailing NaNs and leading NaNs - IS_ARG_PARTITIONED: add descending parameter forwarded to IS_ARR_PARTITIONED

argselect_ always partitions in ascending order. For descending, we must select at position arrsize-1-k so the k-th largest lands at that mirror index; reversing then moves it to position k with the correct left/right partition invariant.

r-devulap · 2026-05-25T07:48:48Z

@seberg @MaanasArora This looks good to me - happy to merge once you sign off.

r-devulap · 2026-05-25T07:57:10Z

Nice, I have to take a slightly closer look. But should it be nans_largest=true or so (I dunno what the best name is). Because I think with the current choice passing descending=true, has_nan=true would default to trailing_nans=true and thus change the result?

NumPy, as of today, does not use a descending flag. I’m not sure which other libraries rely on this flag, and I’m not particularly concerned about that change. The newly introduced trailing_nans flag handles NaN placement independently of sort order—it simply lets users specify where NaNs should appear, regardless of whether the sort is ascending or descending. Personally, I don’t find assigning an order to NaNs meaningful. In that sense, this flag mainly provides users with a clear preference for where they want their NaNs.

seberg · 2026-05-27T09:25:09Z

Sounds good, I prefer the preference to be "at the end" or "at the front", rather than "NaNs are considered larger/smaller".
(IMO, the larger/smaller convention mostly makes sense if you think of it in terms of implementing a comparison function.)

MaanasArora · 2026-05-29T06:06:40Z

Thanks a lot, I had a brief look and this makes sense to me! I like this way of moving nans to the start/end and sorting the rest; within the nans we don't have an order so hopefully this approach could even have performance benefits, may be nice to run the benchmarks!

Agreed that thinking of NaNs as smaller/larger is confusing in most cases, so position is better. Would trailing_nans be a bit confusing still, compared to nans_last for example? Just a nit sorry, probably fine, given the docs...

seberg

Thanks, LGTM. The one ask that I think would be nice is to add all-NaN tests since codex has a point that the logic around it works out but it is a bit fishy.
The compare_arg_* feels like it could be refactored for pretty code, but that doesn't seem too important.

You seem to change the approach for partitioning from NaN replacement to NaN moving. That seems sensible to me, I am just a bit curious what that means speed-wise (not vital either to me, we may also see it in NumPy when bumping)?

seberg · 2026-05-29T08:19:30Z

Accidental commit, please remove.

seberg · 2026-05-29T08:50:35Z

    if (UNLIKELY(hasnan)) {
-        if (descending) {
+        if (!trailing_nans) {
            index_first_elem = move_nans_to_start_of_array(arr, arrsize);


Codex noticed that index_first_elem can overflow to MAX_SIZE_T here and in similar paths. That seems to work out just fine with qselect_ overflowing to consider the range as empty.
(Since this is an index and not a size, -1/MAX_SIZE_T is effectively invalid after all.)

It might be clearer to have inclusive indexes, but honestly, I think it is OK, it would be nice to add an "all_nan" test case.

seberg · 2026-05-29T09:00:00Z

+            else {
+                cmp = compare_arg_nan_begin<T, std::less<T>>(arr);
+            }
+        }


Could be nice to put this into a get_argcmp_func to avoid duplicating this if that works.
Might be possible to refactor it more so that the arg comparisons use the get_cmp_func() to avoid the compare_arg_* family, but I don't think it's vital.

r-devulap mentioned this pull request May 11, 2026

Changing meaning (or adding another flag) for has_nan=True, descending=True #234

Open

seberg mentioned this pull request May 11, 2026

ENH: Descending sorts with built-in dtype implementations numpy/numpy#31345

Merged

seberg mentioned this pull request May 13, 2026

Short-term follow-ups for descending sorts numpy/numpy#31423

Open

r-devulap added 9 commits May 25, 2026 04:01

qsort: partition NaNs to end/start instead of replacing with inf

be7c2da

argsort: fold descending+trailing_nans logic into std_argsort_withnan

8d476a5

tests: add trailing/leading NaN coverage for qsort, argsort, qselect

6265c9a

Clang format files

e9f6e49

docs: update NaN/trailing_nans docs in README files

edce532

Clang format files

81bfcff

Unused parameter

8c4cc0d

Update argselect docs to reflect new descending parameter

030227a

seberg reviewed May 29, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add trailing_nans parameter to control NaN placement in sort results#235

Add trailing_nans parameter to control NaN placement in sort results#235
r-devulap wants to merge 11 commits into
numpy:mainfrom
r-devulap:nan

r-devulap commented May 11, 2026 •

edited

Loading

Uh oh!

seberg commented May 11, 2026

Uh oh!

MaanasArora commented May 11, 2026 •

edited

Loading

Uh oh!

r-devulap commented May 25, 2026

Uh oh!

r-devulap commented May 25, 2026

Uh oh!

r-devulap commented May 25, 2026

Uh oh!

seberg commented May 27, 2026

Uh oh!

MaanasArora commented May 29, 2026 •

edited

Loading

Uh oh!

seberg left a comment

Uh oh!

seberg May 29, 2026

Uh oh!

seberg May 29, 2026 •

edited

Loading

Uh oh!

seberg May 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

r-devulap commented May 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary of changes in nan branch

trailing_nans parameter

descending parameter for argselect

Tests

Uh oh!

seberg commented May 11, 2026

Uh oh!

MaanasArora commented May 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

r-devulap commented May 25, 2026

Uh oh!

r-devulap commented May 25, 2026

Uh oh!

r-devulap commented May 25, 2026

Uh oh!

seberg commented May 27, 2026

Uh oh!

MaanasArora commented May 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

seberg left a comment

Choose a reason for hiding this comment

Uh oh!

seberg May 29, 2026

Choose a reason for hiding this comment

Uh oh!

seberg May 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

seberg May 29, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

r-devulap commented May 11, 2026 •

edited

Loading

Summary of changes in `nan` branch

`trailing_nans` parameter

`descending` parameter for `argselect`

MaanasArora commented May 11, 2026 •

edited

Loading

MaanasArora commented May 29, 2026 •

edited

Loading

seberg May 29, 2026 •

edited

Loading