ENH: Allocate lock only once in StringDType quicksort

MaanasArora · MaanasArora · commit 87f410b7f63b · 2025-06-01T00:01:01.000-04:00
diff --git a/numpy/_core/src/multiarray/stringdtype/dtype.c b/numpy/_core/src/multiarray/stringdtype/dtype.c
@@ -459,6 +459,33 @@ compare(void *a, void *b, void *arr)
     return ret;
 }
 
+int
+_compare_no_mutex(const void *a, const void *b, void *arr)
+{
+    PyArray_StringDTypeObject *descr = (PyArray_StringDTypeObject *)PyArray_DESCR(arr);
+    return _compare((void*) a, (void*) b, descr, descr);
+}
+
+void
+_init_sort_cmp(PyArray_Descr *descr, PyArray_CompareFunc **out_cmp)
+{
+    if (descr->type_num == NPY_VSTRING) {
+        NpyString_acquire_allocator((PyArray_StringDTypeObject *)descr);
+        *out_cmp = _compare_no_mutex;
+    }
+    else {
+        *out_cmp = PyDataType_GetArrFuncs(descr)->compare;
+    }
+}
+
+void
+_end_sort_cmp(PyArray_Descr *descr)
+{
+    if (descr->type_num == NPY_VSTRING) {
+        NpyString_release_allocator(((PyArray_StringDTypeObject *)descr)->allocator);
+    }
+}
+
 int
 _compare(void *a, void *b, PyArray_StringDTypeObject *descr_a,
          PyArray_StringDTypeObject *descr_b)
diff --git a/numpy/_core/src/multiarray/stringdtype/dtype.h b/numpy/_core/src/multiarray/stringdtype/dtype.h
@@ -19,12 +19,21 @@ new_stringdtype_instance(PyObject *na_object, int coerce);
 NPY_NO_EXPORT int
 init_string_dtype(void);
 
+NPY_NO_EXPORT int
+_compare_no_mutex(const void *a, const void *b, void *arr);
+
 // Assumes that the caller has already acquired the allocator locks for both
 // descriptors
 NPY_NO_EXPORT int
 _compare(void *a, void *b, PyArray_StringDTypeObject *descr_a,
          PyArray_StringDTypeObject *descr_b);
 
+NPY_NO_EXPORT void
+_init_sort_cmp(PyArray_Descr *descr, PyArray_CompareFunc **out_cmp);
+
+NPY_NO_EXPORT void
+_end_sort_cmp(PyArray_Descr *descr);
+
 NPY_NO_EXPORT int
 init_string_na_object(PyObject *mod);
 
diff --git a/numpy/_core/src/npysort/quicksort.cpp b/numpy/_core/src/npysort/quicksort.cpp
@@ -44,7 +44,7 @@
  * the below code implements this converted to an iteration and as an
  * additional minor optimization skips the recursion depth checking on the
  * smaller partition as it is always less than half of the remaining data and
- * will thus terminate fast enough
+ * will thus terminate fast enough`
  */
 
 #define NPY_NO_DEPRECATED_API NPY_API_VERSION
@@ -56,6 +56,7 @@
 #include "numpy_tag.h"
 #include "x86_simd_qsort.hpp"
 #include "highway_qsort.hpp"
+#include "stringdtype/dtype.h"
 
 #include <cstdlib>
 #include <utility>
@@ -510,7 +511,7 @@ npy_quicksort(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;
     char *vp;
     char *pl = (char *)start;
     char *pr = pl + (num - 1) * elsize;
@@ -521,6 +522,8 @@ npy_quicksort(void *start, npy_intp num, void *varr)
     int *psdepth = depth;
     int cdepth = npy_get_msb(num) * 2;
 
+    _init_sort_cmp(PyArray_DESCR(arr), &cmp);
+
     /* Items that have zero size don't make sense to sort */
     if (elsize == 0) {
         return 0;
@@ -606,6 +609,9 @@ npy_quicksort(void *start, npy_intp num, void *varr)
     }
 
     free(vp);
+    
+    _end_sort_cmp(PyArray_DESCR(arr));
+
     return 0;
 }
 
@@ -615,7 +621,7 @@ npy_aquicksort(void *vv, npy_intp *tosort, npy_intp num, 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;
     char *vp;
     npy_intp *pl = tosort;
     npy_intp *pr = tosort + num - 1;
@@ -626,6 +632,8 @@ npy_aquicksort(void *vv, npy_intp *tosort, npy_intp num, void *varr)
     int *psdepth = depth;
     int cdepth = npy_get_msb(num) * 2;
 
+    _init_sort_cmp(PyArray_DESCR(arr), &cmp);
+
     /* Items that have zero size don't make sense to sort */
     if (elsize == 0) {
         return 0;
@@ -700,6 +708,8 @@ npy_aquicksort(void *vv, npy_intp *tosort, npy_intp num, void *varr)
         cdepth = *(--psdepth);
     }
 
+    _end_sort_cmp(PyArray_DESCR(arr));
+
     return 0;
 }
 
