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 @@
+* `PyArray_GetSortFunction`, `PyArray_GetArgSortFunction`, and `PyArray_GetSortCompareFunction` have been added to the C-API. These functions return the sorting, argsorting, and sort comparison functions if provided for a given dtype in new slots.
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,28 @@ 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.
+
 API Functions and Typedefs
 ~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -3521,6 +3543,30 @@ 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, int descending, PyArray_SortFuncWithContext **out_sort, \
+        NpyAuxData **out_auxdata, NPY_ARRAYMETHOD_FLAGS *out_flags)
+
+    If defined, sets a custom sorting function for the DType for each of
+    the sort kinds numpy implements. Returns 0 on success.
+
+.. c:macro:: NPY_DT_get_argsort_function
+
+.. c:type:: int *(PyArrayDTypeMeta_GetArgSortFunction)(PyArray_Descr *, \
+        npy_intp sort_kind, int descending, PyArray_ArgSortFuncWithContext **out_argsort, \
+        NpyAuxData **out_auxdata, NPY_ARRAYMETHOD_FLAGS *out_flags)
+
+    If defined, sets a custom argsorting function for the DType for each of
+    the sort kinds numpy implements. Returns 0 on success.
+
+.. c:macro:: NPY_DT_sort_compare
+
+    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 +3593,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
 
@@ -3588,15 +3636,19 @@ DType API slots but for now we have exposed the legacy
 
 .. c:macro:: NPY_DT_PyArray_ArrFuncs_sort
 
-   An array of PyArray_SortFunc of length ``NPY_NSORTS``. If set, allows
+   An array of PyArray_SortFuncWithContext 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,
+   An array of PyArray_ArgSortFuncWithContext 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 +4392,17 @@ 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_FIRST
+
+      Indicates that NaN values should be sorted first.
+
+   .. c:enumerator:: NPY_SORT_NAN_LAST
+
+      Indicates that NaN values should be sorted last.
 
 .. 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
@@ -494,8 +494,8 @@ PyArray_ArrFuncs
            PyArray_NonzeroFunc *nonzero;
            PyArray_FillFunc *fill;
            PyArray_FillWithScalarFunc *fillwithscalar;
-           PyArray_SortFunc *sort[NPY_NSORTS];
-           PyArray_ArgSortFunc *argsort[NPY_NSORTS];
+           PyArray_SortFuncWithContext *sort[NPY_NSORTS];
+           PyArray_ArgSortFuncWithContext *argsort[NPY_NSORTS];
            PyObject *castdict;
            PyArray_ScalarKindFunc *scalarkind;
            int **cancastscalarkindto;
