Skip to content

Commit 0e8b6a5

Browse files
MaanasAroraMaanasArora
committed
ENH: Add new sort func implementations and use in stringdtype
1 parent 84be767 commit 0e8b6a5

File tree

8 files changed

+266
-6
lines changed

8 files changed

+266
-6
lines changed

‎numpy/_core/src/common/npy_sort.h.src

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
#include <Python.h>
66
#include <numpy/npy_common.h>
77
#include <numpy/ndarraytypes.h>
8+
#include <numpy/dtype_api.h>
89

910
#define NPY_ENOMEM 1
1011
#define NPY_ECOMP 2
@@ -97,6 +98,22 @@ NPY_NO_EXPORT int atimsort_@suff@(void *vec, npy_intp *ind, npy_intp cnt, void *
9798
*****************************************************************************
9899
*/
99100

101+
NPY_NO_EXPORT int npy_quicksort_with_context(void *vec, npy_intp cnt,
102+
PyArrayMethod_Context *context, NpyAuxData *auxdata, NpyAuxData **auxdata_out);
103+
NPY_NO_EXPORT int npy_heapsort_with_context(void *vec, npy_intp cnt,
104+
PyArrayMethod_Context *context, NpyAuxData *auxdata, NpyAuxData **auxdata_out);
105+
NPY_NO_EXPORT int npy_mergesort_with_context(void *vec, npy_intp cnt,
106+
PyArrayMethod_Context *context, NpyAuxData *auxdata, NpyAuxData **auxdata_out);
107+
NPY_NO_EXPORT int npy_timsort_with_context(void *vec, npy_intp cnt,
108+
PyArrayMethod_Context *context, NpyAuxData *auxdata, NpyAuxData **auxdata_out);
109+
NPY_NO_EXPORT int npy_aquicksort_with_context(void *vec, npy_intp *ind, npy_intp cnt,
110+
PyArrayMethod_Context *context, NpyAuxData *auxdata, NpyAuxData **auxdata_out);
111+
NPY_NO_EXPORT int npy_aheapsort_with_context(void *vec, npy_intp *ind, npy_intp cnt,
112+
PyArrayMethod_Context *context, NpyAuxData *auxdata, NpyAuxData **auxdata_out);
113+
NPY_NO_EXPORT int npy_amergesort_with_context(void *vec, npy_intp *ind, npy_intp cnt,
114+
PyArrayMethod_Context *context, NpyAuxData *auxdata, NpyAuxData **auxdata_out);
115+
NPY_NO_EXPORT int npy_atimsort_with_context(void *vec, npy_intp *ind, npy_intp cnt,
116+
PyArrayMethod_Context *context, NpyAuxData *auxdata, NpyAuxData **auxdata_out);
100117

101118
NPY_NO_EXPORT int npy_quicksort(void *vec, npy_intp cnt, void *arr);
102119
NPY_NO_EXPORT int npy_heapsort(void *vec, npy_intp cnt, void *arr);

‎numpy/_core/src/multiarray/item_selection.c

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1599,6 +1599,23 @@ PyArray_Sort(PyArrayObject *op, int axis, NPY_SORTKIND which)
15991599
}
16001600

