ENH: Add dtype slot sorting functionality support to all sort kinds

numpy · MaanasArora · Mar 14, 2025 · Mar 26, 2025 · Mar 26, 2025 · Mar 26, 2025
commit 59590d29057fcbde9849c8523dbcbcd2b79bc22b
diff --git a/numpy/_core/src/multiarray/stringdtype/dtype.c b/numpy/_core/src/multiarray/stringdtype/dtype.c
@@ -539,18 +539,78 @@ _stringdtype_quicksort(void *start, npy_intp num, void *varr) {
     return _stringdtype_sort(start, num, varr, &npy_quicksort);
 }
 
+int
+_stringdtype_heapsort(void *start, npy_intp num, void *varr) {
+    return _stringdtype_sort(start, num, varr, &npy_heapsort);
+}
+
+int
+_stringdtype_timsort(void *start, npy_intp num, void *varr) {
+    return _stringdtype_sort(start, num, varr, &npy_timsort);
+}
+
 int
 stringdtype_get_sort_function(
     PyArray_Descr *descr, NPY_SORTKIND sort_kind, PyArray_SortFunc **out_sort) {
 
     switch (sort_kind) {
+        default:
         case NPY_QUICKSORT:
             *out_sort = &_stringdtype_quicksort;
-            return 0;
+            break;
+        case NPY_HEAPSORT:
+            *out_sort = &_stringdtype_heapsort;
+            break;
+        case NPY_STABLESORT:
+            *out_sort = &_stringdtype_timsort;
+            break;
+    }
+
+    return 0;
+}
+
+int
+_stringdtype_argsort(void *vv, npy_intp *tosort, npy_intp num, void *varr,
+                     PyArray_ArgSortFunc *argsort) {
+    PyArray_StringDTypeObject *descr = (PyArray_StringDTypeObject *)PyArray_DESCR(varr);
+
+    NpyString_acquire_allocator(descr);
+    int result = argsort(vv, tosort, num, varr);
+    NpyString_release_allocator(descr->allocator);
+
+    return result;
+}
+
+int
+_stringdtype_aquicksort(void *vv, npy_intp *tosort, npy_intp n, void *varr) {
+    return _stringdtype_argsort(vv, tosort, n, varr, &npy_aquicksort);
+}
+
+int
+_stringdtype_aheapsort(void *vv, npy_intp *tosort, npy_intp n, void *varr) {
+    return _stringdtype_argsort(vv, tosort, n, varr, &npy_aheapsort);
+}
+
+int
+_stringdtype_atimsort(void *vv, npy_intp *tosort, npy_intp n, void *varr) {
+    return _stringdtype_argsort(vv, tosort, n, varr, &npy_atimsort);
+}
+
+int
+stringdtype_get_argsort_function(
+    PyArray_Descr *descr, NPY_SORTKIND sort_kind, PyArray_ArgSortFunc **out_argsort) {
+
+    switch (sort_kind) {
         default:
-            PyErr_Format(PyExc_ValueError,
-                         "Sort kind %d not supported for string dtype", sort_kind);
-            return -1;
+        case NPY_QUICKSORT:
+            *out_argsort = &npy_aquicksort;
+            break;
+        case NPY_HEAPSORT:
+            *out_argsort = &npy_aheapsort;
+            break;
+        case NPY_STABLESORT:
+            *out_argsort = &npy_atimsort;
+            break;
     }
 
     return 0;
@@ -697,6 +757,8 @@ static PyType_Slot PyArray_StringDType_Slots[] = {
         {NPY_DT_setitem, &stringdtype_setitem},
         {NPY_DT_getitem, &stringdtype_getitem},
         {NPY_DT_sort_compare, &stringdtype_sort_compare},
+        {NPY_DT_get_sort_function, &stringdtype_get_sort_function},
+        {NPY_DT_get_argsort_function, &stringdtype_get_argsort_function},
         {NPY_DT_ensure_canonical, &stringdtype_ensure_canonical},
         {NPY_DT_PyArray_ArrFuncs_nonzero, &nonzero},
         {NPY_DT_PyArray_ArrFuncs_compare, &compare},

diff --git a/numpy/_core/src/npysort/heapsort.cpp b/numpy/_core/src/npysort/heapsort.cpp
@@ -54,7 +54,7 @@ npy_heapsort(void *start, npy_intp num, void *varr)
 {
     PyArrayObject *arr = (PyArrayObject *)varr;
     npy_intp elsize = PyArray_ITEMSIZE(arr);
-    PyArray_CompareFunc *cmp = PyDataType_GetArrFuncs(PyArray_DESCR(arr))->compare;
+    PyArray_CompareFunc *cmp = PyArray_SortCompare(PyArray_DESCR(arr));
     if (elsize == 0) {
         return 0;  /* no need for sorting elements of no size */
     }
@@ -116,7 +116,7 @@ npy_aheapsort(void *vv, npy_intp *tosort, npy_intp n, void *varr)
     char *v = (char *)vv;
     PyArrayObject *arr = (PyArrayObject *)varr;
     npy_intp elsize = PyArray_ITEMSIZE(arr);
-    PyArray_CompareFunc *cmp = PyDataType_GetArrFuncs(PyArray_DESCR(arr))->compare;
+    PyArray_CompareFunc *cmp = PyArray_SortCompare(PyArray_DESCR(arr));
     npy_intp *a, i, j, l, tmp;
 
     /* The array needs to be offset by one for heapsort indexing */

diff --git a/numpy/_core/src/npysort/mergesort.cpp b/numpy/_core/src/npysort/mergesort.cpp
@@ -385,7 +385,7 @@ npy_mergesort(void *start, npy_intp num, void *varr)
 {
     PyArrayObject *arr = (PyArrayObject *)varr;
     npy_intp elsize = PyArray_ITEMSIZE(arr);
-    PyArray_CompareFunc *cmp = PyDataType_GetArrFuncs(PyArray_DESCR(arr))->compare;
+    PyArray_CompareFunc *cmp = PyArray_SortCompare(PyArray_DESCR(arr));
     char *pl = (char *)start;
     char *pr = pl + num * elsize;
     char *pw;
@@ -461,7 +461,7 @@ npy_amergesort(void *v, npy_intp *tosort, npy_intp num, void *varr)
 {
     PyArrayObject *arr = (PyArrayObject *)varr;
     npy_intp elsize = PyArray_ITEMSIZE(arr);
-    PyArray_CompareFunc *cmp = PyDataType_GetArrFuncs(PyArray_DESCR(arr))->compare;
+    PyArray_CompareFunc *cmp = PyArray_SortCompare(PyArray_DESCR(arr));
     npy_intp *pl, *pr, *pw;
 
     /* Items that have zero size don't make sense to sort */

diff --git a/numpy/_core/src/npysort/timsort.cpp b/numpy/_core/src/npysort/timsort.cpp
@@ -2250,7 +2250,7 @@ npy_timsort(void *start, npy_intp num, void *varr)
 {
     PyArrayObject *arr = reinterpret_cast<PyArrayObject *>(varr);
     size_t len = PyArray_ITEMSIZE(arr);
-    PyArray_CompareFunc *cmp = PyDataType_GetArrFuncs(PyArray_DESCR(arr))->compare;
+    PyArray_CompareFunc *cmp = PyArray_SortCompare(PyArray_DESCR(arr));
     int ret;
     npy_intp l, n, stack_ptr, minrun;
     run stack[TIMSORT_STACK_SIZE];
@@ -2686,7 +2686,7 @@ npy_atimsort(void *start, npy_intp *tosort, npy_intp num, void *varr)
 {
     PyArrayObject *arr = reinterpret_cast<PyArrayObject *>(varr);
     size_t len = PyArray_ITEMSIZE(arr);
-    PyArray_CompareFunc *cmp = PyDataType_GetArrFuncs(PyArray_DESCR(arr))->compare;
+    PyArray_CompareFunc *cmp = PyArray_SortCompare(PyArray_DESCR(arr));
     int ret;
     npy_intp l, n, stack_ptr, minrun;
     run stack[TIMSORT_STACK_SIZE];