@@ -792,6 +792,40 @@ 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;
+          int reversed;
+          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.
+
+    .. c:member:: int reversed
+
+       A flag indicating whether the sort is reversed.
+
+    .. 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
@@ -367,6 +367,10 @@ typedef int (PyArrayMethod_PromoterFunction)(PyObject *ufunc,
 #define NPY_DT_get_clear_loop 9
 #define NPY_DT_get_fill_zero_loop 10
 #define NPY_DT_finalize_descr 11
+#define NPY_DT_get_sort_function 12
+#define NPY_DT_get_argsort_function 13
+#define NPY_DT_compare 14
+#define NPY_DT_sort_compare 15
 
 // These PyArray_ArrFunc slots will be deprecated and replaced eventually
 // getitem and setitem can be defined as a performance optimization;
@@ -477,4 +481,47 @@ 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_RESULT;
+
+typedef struct PyArrayMethod_SortContext_tag PyArrayMethod_SortContext;
+
+typedef NPY_COMPARE_RESULT (PyArray_CompareFuncWithContext)(
+        const void *a, const void *b, PyArrayMethod_SortContext *context);
+typedef NPY_COMPARE_RESULT (PyArray_SortCompareFunc)(
+        const void *a, const void *b, PyArrayMethod_SortContext *context);
+
+typedef enum {
+    NPY_SORT_NAN_FIRST = 0,
+    NPY_SORT_NAN_LAST = 1,
+} NPY_SORT_NAN_POSITION;
+
+struct PyArrayMethod_SortContext_tag {
+    PyArray_Descr *descriptor;
+    PyArray_SortCompareFunc *compare;
+    int descending;
+    int reversed;
+    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, int, PyArray_SortFuncWithContext **, NpyAuxData **,
+        NPY_ARRAYMETHOD_FLAGS *out_flags);
+typedef int *(PyArrayDTypeMeta_GetArgSortFunction)(PyArray_Descr *, 
+        npy_intp, int, PyArray_ArgSortFuncWithContext **, NpyAuxData **,
+        NPY_ARRAYMETHOD_FLAGS *out_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/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
@@ -192,6 +192,10 @@ dtypemeta_initialize_struct_from_spec(
     NPY_DT_SLOTS(DType)->common_instance = NULL;
     NPY_DT_SLOTS(DType)->setitem = NULL;
     NPY_DT_SLOTS(DType)->getitem = NULL;
+    NPY_DT_SLOTS(DType)->get_sort_function = NULL;
+    NPY_DT_SLOTS(DType)->get_argsort_function = NULL;
+    NPY_DT_SLOTS(DType)->compare = NULL;
+    NPY_DT_SLOTS(DType)->sort_compare = NULL;
     NPY_DT_SLOTS(DType)->get_clear_loop = NULL;
     NPY_DT_SLOTS(DType)->get_fill_zero_loop = NULL;
     NPY_DT_SLOTS(DType)->finalize_descr = NULL;
@@ -1230,6 +1234,20 @@ 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;
+        }
+    }
+
+    /* Auto-fill compare slot with sort-compare as default */
+    if (dt_slots->compare == NULL && dt_slots->sort_compare != NULL) {
+        dt_slots->compare = dt_slots->sort_compare;
+    }
+
     if (_PyArray_MapPyTypeToDType(dtype_class, descr->typeobj,
             PyTypeNum_ISUSERDEF(dtype_class->type_num)) < 0) {
         Py_DECREF(dtype_class);

diff --git a/numpy/_core/src/multiarray/dtypemeta.h b/numpy/_core/src/multiarray/dtypemeta.h
@@ -67,6 +67,13 @@ typedef struct {
      * parameters, if any, as the operand dtype.
      */
     PyArrayDTypeMeta_FinalizeDescriptor *finalize_descr;
+
+    /* DType sorting methods. */
+    PyArrayDTypeMeta_GetSortFunction *get_sort_function;
+    PyArrayDTypeMeta_GetArgSortFunction *get_argsort_function;
+    PyArray_CompareFuncWithContext *compare;
+    PyArray_SortCompareFunc *sort_compare;
+
     /*
      * The casting implementation (ArrayMethod) to convert between two
      * instances of this DType, stored explicitly for fast access:
@@ -89,7 +96,7 @@ typedef struct {
 
 // This must be updated if new slots before within_dtype_castingimpl
 // are added
-#define NPY_NUM_DTYPE_SLOTS 11
+#define NPY_NUM_DTYPE_SLOTS 15
 #define NPY_NUM_DTYPE_PYARRAY_ARRFUNCS_SLOTS 22
 #define NPY_DT_MAX_ARRFUNCS_SLOT \
   NPY_NUM_DTYPE_PYARRAY_ARRFUNCS_SLOTS + _NPY_DT_ARRFUNCS_OFFSET
@@ -291,6 +298,45 @@ PyArray_SETITEM(PyArrayObject *arr, char *itemptr, PyObject *v)
                 Py_XSETREF(descr, _new_);  \
         } while(0)
 
+static inline int
+PyArray_GetSortFunction(PyArray_Descr *descr, 
+    NPY_SORTKIND which, int descending, PyArray_SortFuncWithContext **out_sort,
+    NpyAuxData **out_auxdata, NPY_ARRAYMETHOD_FLAGS *out_flags)
+{
+    if (NPY_DT_SLOTS(NPY_DTYPE(descr))->get_sort_function == NULL) {
+        return -1;
+    }
+
+    NPY_DT_SLOTS(NPY_DTYPE(descr))->get_sort_function(
+        descr, which, descending, out_sort, out_auxdata, out_flags);
+    return 0;
+}
+
+static inline int
+PyArray_GetArgSortFunction(PyArray_Descr *descr, 
+    NPY_SORTKIND which, int descending, PyArray_ArgSortFuncWithContext **out_argsort,
+    NpyAuxData **out_auxdata, NPY_ARRAYMETHOD_FLAGS *out_flags)
+{
+    if (NPY_DT_SLOTS(NPY_DTYPE(descr))->get_argsort_function == NULL) {
+        return -1;
+    }
+
+    NPY_DT_SLOTS(NPY_DTYPE(descr))->get_argsort_function(
+        descr, which, descending, out_argsort, out_auxdata, out_flags);
+    return 0;
+}
+
+static inline PyArray_CompareFuncWithContext *
+PyArray_GetCompareFunction(PyArray_Descr *descr)
+{
+    return NPY_DT_SLOTS(NPY_DTYPE(descr))->compare;
+}
+
+static inline PyArray_SortCompareFunc *
+PyArray_GetSortCompareFunction(PyArray_Descr *descr)
+{
+    return NPY_DT_SLOTS(NPY_DTYPE(descr))->sort_compare;
+}
 
 // Get the pointer to the PyArray_DTypeMeta for the type associated with the typenum.
 static inline PyArray_DTypeMeta *
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		* `PyArray_GetSortFunction`, `PyArray_GetArgSortFunction`, and `PyArray_GetSortCompareFunction` have been added to the C-API. These functions return the sorting, argsorting, and sort comparison functions if provided for a given dtype in new slots.
Copy link Member seberg Apr 1, 2025 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others. Learn more. You did not actually add them to the public C-API. Which is totally fine, though. (I might start with adding a `SortBuffer()` function or so.) Copy link Member seberg Jul 1, 2025 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others. Learn more. This comment is still valid, I think we might as well remove this for now. MaanasArora reacted with thumbs up emoji