16011601
if (sort_with_array == NULL) {
1602+
if (PyArray_GetSortCompareFunction(PyArray_DESCR(op)) != NULL) {
1603+
switch (which) {
1604+
default:
1605+
case NPY_QUICKSORT:
1606+
sort = npy_quicksort_with_context;
1607+
break;
1608+
case NPY_HEAPSORT:
1609+
sort = npy_heapsort_with_context;
1610+
break;
1611+
case NPY_STABLESORT:
1612+
sort = npy_timsort_with_context;
1613+
break;
1614+
}
1615+
}
1616+
}
1617+
1618+
if (sort == NULL) {
16021619
if (PyDataType_GetArrFuncs(PyArray_DESCR(op))->compare != NULL) {
16031620
switch (which) {
16041621
default:
@@ -1755,6 +1772,23 @@ PyArray_ArgSort(PyArrayObject *op, int axis, NPY_SORTKIND which)
17551772
}
17561773

17571774
if (argsort_with_array == NULL) {
1775+
if (PyArray_GetSortCompareFunction(PyArray_DESCR(op)) != NULL) {
1776+
switch (which) {
1777+
default:
1778+
case NPY_QUICKSORT:
1779+
argsort = npy_aquicksort_with_context;
1780+
break;
1781+
case NPY_HEAPSORT:
1782+
argsort = npy_aheapsort_with_context;
1783+
break;
1784+
case NPY_STABLESORT:
1785+
argsort = npy_atimsort_with_context;
1786+
break;
1787+
}
1788+
}
1789+
}
1790+
1791+
if (argsort == NULL) {
17581792
if (PyDataType_GetArrFuncs(PyArray_DESCR(op))->compare != NULL) {
17591793
switch (which) {
17601794
default:

‎numpy/_core/src/multiarray/stringdtype/dtype.c

Lines changed: 115 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -518,9 +518,121 @@ _compare(void *a, void *b, PyArray_StringDTypeObject *descr_a,
518518
}
519519

520520
static int
521-
stringdtype_sort_compare(void *a, void *b, void *arr) {
522-
PyArray_StringDTypeObject *descr = (PyArray_StringDTypeObject *)PyArray_DESCR(arr);
523-
return _compare(a, b, descr, descr);
521+
stringdtype_sort_compare(void *a, void *b, PyArray_Descr *descr) {
522+
PyArray_StringDTypeObject *string_descr = (PyArray_StringDTypeObject *)descr;
523+
return _compare(a, b, string_descr, string_descr);
524+
}
525+
526+
int
527+
_stringdtype_sort(void *start, npy_intp num, PyArrayMethod_Context *context,
528+
NpyAuxData *auxdata, NpyAuxData **out_auxdata,
529+
PyArray_SortFunc *sort) {
530+
PyArray_StringDTypeObject *descr = (PyArray_StringDTypeObject *)context->descriptors[0];
531+
532+
NpyString_acquire_allocator(descr);
533+
int result = sort(start, num, context, auxdata, out_auxdata);
534+
NpyString_release_allocator(descr->allocator);
535+
536+
return result;
537+
}
538+
539+
int
540+
_stringdtype_quicksort(void *start, npy_intp num, PyArrayMethod_Context *context,
541+
NpyAuxData *auxdata, NpyAuxData **out_auxdata) {
542+
return _stringdtype_sort(start, num, context, auxdata, out_auxdata,
543+
&npy_quicksort_with_context);
544+
}
545+
546+
int
547+
_stringdtype_heapsort(void *start, npy_intp num, PyArrayMethod_Context *context,
548+
NpyAuxData *auxdata, NpyAuxData **out_auxdata) {
549+
return _stringdtype_sort(start, num, context, auxdata, out_auxdata,
550+
&npy_heapsort_with_context);
551+
}
552+
553+
int
554+
_stringdtype_timsort(void *start, npy_intp num, PyArrayMethod_Context *context,
555+
NpyAuxData *auxdata, NpyAuxData **out_auxdata) {
556+
return _stringdtype_sort(start, num, context, auxdata, out_auxdata,
557+
&npy_timsort_with_context);
558+
}
559+
560+
int
561+
stringdtype_get_sort_function(PyArray_Descr *descr,
562+
NPY_SORTKIND sort_kind, int descending, PyArray_SortFunc **out_sort) {
563+
564+
switch (sort_kind) {
565+
default:
566+
case NPY_QUICKSORT:
567+
*out_sort = &_stringdtype_quicksort;
568+
break;
569+
case NPY_HEAPSORT:
570+
*out_sort = &_stringdtype_heapsort;
571+
break;
572+
case NPY_STABLESORT:
573+
*out_sort = &_stringdtype_timsort;
574+
break;
575+
}
576+
577+
return 0;
578+
}
579+
580+
int
581+
_stringdtype_argsort(void *vv, npy_intp *tosort, npy_intp num,
582+
PyArrayMethod_Context *context, NpyAuxData *auxdata,
583+
NpyAuxData **out_auxdata,
584+
PyArray_ArgSortFunc *argsort) {
585+
PyArray_StringDTypeObject *descr = (PyArray_StringDTypeObject *)context->descriptors[0];
586+
587+
NpyString_acquire_allocator(descr);
588+
int result = argsort(vv, tosort, num, context, auxdata, out_auxdata);
589+
NpyString_release_allocator(descr->allocator);
590+
591+
return result;
592+
}
593+
594+
int
595+
_stringdtype_aquicksort(void *vv, npy_intp *tosort, npy_intp n,
596+
PyArrayMethod_Context *context, NpyAuxData *auxdata,
597+
NpyAuxData **out_auxdata) {
598+
return _stringdtype_argsort(vv, tosort, n, context, auxdata, out_auxdata,
599+
&npy_aquicksort_with_context);
600+
}
601+
602+
int
603+
_stringdtype_aheapsort(void *vv, npy_intp *tosort, npy_intp n,
604+
PyArrayMethod_Context *context, NpyAuxData *auxdata,
605+
NpyAuxData **out_auxdata) {
606+
return _stringdtype_argsort(vv, tosort, n, context, auxdata, out_auxdata,
607+
&npy_aheapsort_with_context);
608+
}
609+
610+
int
611+
_stringdtype_atimsort(void *vv, npy_intp *tosort, npy_intp n,
612+
PyArrayMethod_Context *context, NpyAuxData *auxdata,
613+
NpyAuxData **out_auxdata) {
614+
return _stringdtype_argsort(vv, tosort, n, context, auxdata, out_auxdata,
615+
&npy_atimsort_with_context);
616+
}
617+
618+
int
619+
stringdtype_get_argsort_function(PyArray_Descr *descr,
620+
NPY_SORTKIND sort_kind, int descending, PyArray_ArgSortFunc **out_argsort) {
621+
622+
switch (sort_kind) {
623+
default:
624+
case NPY_QUICKSORT:
625+
*out_argsort = &_stringdtype_aquicksort;
626+
break;
627+
case NPY_HEAPSORT:
628+
*out_argsort = &_stringdtype_aheapsort;
629+
break;
630+
case NPY_STABLESORT:
631+
*out_argsort = &_stringdtype_atimsort;
632+
break;
633+
}
634+
635+
return 0;
524636
}
525637

526638
// PyArray_ArgFunc

‎numpy/_core/src/npysort/heapsort.cpp

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,14 +49,32 @@
4949
*****************************************************************************
5050
*/
5151

52+
NPY_NO_EXPORT int
53+
npy_heapsort_with_context(void *start, npy_intp num,
54+
PyArrayMethod_Context *context, NpyAuxData *auxdata,
55+
NpyAuxData **out_auxdata)
56+
{
57+
return handle_npysort_with_context(start, num, context, auxdata,
58+
out_auxdata, &npy_heapsort);
59+
}
60+
61+
NPY_NO_EXPORT int
62+
npy_aheapsort_with_context(void *vv, npy_intp *tosort, npy_intp num,
63+
PyArrayMethod_Context *context, NpyAuxData *auxdata,
64+
NpyAuxData **out_auxdata)
65+
{
66+
return handle_npyasort_with_context(vv, tosort, num, context, auxdata,
67+
out_auxdata, &npy_aheapsort);
68+
}
69+
5270
NPY_NO_EXPORT int
5371
npy_heapsort(void *start, npy_intp num, void *varr)
5472
{
5573
void *arr;
5674
npy_intp elsize;
5775
PyArray_CompareFunc *cmp;
5876
fill_sort_data_from_arr_or_descr(varr, &arr, &elsize, &cmp);
59-
77+
6078
if (elsize == 0) {
6179
return 0; /* no need for sorting elements of no size */
6280
}

‎numpy/_core/src/npysort/mergesort.cpp

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -335,6 +335,24 @@ string_amergesort_(type *v, npy_intp *tosort, npy_intp num, void *varr)
335335
*****************************************************************************
336336
*/
337337

338+
NPY_NO_EXPORT int
339+
npy_mergesort_with_context(void *start, npy_intp num,
340+
PyArrayMethod_Context *context, NpyAuxData *auxdata,
341+
NpyAuxData **out_auxdata)
342+
{
343+
return handle_npysort_with_context(start, num, context, auxdata,
344+
out_auxdata, &npy_mergesort);
345+
}
346+
347+
NPY_NO_EXPORT int
348+
npy_amergesort_with_context(void *vv, npy_intp *tosort, npy_intp num,
349+
PyArrayMethod_Context *context, NpyAuxData *auxdata,
350+
NpyAuxData **out_auxdata)
351+
{
352+
return handle_npyasort_with_context(vv, tosort, num, context, auxdata,
353+
out_auxdata, &npy_amergesort);
354+
}
355+
338356
static void
339357
npy_mergesort0(char *pl, char *pr, char *pw, char *vp, npy_intp elsize,
340358
PyArray_CompareFunc *cmp, PyArrayObject *arr)

‎numpy/_core/src/npysort/npysort_common.h

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,33 @@
1010
extern "C" {
1111
#endif
1212

13+
14+
/*
15+
*****************************************************************************
16+
** NEW SORTFUNC HANDLERS **
17+
*****************************************************************************
18+
*/
19+
20+
static inline int
21+
handle_npysort_with_context(void *start, npy_intp num,
22+
PyArrayMethod_Context *context, NpyAuxData *auxdata,
23+
NpyAuxData **out_auxdata,
24+
PyArray_SortFuncWithArray *sort)
25+
{
26+
PyArray_Descr *descr = context->descriptors[0];
27+
return sort(start, num, descr);
28+
}
29+
30+
static inline int
31+
handle_npyasort_with_context(void *vv, npy_intp *tosort, npy_intp num,
32+
PyArrayMethod_Context *context, NpyAuxData *auxdata,
33+
NpyAuxData **out_auxdata,
34+
PyArray_ArgSortFuncWithArray *asort)
35+
{
36+
PyArray_Descr *descr = context->descriptors[0];
37+
return asort(vv, tosort, num, descr);
38+
}
39+
1340
static inline void
1441
fill_sort_data_from_arr_or_descr(void *arr_or_descr, void **out_arr_or_descr,
1542
npy_intp *elsize, PyArray_CompareFunc **out_cmp)
@@ -28,8 +55,6 @@ fill_sort_data_from_arr_or_descr(void *arr_or_descr, void **out_arr_or_descr,
2855
}
2956
}
3057

31-
32-
3358
/*
3459
*****************************************************************************
3560
** SWAP MACROS **

‎numpy/_core/src/npysort/quicksort.cpp

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -511,6 +511,24 @@ string_aquicksort_(type *vv, npy_intp *tosort, npy_intp num, void *varr)
511511
*****************************************************************************
512512
*/
513513

514+
NPY_NO_EXPORT int
515+
npy_quicksort_with_context(void *start, npy_intp num,
516+
PyArrayMethod_Context *context, NpyAuxData *auxdata,
517+
NpyAuxData **out_auxdata)
518+
{
519+
return handle_npysort_with_context(start, num, context, auxdata,
520+
out_auxdata, &npy_quicksort);
521+
}
522+
523+
NPY_NO_EXPORT int
524+
npy_aquicksort_with_context(void *vv, npy_intp *tosort, npy_intp num,
525+
PyArrayMethod_Context *context, NpyAuxData *auxdata,
526+
NpyAuxData **out_auxdata)
527+
{
528+
return handle_npyasort_with_context(vv, tosort, num, context, auxdata,
529+
out_auxdata, &npy_aquicksort);
530+
}
531+
514532
NPY_NO_EXPORT int
515533
npy_quicksort(void *start, npy_intp num, void *varr)
516534
{

‎numpy/_core/src/npysort/timsort.cpp

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1851,6 +1851,24 @@ string_atimsort_(void *start, npy_intp *tosort, npy_intp num, void *varr)
18511851
*****************************************************************************
18521852
*/
18531853

1854+
NPY_NO_EXPORT int
1855+
npy_timsort_with_context(void *start, npy_intp num,
1856+
PyArrayMethod_Context *context, NpyAuxData *auxdata,
1857+
NpyAuxData **out_auxdata)
1858+
{
1859+
return handle_npysort_with_context(start, num, context, auxdata,
1860+
out_auxdata, &npy_timsort);
1861+
}
1862+
1863+
NPY_NO_EXPORT int
1864+
npy_atimsort_with_context(void *vv, npy_intp *tosort, npy_intp num,
1865+
PyArrayMethod_Context *context, NpyAuxData *auxdata,
1866+
NpyAuxData **out_auxdata)
1867+
{
1868+
return handle_npyasort_with_context(vv, tosort, num, context, auxdata,
1869+
out_auxdata, &npy_atimsort);
1870+
}
1871+
18541872
typedef struct {
18551873
char *pw;
18561874
npy_intp size;

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy