ENH: Add extended sorting APIs

charris · charris · commit fe16eac02af2 · 2025-09-02T11:59:28.000-06:00
This PR adds new sorting APIs that can be used for extending our current
sort types. The new C-APIs do not have the *which* argument of the
legacy C-APIs, instead they have *flags* request flag API. In this
transition, heapsort can no longer be called directly, introsort
(quicksort) is called instead. The sort/argsort methods get new keyword
arguments used to set the flags when *kind* is not given, so *flags* is
invisible at the Python level. The new keywords are:

- stable -- whether the sort should stable, default is False,
- descending -- whether to sort in descending order, default is False,
- nanfirst -- whether NaNs sort to the last or first, default is False.

The main difference here is that the keywords select by functionality,
whereas *kind* selects by algorithm. Note that descending and nanfirst
are not yet implemented, so an error is raised if either is specified.

The added C-API functions are:

- (API)PyArray_SortEx,
- (API)PyArray_ArgSortEx,

Those are the two functions that need changes depending on how we want
to implement adding new sort algorithms.  They should be pretty easy to
adapt to whatever we choose to do.

The legacy C-API functions, PyArray_Sort and PyArray_ArgSort, have been
modified to call through the new C-API functions. In order to do so
the *kind* function has been encoded in the *flags* like so:

- "quicksort" - default (0)
- "heapsort"  - default (0)
- "mergesort" - stable  (1)

I note that the partitioning/selection functions have not been modified,
we may want to do that a some point.
diff --git a/doc/release/upcoming_changes/29642.c_api.rst b/doc/release/upcoming_changes/29642.c_api.rst
@@ -0,0 +1,14 @@
+The NPY_SORTKIND enum has been enhanced with new variables
+----------------------------------------------------------
+This is of interest if you are using ``PyArray_Sort`` or ``PyArray_ArgSort``.
+We have changed the semantics of the old names in the NPY_SORTKIND enum and
+added new ones. The changes are backward compatible, and no recompilation is
+needed. The new names of interest are:
+
+* NPY_SORT_DEFAULT -- default sort (same value as NPY_QUICKSORT)
+* NPY_SORT_STABLE  -- the sort must be stable (same value as NPY_MERGESORT)
+* NPY_SORT_DESCENDING -- the sort must be descending
+* NPY_SORT_NANFIRST -- NaNs sort to the beginning
+
+The semantic change is that NPY_HEAPSORT is mapped to NPY_QUICKSORT when used.
+
diff --git a/doc/release/upcoming_changes/29642.change.rst b/doc/release/upcoming_changes/29642.change.rst
@@ -0,0 +1,4 @@
+Sorting ``kind='heapsort'`` now maps to ``kind='quicksort'``
+------------------------------------------------------------
+It is unlikely that this change will be noticed, but if you do see a change in
+execution time or unstable argsort order, this is likely the cause.
diff --git a/doc/source/reference/c-api/array.rst b/doc/source/reference/c-api/array.rst
@@ -2308,9 +2308,9 @@ Item selection and manipulation
     is sorted using the algorithm denoted by *kind*, which is an integer/enum pointing
     to the type of sorting algorithms used.
 
-.. c:function:: PyObject* PyArray_ArgSort(PyArrayObject* self, int axis)
+.. c:function:: PyObject* PyArray_ArgSort(PyArrayObject* self, int axis, NPY_SORTKIND kind)
 
-    Equivalent to :meth:`ndarray.argsort<numpy.ndarray.argsort>` (*self*, *axis*).
+    Equivalent to :meth:`ndarray.argsort<numpy.ndarray.argsort>` (*self*, *axis*, *kind*).
     Return an array of indices such that selection of these indices
     along the given ``axis`` would return a sorted version of *self*. If *self* ->descr
     is a data-type with fields defined, then self->descr->names is used
@@ -4321,7 +4321,11 @@ Enumerated Types
 .. c:enum:: NPY_SORTKIND
 
     A special variable-type which can take on different values to indicate
-    the sorting algorithm being used.
+    the sorting algorithm being used. These algorithm types have not been
+    treated strictly for some time, but rather treated as stable/not stable.
+    In NumPy 2.4 they are replaced by requirements (see below), but done in a
+    backwards compatible way. These values will continue to work, except that
+    that NPY_HEAPSORT will do the same thing as NPY_QUICKSORT.
 
     .. c:enumerator:: NPY_QUICKSORT
 
@@ -4335,11 +4339,41 @@ Enumerated Types
 
     .. c:enumerator:: NPY_NSORTS
 
-       Defined to be the number of sorts. It is fixed at three by the need for
-       backwards compatibility, and consequently :c:data:`NPY_MERGESORT` and
-       :c:data:`NPY_STABLESORT` are aliased to each other and may refer to one
-       of several stable sorting algorithms depending on the data type.
+        Defined to be the number of sorts. It is fixed at three by the need for
+        backwards compatibility, and consequently :c:data:`NPY_MERGESORT` and
+        :c:data:`NPY_STABLESORT` are aliased to each other and may refer to one
+        of several stable sorting algorithms depending on the data type.
 
+    In NumPy 2.4 the algorithm names are replaced by requirements. You can still use
+    the old values, a recompile is not needed, but they are reinterpreted such that
+
+    * NPY_QUICKSORT and NPY_HEAPSORT -> NPY_SORT_DEFAULT
+    * NPY_MERGESORT and NPY_STABLE -> NPY_SORT_STABLE
+
+    .. c:enumerator:: NPY_SORT_DEFAULT
+
+        The default sort for the type. For the NumPy builtin types it may be
+        stable or not, but will be ascending and sort NaN types to the end. It
+        is usually chosen for speed and/or low memory.
+
+    .. c:enumerator:: NPY_SORT_STABLE
+
+        (Requirement) Specifies that the sort must be stable.
+
+    .. c:enumerator:: NPY_SORT_DESCENDING
+
+        (Requirement) Specifies that the sort must be in descending order.
+        This functionality is not yet implemented for any of the NumPy types
+        and cannot yet be set from the Python interface.
+
+    .. c:enumerator:: NPY_SORT_NANFIRST
+
+        (Requirement) Specifies that the sort must sort NaNs to
+        the beginning, e.g. NaNs compare less than any other value for the
+        type. This is only relevant if the type has NaN equivalents.
+        This functionality is not yet implemented for any of the NumPy types
+        and cannot yet be set from the Python interface. It may change in the
+        future.
 
 .. c:enum:: NPY_SCALARKIND
 
diff --git a/doc/source/reference/c-api/types-and-structures.rst b/doc/source/reference/c-api/types-and-structures.rst
@@ -650,7 +650,7 @@ PyArray_ArrFuncs
         address is given. The final argument is the array which is
         needed to get the itemsize for variable-length arrays.
 
-    .. c:member:: int sort(void* start, npy_intp length, void* arr)
+    .. c:member:: int sort(void* start, npy_intp length, void *arr)
 
         An array of function pointers to a particular sorting
         algorithms. A particular sorting algorithm is obtained using a
@@ -659,7 +659,7 @@ PyArray_ArrFuncs
         in-place assuming contiguous and aligned data.
 
     .. c:member:: int argsort( \
-            void* start, npy_intp* result, npy_intp length, void *arr)
+            void* start, npy_intp* result, npy_intp length, void* arr)
 
         An array of function pointers to sorting algorithms for this
         data type. The same sorting algorithms as for sort are
diff --git a/numpy/_core/include/numpy/ndarraytypes.h b/numpy/_core/include/numpy/ndarraytypes.h
@@ -162,18 +162,39 @@ enum NPY_TYPECHAR {
 };
 
 /*
- * Changing this may break Numpy API compatibility
- * due to changing offsets in PyArray_ArrFuncs, so be
- * careful. Here we have reused the mergesort slot for
- * any kind of stable sort, the actual implementation will
- * depend on the data type.
+ * Changing this may break Numpy API compatibility due to changing offsets in
+ * PyArray_ArrFuncs, so be careful. Here we have reused the mergesort slot for
+ * any kind of stable sort, the actual implementation will depend on the data
+ * type.
+ *
+ * Updated in NumPy 2.4
+ *
+ * Updated with new names denoting requirements rather than the algorithm. All
+ * the previous values are reused in a way that should be downstream
+ * compatible, but the actual algorithms used may be different than before. The
+ * new approach should be more flexible and easier to update. The idea is that
+ * NPY_SORT_STABLE | NPY_SORT_DESCENDING | NPY_SORT_NANFIRST should provide an
+ * index.
+ * 
+ * Names with a leading underscore are private, and should only be used
+ * internally by NumPy.
+ *
+ * NPY_NSORTS remains the same for backwards compatibility, it should not be
+ * changed.
  */
+
 typedef enum {
-        _NPY_SORT_UNDEFINED=-1,
-        NPY_QUICKSORT=0,
-        NPY_HEAPSORT=1,
-        NPY_MERGESORT=2,
-        NPY_STABLESORT=2,
+        _NPY_SORT_UNDEFINED = -1,
+        NPY_QUICKSORT = 0,
+        NPY_HEAPSORT = 1,
+        NPY_MERGESORT = 2,
+        NPY_STABLESORT = 2,
+        // new style names 
+        _NPY_SORT_HEAPSORT = 1,
+        NPY_SORT_DEFAULT = 0,
+        NPY_SORT_STABLE = 2,
+        NPY_SORT_DESCENDING = 4,
+        NPY_SORT_NANFIRST = 8,
 } NPY_SORTKIND;
 #define NPY_NSORTS (NPY_STABLESORT + 1)
 
diff --git a/numpy/_core/src/multiarray/_multiarray_tests.c.src b/numpy/_core/src/multiarray/_multiarray_tests.c.src
@@ -2089,10 +2089,18 @@ run_sortkind_converter(PyObject* NPY_UNUSED(self), PyObject *args)
         return NULL;
     }
     switch (kind) {
-        case _NPY_SORT_UNDEFINED: return PyUnicode_FromString("_NPY_SORT_UNDEFINED");
-        case NPY_QUICKSORT: return PyUnicode_FromString("NPY_QUICKSORT");
-        case NPY_HEAPSORT: return PyUnicode_FromString("NPY_HEAPSORT");
-        case NPY_STABLESORT: return PyUnicode_FromString("NPY_STABLESORT");
+        case _NPY_SORT_UNDEFINED:
+            return PyUnicode_FromString("_NPY_SORT_UNDEFINED");
+        case NPY_QUICKSORT:
+            return PyUnicode_FromString("NPY_QUICKSORT");
+        case NPY_HEAPSORT:
+            return PyUnicode_FromString("NPY_HEAPSORT");
+        case NPY_STABLESORT:
+            return PyUnicode_FromString("NPY_STABLESORT");
+        default:
+            // the other possible values in NPY_SORTKIND can only
+            // be set with keywords.
+            break;
     }
     return PyLong_FromLong(kind);
 }
diff --git a/numpy/_core/src/multiarray/item_selection.c b/numpy/_core/src/multiarray/item_selection.c
diff --git a/numpy/_core/src/multiarray/methods.c b/numpy/_core/src/multiarray/methods.c
diff --git a/numpy/_core/tests/test_multiarray.py b/numpy/_core/tests/test_multiarray.py
