numpy · MaanasArora · Mar 14, 2025 · Mar 26, 2025 · Mar 26, 2025 · Mar 26, 2025
diff --git a/doc/release/upcoming_changes/28516.c_api.rst b/doc/release/upcoming_changes/28516.c_api.rst
@@ -0,0 +1,17 @@
+New comparison and null handling enums for sorting in dtype API
+---------------------------------------------------------------
+
+Using the new `NPY_DT_sort_compare` slot, user-defined dtypes can
+now specify how to compare elements during sorting operations.
+The sort compare function should return a member of the
+`NPY_COMPARE_RESULT` enum to indicate the result of the comparison,
+including support for unordered comparisons.
+
+The sorting context of type `PyArrayMethod_SortContext` is passed
+to the `NPY_DT_get_sort_function` and `NPY_DT_get_argsort_function`
+functions and contains a boolean `descending` flag and a
+`nan_position` of type `NPY_SORT_NAN_POSITION`, which can be used to
+control the behavior of sorting with respect to NaN values.
+Currently, sorts are always ascending and nulls are always sorted last,
+but this must be checked in the context passed to the sort function
+to allow for future features.
diff --git a/doc/release/upcoming_changes/28516.new_feature.rst b/doc/release/upcoming_changes/28516.new_feature.rst
@@ -0,0 +1,11 @@
+New sorting function slots `NPY_DT_get_sort_function`, `NPY_DT_get_argsort_function` for dtype API
+---------------------------------------------------------------------------------------------------
+
+User-defined dtypes can now provide specific sorting functions for use with NumPy's sort methods.
+The new slots `NPY_DT_get_sort_function` and `NPY_DT_get_argsort_function` should be functions that
+return function pointers implementing sorting functionality for the dtype, while considering the
+sort-kind and order. The old arrfunc slots ``NPY_DT_PyArray_ArrFuncs_sort`` and
+``NPY_DT_PyArray_ArrFuncs_argsort`` may be deprecated in the future.
+
+Additionally, the new `NPY_DT_sort_compare` slot can be used to provide a comparison function for
+sorting, which will replace the default comparison function for the dtype in sorting functions.
diff --git a/doc/source/reference/c-api/array.rst b/doc/source/reference/c-api/array.rst
@@ -1873,6 +1873,36 @@ described below.
    pointer. Currently this is used for zero-filling and clearing arrays storing
    embedded references.
 
+.. c:type:: int (PyArray_SortFuncWithContext)( \
+                 PyArrayMethod_SortContext *context, void *data, \
+                 npy_intp num, NpyAuxData *auxdata)
+
+    A function to sort a buffer of data. The *data* is a pointer to the
+    beginning of the contiguous buffer containing *num* elements. A function
+    of this type is returned by the `get_sort_function` function in the DType
+    slots, where *context* is passed in containing the descriptor for the
+    array. Returns 0 on success, -1 on failure.
+
+.. c:type:: int (PyArray_ArgSortFuncWithContext)( \
+                 PyArrayMethod_SortContext *context, void *data, \
+                 npy_intp *tosort, npy_intp num, NpyAuxData *auxdata)
+
+    A function to arg-sort a buffer of data. The *data* is a pointer to the
+    beginning of the buffer containing *num* elements. The *tosort* is a
+    pointer to an array of indices that will be filled in with the
+    indices of the sorted elements. A function of this type is returned by
+    the `get_argsort_function` function in the DType slots, where
+    *context* is passed in containing the descriptor for the array.
+    Returns 0 on success, -1 on failure.
+
+.. c:type:: NPY_COMPARE_RESULT (PyArray_SortCompareFunc) ( \
+                 const void *a, const void *b, PyArray_Descr *descr)
+
+    A function to compare two elements of an array for sorting. The *a* and *b*
+    pointers point to the elements to compare, and *descr* is the descriptor for
+    the array. Returns a value of type :c:type:`NPY_COMPARE_RESULT` indicating
+    the result of the comparison, including whether each element is unordered.
+
 API Functions and Typedefs
 ~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -3521,6 +3551,40 @@ member of ``PyArrayDTypeMeta_Spec`` struct.
    force newly created arrays to have a newly created descriptor
    instance, no matter what input descriptor is provided by a user.
 
+.. c:macro:: NPY_DT_get_sort_function
+
+.. c:type:: int *(PyArrayDTypeMeta_GetSortFunction)(PyArray_Descr *, \
+        npy_intp sort_kind, PyArray_SortFuncWithContext **out_sort, \
+        NpyAuxData **out_auxdata, NPY_ARRAYMETHOD_FLAGS *out_flags)
+
+    .. versionadded:: 2.4
+
+    If defined, sets a custom sorting function for the DType for each of
+    the sort kinds numpy implements. Currently, sorts are always descending
+    and always use nulls to the end, and this must be checked in the
+    implementation. Returns 0 on success.
+
+.. c:macro:: NPY_DT_get_argsort_function
+
+.. c:type:: int *(PyArrayDTypeMeta_GetArgSortFunction)(PyArray_Descr *, \
+        npy_intp sort_kind, PyArray_ArgSortFuncWithContext **out_argsort, \
+        NpyAuxData **out_auxdata, NPY_ARRAYMETHOD_FLAGS *out_flags)
+
+    .. versionadded:: 2.4
+
+    If defined, sets a custom argsorting function for the DType for each of
+    the sort kinds numpy implements. Currently, sorts are always descending
+    and always use nulls to the end, and this must be checked in the
+    implementation. Returns 0 on success.
+
+.. c:macro:: NPY_DT_sort_compare
+
+    .. versionadded:: 2.4
+
+    If defined, sets a custom comparison function for the DType for use in
+    sorting, which will replace `NPY_DT_PyArray_ArrFuncs_compare`. Implements
+    ``PyArray_CompareFunc``.
+
 PyArray_ArrFuncs slots
 ^^^^^^^^^^^^^^^^^^^^^^
 
@@ -3547,6 +3611,8 @@ DType API slots but for now we have exposed the legacy
 .. c:macro:: NPY_DT_PyArray_ArrFuncs_compare
 
    Computes a comparison for `numpy.sort`, implements ``PyArray_CompareFunc``.
+   If `NPY_DT_sort_compare` is defined, it will be used instead. This slot may
+   be deprecated in the future.
 
 .. c:macro:: NPY_DT_PyArray_ArrFuncs_argmax
 
@@ -3590,13 +3656,17 @@ DType API slots but for now we have exposed the legacy
 
    An array of PyArray_SortFunc of length ``NPY_NSORTS``. If set, allows
    defining custom sorting implementations for each of the sorting
-   algorithms numpy implements.
+   algorithms numpy implements. If `NPY_DT_get_sort_function` is
+   defined, it will be used instead. This slot may be deprecated in the
+   future.
 
 .. c:macro:: NPY_DT_PyArray_ArrFuncs_argsort
 
    An array of PyArray_ArgSortFunc of length ``NPY_NSORTS``. If set,
    allows defining custom argsorting implementations for each of the
-   sorting algorithms numpy implements.
+   sorting algorithms numpy implements. If `NPY_DT_get_argsort_function`
+   is defined, it will be used instead. This slot may be deprecated in
+   the future.
 
 Macros and Static Inline Functions
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -4340,6 +4410,40 @@ Enumerated Types
        :c:data:`NPY_STABLESORT` are aliased to each other and may refer to one
        of several stable sorting algorithms depending on the data type.
 
+.. c:enum:: NPY_SORT_NAN_POSITION
+
+   An enum used to indicate the position of NaN values in sorting.
+
+   .. c:enumerator:: NPY_SORT_NAN_TO_START
+
+      Indicates that NaN values should be sorted to the start.
+
+   .. c:enumerator:: NPY_SORT_NAN_TO_END
+
+      Indicates that NaN values should be sorted to the end.
+
+.. c:enum:: NPY_COMPARE_RESULT
+
+    An enum used to indicate the result of a comparison operation.
+    The unordered comparisons are used to indicate that the
+    comparison is not well-defined for one or both of the operands,
+    such as when comparing NaN values.
+
+    .. c:enumerator:: NPY_LESS
+
+    .. c:enumerator:: NPY_EQUAL
+
+    .. c:enumerator:: NPY_GREATER
+
+    .. c:enumerator:: NPY_UNORDERED_LEFT
+
+    .. c:enumerator:: NPY_UNORDERED_RIGHT
+
+    .. c:enumerator:: NPY_UNORDERED_BOTH
+
+    .. c:enumerator:: NPY_COMPARE_ERROR
+
+        Indicates that an error occurred during the comparison operation.
 
 .. 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
@@ -792,6 +792,41 @@ PyArrayMethod_Context and PyArrayMethod_Spec
       An array of slots for the method. Slot IDs must be one of the values
       below.
 
+.. _arraymethod-sort-context:
+
+PyArrayMethod_SortContext
+-------------------------
+
+.. c:type:: PyArrayMethod_SortContext
+
+   A struct used to provide context for sorting methods.
+
+   .. code-block:: c
+
+      typedef struct {
+          PyArray_Descr *descriptor;
+          PyArray_SortCompareFunc *compare;
+          npy_bool descending;
+          NPY_SORT_NAN_POSITION nan_position;
+      } PyArrayMethod_SortContext
+
+    .. c:member:: PyArray_Descr *descriptor
+
+       The descriptor for the data type being sorted.
+
+    .. c:member:: PyArray_SortCompareFunc *compare
+
+       A pointer to the comparison function used for sorting. This function
+       can be NULL if the sort is not based on a comparison function.
+
+    .. c:member:: npy_bool descending
+
+       A flag indicating whether the sort is descending.
+
+    .. c:member:: NPY_SORT_NAN_POSITION nan_position
+
+       The position of NaN values in the sort order.
+
 .. _dtypemeta:
 
 PyArray_DTypeMeta and PyArrayDTypeMeta_Spec

diff --git a/numpy/_core/include/numpy/dtype_api.h b/numpy/_core/include/numpy/dtype_api.h
@@ -368,6 +368,12 @@ typedef int (PyArrayMethod_PromoterFunction)(PyObject *ufunc,
 #define NPY_DT_get_fill_zero_loop 10
 #define NPY_DT_finalize_descr 11
 
+#if NPY_API_VERSION >= NPY_2_4_API_VERSION
+#define NPY_DT_get_sort_function 12
+#define NPY_DT_get_argsort_function 13
+#define NPY_DT_sort_compare 14
+#endif
+
 // These PyArray_ArrFunc slots will be deprecated and replaced eventually
 // getitem and setitem can be defined as a performance optimization;
 // by default the user dtypes call `legacy_getitem_using_DType` and
@@ -477,4 +483,45 @@ typedef PyArray_Descr *(PyArrayDTypeMeta_FinalizeDescriptor)(PyArray_Descr *dtyp
 typedef int(PyArrayDTypeMeta_SetItem)(PyArray_Descr *, PyObject *, char *);
 typedef PyObject *(PyArrayDTypeMeta_GetItem)(PyArray_Descr *, char *);
 
+typedef enum {
+        NPY_LESS = -1,
+        NPY_EQUAL = 0,
+        NPY_GREATER = 1,
+        NPY_UNORDERED_LEFT = 2,
+        NPY_UNORDERED_RIGHT = 3,
+        NPY_UNORDERED_BOTH = 4,
+        NPY_COMPARE_ERROR = 5,
+} NPY_COMPARE_RESULT;
+
+typedef struct PyArrayMethod_SortContext_tag PyArrayMethod_SortContext;
+
+typedef NPY_COMPARE_RESULT (PyArray_SortCompareFunc)(
+        const void *a, const void *b, PyArray_Descr *descr);
+
+typedef enum {
+    NPY_SORT_NAN_TO_START = 0,
+    NPY_SORT_NAN_TO_END = 1,
+} NPY_SORT_NAN_POSITION;
+
+struct PyArrayMethod_SortContext_tag {
+    PyArray_Descr *descriptor;
+    PyArray_SortCompareFunc *compare;
+    npy_bool descending;
+    NPY_SORT_NAN_POSITION nan_position;
+};
+
+typedef int (PyArray_SortFuncWithContext)(PyArrayMethod_SortContext *,
+                               void *, npy_intp,
+                               NpyAuxData *);
+typedef int (PyArray_ArgSortFuncWithContext)(PyArrayMethod_SortContext *, 
+                                  void *, npy_intp *, npy_intp, 
+                                  NpyAuxData *);
+
+typedef int *(PyArrayDTypeMeta_GetSortFunction)(PyArray_Descr *, 
+        npy_intp, PyArray_SortFuncWithContext **, NpyAuxData **,
+        NPY_ARRAYMETHOD_FLAGS *);
+typedef int *(PyArrayDTypeMeta_GetArgSortFunction)(PyArray_Descr *, 
+        npy_intp, PyArray_ArgSortFuncWithContext **, NpyAuxData **,
+        NPY_ARRAYMETHOD_FLAGS *);
+
 #define NPY_DT_MAX_ARRFUNCS_SLOT \ 
   NPY_NUM_DTYPE_PYARRAY_ARRFUNCS_SLOTS + _NPY_DT_ARRFUNCS_OFFSET 
 #define NPY_DT_MAX_ARRFUNCS_SLOT \ 
   NPY_NUM_DTYPE_PYARRAY_ARRFUNCS_SLOTS + _NPY_DT_ARRFUNCS_OFFSET 
 #endif  /* NUMPY_CORE_INCLUDE_NUMPY___DTYPE_API_H_ */
diff --git a/numpy/_core/include/numpy/numpyconfig.h b/numpy/_core/include/numpy/numpyconfig.h
@@ -84,6 +84,7 @@
 #define NPY_2_1_API_VERSION 0x00000013
 #define NPY_2_2_API_VERSION 0x00000013
 #define NPY_2_3_API_VERSION 0x00000014
+#define NPY_2_4_API_VERSION 0x00000015
 
 
 /*

diff --git a/numpy/_core/src/common/npy_sort.h.src b/numpy/_core/src/common/npy_sort.h.src
@@ -5,6 +5,7 @@
 #include <Python.h>
 #include <numpy/npy_common.h>
 #include <numpy/ndarraytypes.h>
+#include <numpy/dtype_api.h>
 
 #define NPY_ENOMEM 1
 #define NPY_ECOMP 2
@@ -97,6 +98,22 @@ NPY_NO_EXPORT int atimsort_@suff@(void *vec, npy_intp *ind, npy_intp cnt, void *
  *****************************************************************************
  */
 
+NPY_NO_EXPORT int npy_quicksort_with_context(PyArrayMethod_SortContext *context, void *vec,
+    npy_intp cnt, NpyAuxData *auxdata);
+NPY_NO_EXPORT int npy_heapsort_with_context(PyArrayMethod_SortContext *context, void *vec,
+    npy_intp cnt, NpyAuxData *auxdata);
+NPY_NO_EXPORT int npy_mergesort_with_context(PyArrayMethod_SortContext *context, void *vec,
+    npy_intp cnt, NpyAuxData *auxdata);
+NPY_NO_EXPORT int npy_timsort_with_context(PyArrayMethod_SortContext *context, void *vec,
+    npy_intp cnt, NpyAuxData *auxdata);
+NPY_NO_EXPORT int npy_aquicksort_with_context(PyArrayMethod_SortContext *context, void *vec,
+    npy_intp *ind, npy_intp cnt, NpyAuxData *auxdata);
+NPY_NO_EXPORT int npy_aheapsort_with_context(PyArrayMethod_SortContext *context, void *vec,
+    npy_intp *ind, npy_intp cnt, NpyAuxData *auxdata);
+NPY_NO_EXPORT int npy_amergesort_with_context(PyArrayMethod_SortContext *context, void *vec,
+    npy_intp *ind, npy_intp cnt, NpyAuxData *auxdata);
+NPY_NO_EXPORT int npy_atimsort_with_context(PyArrayMethod_SortContext *context, void *vec,
+    npy_intp *ind, npy_intp cnt, NpyAuxData *auxdata);
 
 NPY_NO_EXPORT int npy_quicksort(void *vec, npy_intp cnt, void *arr);
 NPY_NO_EXPORT int npy_heapsort(void *vec, npy_intp cnt, void *arr);
@@ -107,6 +124,15 @@ NPY_NO_EXPORT int npy_aheapsort(void *vec, npy_intp *ind, npy_intp cnt, void *ar
 NPY_NO_EXPORT int npy_amergesort(void *vec, npy_intp *ind, npy_intp cnt, void *arr);
 NPY_NO_EXPORT int npy_atimsort(void *vec, npy_intp *ind, npy_intp cnt, void *arr);
 
+NPY_NO_EXPORT int npy_quicksort_impl(void *vec, npy_intp cnt, void *arr, PyArrayMethod_SortContext *context);
+NPY_NO_EXPORT int npy_heapsort_impl(void *vec, npy_intp cnt, void *arr, PyArrayMethod_SortContext *context);
+NPY_NO_EXPORT int npy_mergesort_impl(void *vec, npy_intp cnt, void *arr, PyArrayMethod_SortContext *context);
+NPY_NO_EXPORT int npy_timsort_impl(void *vec, npy_intp cnt, void *arr, PyArrayMethod_SortContext *context);
+NPY_NO_EXPORT int npy_aquicksort_impl(void *vec, npy_intp *ind, npy_intp cnt, void *arr, PyArrayMethod_SortContext *context);
+NPY_NO_EXPORT int npy_aheapsort_impl(void *vec, npy_intp *ind, npy_intp cnt, void *arr, PyArrayMethod_SortContext *context);
+NPY_NO_EXPORT int npy_amergesort_impl(void *vec, npy_intp *ind, npy_intp cnt, void *arr, PyArrayMethod_SortContext *context);
+NPY_NO_EXPORT int npy_atimsort_impl(void *vec, npy_intp *ind, npy_intp cnt, void *arr, PyArrayMethod_SortContext *context);
+
 #ifdef __cplusplus
 }
 #endif

diff --git a/numpy/_core/src/multiarray/dtypemeta.c b/numpy/_core/src/multiarray/dtypemeta.c
@@ -194,6 +194,9 @@ dtypemeta_initialize_struct_from_spec(
     NPY_DT_SLOTS(DType)->getitem = NULL;
     NPY_DT_SLOTS(DType)->get_clear_loop = NULL;
     NPY_DT_SLOTS(DType)->get_fill_zero_loop = NULL;
+    NPY_DT_SLOTS(DType)->get_sort_function = NULL;
+    NPY_DT_SLOTS(DType)->get_argsort_function = NULL;
+    NPY_DT_SLOTS(DType)->sort_compare = NULL;
     NPY_DT_SLOTS(DType)->finalize_descr = NULL;
     NPY_DT_SLOTS(DType)->f = default_funcs;
 
@@ -1230,6 +1233,15 @@ dtypemeta_wrap_legacy_descriptor(
         dtype_class->flags |= NPY_DT_NUMERIC;
     }
 
+    if (dt_slots->sort_compare == NULL) {
+        if (!NPY_DT_is_legacy(dtype_class) && !NPY_DT_is_user_defined(dtype_class)) {
+            PyErr_SetString(PyExc_RuntimeError,
+                    "DType has no sort_compare function.");
+            Py_DECREF(dtype_class);
+            return -1;
+        }
+    }
+
     if (_PyArray_MapPyTypeToDType(dtype_class, descr->typeobj,
             PyTypeNum_ISUSERDEF(dtype_class->type_num)) < 0) {
         Py_DECREF(dtype_class);