ENH, API: New sorting mechanism for DType API #28516

MaanasArora · 2025-03-14T19:35:19Z

Resolves #26510.

Allocates the lock for the StringDType array before sort and releases after.

I noticed the sorting algorithms independently get the compare function from the descriptor, so I have created new helper functions in stringdtype/dtype.c but not sure if that's the right place. Changes have been made only in quicksort.cpp (but will add others later), so this is a draft but would appreciate feedback.

Posting simple benchmarks in a comment below. Thank you for reviewing!

MaanasArora · 2025-03-14T19:39:33Z

Running this script on both branches:

import numpy
import random
import timeit

print(numpy.__version__)  # '2.0.0rc2'

options = ["a", "bb", "ccc", "dddd"]
lst = random.choices(options, k=1000)
arr_s = numpy.fromiter(lst, dtype="T", count=len(lst))

print(timeit.timeit(lambda: numpy.unique(arr_s), number=10000))

produces on master:

2.3.0.dev0+git20250310.c275e25
3.481267879999905

and on this branch:

2.3.0.dev0+git20250314.0eb0c8e
1.1663875859994732

seberg · 2025-03-25T19:25:20Z

@ngoldbaum just to let you know, I'll let you decide on whether you want this. I am starting to think it is time to implement a get_sortfunction slot but that doesn't mean we can't do this in the mean-time as it's a pretty big speed advantage.

ngoldbaum · 2025-03-25T20:33:09Z

I am starting to think it is time to implement a get_sortfunction slot

Agreed. @MaanasArora as you said this was a draft, would you be up for a bigger refactor? IMO this functionality deserves support in the new DType system without relying on the ArrFuncs baggage.

but that doesn't mean we can't do this in the mean-time as it's a pretty big speed advantage.

Also agreed if you don't want to take this further.

MaanasArora · 2025-03-25T20:36:47Z

Yes, agreed, and willing to do a larger refactor! I actually began by considering special casing array sorting for strings overall, but wondered what the preferred approach would be. I think the sorting routines are not very flexible and could use an overhaul.

PS just to clarify, adding a slot to the dtype will mean we will restructure the sorting to be more generic and allow replacing or extending compare etc.? I'll look further into this but would appreciate pointers.

ngoldbaum · 2025-03-25T20:50:51Z

adding a slot to the dtype will mean we will restructure the sorting to be more generic and allow replacing or extending compare etc.?

Take a look at numpy/_core/src/multiarray/dtypemeta.h - I'm talking about adding a new entry in NPY_DType_Slots that handles comparison. Adding a new slot takes some ceremony - there are some magic constants doing offsets on structs elsewhere in NumPy that need to be updated alongside any changes - but there are comments that should hopefully guide you along your way.

We already have entries for getitem and setitem as well as the legacy arrfuncs slots. You'd be migrating from using NPY_DT_SLOTS(dtype)->f->compare to some new api like NPY_DT_COMPARE(dtype) that uses its own slot and allows per-dtype setup for sorting.

seberg · 2025-03-25T20:51:37Z

PS just to clarify, adding a slot to the dtype will mean we will restructure the sorting to be more generic and allow replacing or extending compare etc.? I'll look further into this but would appreciate pointers.

Yes, if you look around, you will find for example get_clear_loop and the way users can specify it. I think sorting would look similar.
(I.e. the core sort loop probably always works on a contiguous chunk of memory that is aligned -- that part may be simpler would have to check. I assume it would work in-place, but we will also need an argsort.)
The get_sort_function -- or what we name it -- would then get the desired sort-kind passed in, that way we will also have an easier time of adding new sort methods in the future.
(We may want to provision for an ascending/descending flag even if we don't use it).

I see nathan has some other pointers, I don't expect mine to be quite enough, so please ask!

MaanasArora · 2025-03-25T20:56:37Z

Thank you both, this was helpful! Starting to plan this now and will surely clarify if needed.

MaanasArora · 2025-03-26T17:55:56Z

I've added the slots and done some patchy work around using it, and the stringdtype integration. Looking into how to better relate to the array funcs. WIP, but hopefully this is in the right direction!

ngoldbaum · 2025-03-26T17:58:21Z

Yeah, I think you see we still have a lot of other functionality that we should have slots for. Definitely nonzero, at least.

MaanasArora · 2025-03-27T00:02:39Z

Yes, nonzero and some other arrayfuncs could definitely use a slot!

Thank you for the guidance--I've completed most of the missing pieces I think. I assume we would deprecate some of the earlier uses more gradually, so I've fallen back on array funcs as defaults in some places. I think this should be ready for a first pass!

MaanasArora · 2025-03-27T00:04:09Z

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

+    }
+
+    return 0;
+}


This is probably a lot of somewhat redundant code, but I added it here as a 'test' use which provides a boilerplate to replace the routines with a more efficient special-case indirect sort in a future PR. As a "bonus" it allows us to at least temporarily do the allocation this PR was intended for

Nice! I'll let @seberg evaluate whether this API needs adjustments to fit in with the broader DType API but at a first glance it looks reasonable to me, especially if the common case with no specializations can be done with less boilerplate.

That said - there definitely are specialized string sorting implementations we could be using here.

This is good, the thing that I would like to change is the PyArray_(Arg)SortFunc itself, so that it gets a context and auxdata instead of the "array". And this function would get NpyAuxData *out_auxdata in also.
(Although, maybe for sorting this is less interesting as it is not as common to sort many arrays in one.)

The unfortunate thing is that you need to wrap the existing functions if you do this or have a second path for the old function.

I am also considering if we should have a return -2 or so to indicate that the sort-kind is not supported (no error set), to allow NumPy to fall back to a different one.
(But I am not sure we need it, it is useful only for somtehing like mergesort/stablesort, explicitly.)
@charris may have a thought on that.

Looking into passing context now! It looks like a good idea, will try to implement. Thanks.

This change would be really nice; unfortunately both the PyArray_CompareFunc and the sort functions use the PyArrayObject right now, which we probably (!) shouldn't get from the context.

If I'm thinking right, we can define a new type such as PyArray_SortCompareFunc that uses the descr instead of the array and make new sort functions that do not use the array somehow (as we can no longer interchange the SortCompareFunc and the CompareFunc!), but we would still probably need the old functions to use with the older compare functions; I think the duplication will be quite complicated.

At the same time, this would be a missed opportunity to have the new CompareFunc type if we do deprecate later and want to go down this route...

There are two possible approaches here:

We deal with it in the sort function, and just have to different calls depending on whether it is an old or new sort function.

You ignore the fact that an array is currently passed (effectively). We do that in some other places as well, due to how terrible it is.
That is, we wrap it into a dummy object for which basically the only valid field is arr->descr (and maybe arr->flags, don't recall). (See get_dummy_stack_array, yes this is terrible and even reviewers stumble over it, but...)
If you do that, you can write a short function that wraps the old call into a function taking the new one.

I did the second for ufuncs (not sure if that was the easier!), so I suspect the first is likely the simplest here.

Allowing to set a compare function, seems like a nice idea (also to have a simple default).

It would be nice to move that into a default slot function. I.e. rather than setting it for StringDType here, auto-fill the slot with the function that tries to use the SortCompareFunc (If that slot is undefined, we can keep the slot filled with NULL).
That also removes the second check later.

About this SortCompareFunc, it may make sense to keep it "light-weight" (i.e. a single function even if that may not be ideal if you have to inspect the dtype to do the comparisons -- for example structured has to do this).

But I would like to think about what we need to sort things like NaNs if possible. Unfortunately, I am not immediately sure, i.e. <=> in C++ can return a partial order, which means that:

For us probably an "error" is a valid return (right now we can't propagate errors!).

"unordered" is a valid return, although I am not sure how to deal with it. If we have "compare(a, b) == unordered" (i.e. one or both are NaN), we don't yet know how to swap them. That may be possible to resolve with compare(b, b) and compare(a, a).
But the only way I am quite sure how to resolve possible reversed sorting, etc. might be to have unordered_left, unordered_both, unordered_right.

Or we should keep it roughly as is, and accept that this function doesn't exist for all dtypes... A neat thing about having a clear order with "unordered" is that we could also use it from the comparison (u)funcs.

I'll try the first approach I think! I think it will help isolate the new API in a way that makes deprecation easier later too. I will also do this default logic for both compare and sort compare funcs.

As for allowing partial order: yes that could break the symmetry, I suppose. But it could be useful to make sorting more precise in the long run too, and so "unordered" does seem a better way to allow for those kinds of extensions. And then we can use this machinery as the go-to for anywhere comparison decisions need to be made as you said.

Just an update: after some thought, I think it might be quite nice to go with unordered_left, unordered_both, unordered_right, mainly because it saves us any issues with reverse sorting down the line, might as well get everything in, especially as you mentioned dataframes and that's clearly a very important use case. Working on this! I'll try to draft an API that can easily fill in 'defaults' somehow, so that the user-facing side can be used at different levels of complexity / customization.

ngoldbaum · 2025-03-27T14:46:54Z

numpy/_core/include/numpy/dtype_api.h

+        npy_intp, int, PyArray_SortFunc **);
+typedef int *(PyArrayDTypeMeta_GetArgSortFunction)(PyArray_Descr *, 
+        npy_intp, int, PyArray_ArgSortFunc **);
+


New stuff in the public API needs new API docs as well as a release note describing the new features.

Maybe also as a proof-of-concept, it looks like both quaddtype and mpfdtype in numpy-user-dtypes implement sorting - would you be willing to update them to use the new API in a PR to numpy-user-dtypes that depends on this PR to numpy? That should give you a feeling for whether this API is helpful for someone writing a new user dtype. It'll also be a form of documentation - we don't have great docs for writing user dtypes besides the examples in numpy-user-dtypes.

Also what should we do about the flags that got added before we made the dtype API public, e.g. NPY_DT_PyArray_ArrFuncs_compare? I guess we can deprecate them although I don't know how hard it would be to generate deprecation warnings if those are used.

It's easy to generate a deprecation warning during registration (a bit tedious maybe, as you need explicit check).

Sure, I'll add API docs and a release note, and willing to make a PR to numpy-user-dtypes! Will look into that.

Just to be clear, NPY_DT_PyArray_ArrFuncs_compare is still needed right? We can move it to a new slot rather than an arrayfunc but it's going to be different from the sort comparison for now if I'm thinking right (as it is user-facing rather than used in the sorting). Do we need to do this another way?

We can't change slot numbers (unless they are guarded as private)! So the numbers are fixed (until they have not been used for a bit at least).
So yeah, I think we should keep it the old slot for now, maybe easier to make the deprecation a follow up.[^depr]

So, we just have to live with the numbering we got, I half thought I asked for an offset for the NPY_DT_PyArray_ArrFuncs slots, but maybe I didn't bother.
(It's not a big issue, the only thing is the convenience if slot numbers == slot offset so you don't need to translate it.)

[^depr] I think this is as simple as asking users to compile with the new NumPy, and then adding PyArray_RUNTIME_VERSION, but this PR is complicated enough due to API decisions for the new loops, etc.

So, we just have to live with the numbering we got, I half thought I asked for an offset for the NPY_DT_PyArray_ArrFuncs slots, but maybe I didn't bother.

There is an offset, _NPY_DT_ARRFUNCS_OFFSET:

numpy/numpy/_core/src/multiarray/dtypemeta.h

Lines 94 to 95 in 9389862

#define NPY_DT_MAX_ARRFUNCS_SLOT \

NPY_NUM_DTYPE_PYARRAY_ARRFUNCS_SLOTS + _NPY_DT_ARRFUNCS_OFFSET

numpy/_core/src/multiarray/dtypemeta.c

ngoldbaum · 2025-03-27T14:49:54Z

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

+    }
+
+    return 0;
+}


Nice! I'll let @seberg evaluate whether this API needs adjustments to fit in with the broader DType API but at a first glance it looks reasonable to me, especially if the common case with no specializations can be done with less boilerplate.

That said - there definitely are specialized string sorting implementations we could be using here.

ngoldbaum · 2025-03-27T14:50:56Z

numpy/_core/src/multiarray/dtypemeta.h

+}
+
+static inline PyArray_CompareFunc *
+PyArray_SortCompare(PyArray_Descr *descr)


I'd call this PyArray_GetSortCompareFunction

Yes, made this change. Thanks

ngoldbaum · 2025-03-27T14:51:53Z

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`


I think this was added by mistake?

Yes, sorry!

ngoldbaum

Thanks for iterating so quickly :)

I think the new error path needs a little more thought, sorry for the back-and-forth.

ngoldbaum · 2025-03-28T14:23:44Z

doc/release/upcoming_changes/28516.new_feature.rst

+sort-kind and order.
+
+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.


maybe a note that the old arrfuncs slots may be deprecated in the future.

Added, thanks!

ngoldbaum · 2025-03-28T14:24:50Z

numpy/_core/include/numpy/dtype_api.h

+#define NPY_DT_sort_compare 11
+#define NPY_DT_get_clear_loop 12
+#define NPY_DT_get_fill_zero_loop 13
+#define NPY_DT_finalize_descr 14


can you re-order these so the slots that were already in the struct keep their old values? I don't know offhand if changing this order is problematic but it seems more consistent to not change the old values even if it's fine.

numpy/_core/src/multiarray/dtypemeta.c

seberg

Sorry, long comments, and I realize this is becoming a lot more complex than it may have looked initially. But, I would really like to see the context passed in/a new signature.

That also probably includes filling in/returning ARRAY_METHODFLAGS, even if the only useful flag is "requires GIL".

I also still tend to think it may make sense to have a magic return for "unsupported sort method", although should maybe ask Chuck once in a meeting about that.
(in principle I agree we usually just need stable and not-stable, but if we want users to be able to choose more precisely, I think it may make sense to allow us to fallback here. We could still use something like "no error indicated, but func == NULL for it even, but maybe a special return is easier.)

If needed, maybe we have to talk briefly about it synchronously? Or maybe just write the docs/signatures first that we want for the public API.

seberg · 2025-04-01T09:40:25Z

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.


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.)

seberg · 2025-04-01T10:02:41Z

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

+    }
+
+    return 0;
+}


There are two possible approaches here:

We deal with it in the sort function, and just have to different calls depending on whether it is an old or new sort function.

You ignore the fact that an array is currently passed (effectively). We do that in some other places as well, due to how terrible it is.
That is, we wrap it into a dummy object for which basically the only valid field is arr->descr (and maybe arr->flags, don't recall). (See get_dummy_stack_array, yes this is terrible and even reviewers stumble over it, but...)
If you do that, you can write a short function that wraps the old call into a function taking the new one.

I did the second for ufuncs (not sure if that was the easier!), so I suspect the first is likely the simplest here.

Allowing to set a compare function, seems like a nice idea (also to have a simple default).

It would be nice to move that into a default slot function. I.e. rather than setting it for StringDType here, auto-fill the slot with the function that tries to use the SortCompareFunc (If that slot is undefined, we can keep the slot filled with NULL).
That also removes the second check later.

About this SortCompareFunc, it may make sense to keep it "light-weight" (i.e. a single function even if that may not be ideal if you have to inspect the dtype to do the comparisons -- for example structured has to do this).

But I would like to think about what we need to sort things like NaNs if possible. Unfortunately, I am not immediately sure, i.e. <=> in C++ can return a partial order, which means that:

For us probably an "error" is a valid return (right now we can't propagate errors!).

"unordered" is a valid return, although I am not sure how to deal with it. If we have "compare(a, b) == unordered" (i.e. one or both are NaN), we don't yet know how to swap them. That may be possible to resolve with compare(b, b) and compare(a, a).
But the only way I am quite sure how to resolve possible reversed sorting, etc. might be to have unordered_left, unordered_both, unordered_right.

Or we should keep it roughly as is, and accept that this function doesn't exist for all dtypes... A neat thing about having a clear order with "unordered" is that we could also use it from the comparison (u)funcs.

seberg · 2025-04-01T10:09:43Z

numpy/_core/src/multiarray/dtypemeta.h

+    NPY_SORTKIND which, int descending, PyArray_SortFunc **out_sort)
+{
+    if (NPY_DT_SLOTS(NPY_DTYPE(descr))->get_sort_function == NULL) {
+        return -1;


This needs to set an error (TypeError or DTypeTypeError, which is defined somewhere I think.)

(An error here will make sense after you move the fallback logic into a default slot filling. Or you could have the default slot raise an error, that is also completely fine.)

MaanasArora · 2025-04-01T10:41:10Z

No worries; thank you for the detailed feedback actually! It's nice to be able to iron out the direction for the API. I'll address the docs and public API changes and keep working away at the SortFunc changes. Happy to have a synchronous chat if it seems useful.

MaanasArora · 2025-04-05T00:30:10Z

numpy/_core/include/numpy/dtype_api.h

@@ -477,4 +481,18 @@ typedef PyArray_Descr *(PyArrayDTypeMeta_FinalizeDescriptor)(PyArray_Descr *dtyp
 typedef int(PyArrayDTypeMeta_SetItem)(PyArray_Descr *, PyObject *, char *);
 typedef PyObject *(PyArrayDTypeMeta_GetItem)(PyArray_Descr *, char *);

+typedef int (PyArray_CompareFuncWithDescr)(const void *, const void *,
+                                           PyArray_Descr *);


The naming is a bit weird here, but I didn't want to disturb the original type as it's used a lot. I think the SortCompareFunc should still be a unique type so will do that (even if only a clone of this type).

I have slightly mixed feelings. On the one hand, I think this is the pragmatic thing to have.
On the other hand, we could also look this function from the np.less_than or np.great_than ufunc to implement sorting, I think.
(The problem there is still how to deal with unordered elements, a compare ufunc would work better...)

But, on the other hand, it seems pragmatic even if it won't work well e.g. for structured dtypes (performance issues), it will always work and provides an easy entry-point (we can also use this to define default comparison ufuncs).

So overall, I think I end up at just doing this, although I could imaging punting if we don't need it for StringDType (I suspect we do, though).

Would like to hear if @ngoldbaum has an opinion.

(A neater future path would also be if this was more of a header-only code binding generator job with us making the sorting patterns available maybe. I.e. if this was defined in a C++ class and our sort code available, the DType could compile the full loop and avoid calling such a helper everywhere.)

IMO this is fine, if only because it exists right now 😄

MaanasArora · 2025-04-05T00:39:31Z

Sorry for the bit of delay, I was thinking through this and essentially ended up with separating more the legacy sorting machinery from this API. This way, the new signatures can freely use context-related features and we do not have to create some sort of empty array or refactor a very large number of sorting-related files. Aside from how nice the feature is, I think this separation is actually a plus and even rolling back the stringdtype integration was worth it (in any case, user dtypes are not to define sorting with the internal, now legacy functions, so may be best to add a specialized routine).

We also have the new compare slot which defaults to sort_compare and related features now, though they're not used yet. Hopefully this is in the right direction. If it is, we can gradually move the older sorting machinery to the context signatures, thus converging. Thank you!

ngoldbaum · 2025-04-11T19:17:10Z

Sorry for not getting to this yet. I'm going to try to make sure to give this a once-over next week.

I think you can fix the test failures by rebasing?

MaanasArora · 2025-04-11T20:02:02Z

No worries, I have some things to address as well.

Just rebased--sorry not sure if things went perfectly smoothly.

MaanasArora · 2025-04-11T21:23:37Z

Just brought this implementation with the new signature to parity with the previous one, including the usage in StringDType and ensuring use cases for the new and old sort functions are handled properly in the handlers. There is repetitive code but I guess we will phase out the legacy slots. Now we can make the context nicer if needed, and incorporate the auxdata!

MaanasArora · 2025-05-08T20:58:03Z

Hello! Getting back to this, is there anything I need to address? Thinking of adding the functions to the public C-API if things look fine.

Would we need to create a new C-API version (regenerate the hashes and such), and I guess it would come under 2.4, given how close 2.3 is?

seberg

A few comments, yeah, this won't make 2.3, sorry. I think it might be good to discuss a bit in depth with @ngoldbaum some time (not next week, sorry).

Another thing that I would like addressed/discussed is the problem of reverse sorting.
I do think we need at least a reverse=True, I think it might make sense to also provision for a nan_position (if nan goes first or last).

(NULL/NA ordering is very important in dataframe world, and I am tempted to include this, even if we say that the value for now is always "last").

seberg · 2025-05-09T07:33:55Z

doc/source/reference/c-api/array.rst

@@ -1873,6 +1873,29 @@ described below.
   pointer. Currently this is used for zero-filling and clearing arrays storing
   embedded references.

+.. c:type:: int (PyArray_SortFunc)( \
+                 void *start, npy_intp num, PyArrayMethod_Context *context, \


Let's move the context to the first spot just for similarity. I think I added a context for traversal functions, I am not sure that was smart, but since we have it, it may be a slightly better fit.

I might call start, data (not that it matters).

This is done!

seberg · 2025-05-09T07:36:34Z

doc/source/reference/c-api/array.rst

+                 NpyAuxData *auxdata, NpyAuxData **out_auxdata)
+    
+    A function to sort a buffer of data. The *start* is a pointer to the
+    beginning of the buffer containing *num* elements. A function of this


Suggested change

beginning of the buffer containing *num* elements. A function of this

beginning of the contiguous containing *num* elements. A function of this

It also should be aligned, but I have to think whether we should allow indicating support for unaligned data here. (Which would require flags, for ufuncs "supports unaligned" is flagged before get_loop(), although since here we always do contiguous, flagging it inside get_loop() is OK also -- that is, becuase get_loop() is not passed any strides).

Makes sense! Committed (modulo typo).

seberg · 2025-05-09T07:37:24Z

doc/source/reference/c-api/array.rst

+                 NpyAuxData **out_auxdata)
+    
+    A function to arg-sort a buffer of data. The *start* is a pointer to the
+    beginning of the buffer containing *num* elements. The *tosort* is a


@charris to confirm, even for argsorting it probably makes sense to always use a contiguous buffer for sorting?

seberg · 2025-05-09T07:38:18Z

doc/source/reference/c-api/array.rst

+                 PyArrayMethod_Context *context, NpyAuxData *auxdata, \
+                 NpyAuxData **out_auxdata)


Suggested change

PyArrayMethod_Context *context, NpyAuxData *auxdata, \

NpyAuxData **out_auxdata)

PyArrayMethod_Context *context, NpyAuxData *auxdata)

The out_auxdata belongs on the get_loop function!

This is done, thank you :)

seberg · 2025-05-09T07:39:43Z

doc/source/reference/c-api/array.rst

+.. c:macro:: NPY_DT_get_sort_function
+
+.. c:type:: int *(PyArrayDTypeMeta_GetSortFunction)(PyArray_Descr *, \
+        npy_intp sort_kind, int descending, PyArray_SortFunc **out_sort);


This needs *out_flags, since we need the ability to indicate whether the GIL is required (I think we can ignore FPEs), but who knows if we'll have a reason for other flags eventually.

It also needs **out_auxdata, since auxdata needs to come from somewhere :).

Done, thanks!

seberg · 2025-05-09T07:47:40Z

numpy/_core/src/multiarray/item_selection.c

@@ -1570,20 +1592,41 @@ PyArray_Sort(PyArrayObject *op, int axis, NPY_SORTKIND which)
        return -1;
    }

-    sort = PyDataType_GetArrFuncs(PyArray_DESCR(op))->sort[which];
+    PyArray_GetSortFunction(PyArray_DESCR(op), which, 0, &sort);


Hmmm, let's just use < 0 to decide if it's an error. In which case sort != NULL is assumed.

Makes sense, this is done!

seberg · 2025-05-09T08:06:00Z

numpy/_core/include/numpy/dtype_api.h

@@ -477,4 +481,18 @@ typedef PyArray_Descr *(PyArrayDTypeMeta_FinalizeDescriptor)(PyArray_Descr *dtyp
 typedef int(PyArrayDTypeMeta_SetItem)(PyArray_Descr *, PyObject *, char *);
 typedef PyObject *(PyArrayDTypeMeta_GetItem)(PyArray_Descr *, char *);

+typedef int (PyArray_CompareFuncWithDescr)(const void *, const void *,
+                                           PyArray_Descr *);


I have slightly mixed feelings. On the one hand, I think this is the pragmatic thing to have.
On the other hand, we could also look this function from the np.less_than or np.great_than ufunc to implement sorting, I think.
(The problem there is still how to deal with unordered elements, a compare ufunc would work better...)

But, on the other hand, it seems pragmatic even if it won't work well e.g. for structured dtypes (performance issues), it will always work and provides an easy entry-point (we can also use this to define default comparison ufuncs).

So overall, I think I end up at just doing this, although I could imaging punting if we don't need it for StringDType (I suspect we do, though).

Would like to hear if @ngoldbaum has an opinion.

(A neater future path would also be if this was more of a header-only code binding generator job with us making the sorting patterns available maybe. I.e. if this was defined in a C++ class and our sort code available, the DType could compile the full loop and avoid calling such a helper everywhere.)

MaanasArora · 2025-05-10T01:28:45Z

Thanks for the comments! And no worries, I was just making sure I wasn't missing something to do :)

I think I need to think a bit more about the best way to adjust this for the extra features you mentioned, yes. Unordered elements is definitely something to consider at this stage, so I might try to draft something for that soon enough; that should hopefully create a clearer story around these features!

ngoldbaum · 2025-05-14T01:28:22Z

I want to call your attention to this suggestion: #28516 (comment).

Did you ever take a look at numpy-user-dtypes? A worked example would help.

MaanasArora · 2025-05-14T05:51:18Z

Yes, sorry, I took a look actually--but was having some trouble with installing the dtype packages over the editable install of numpy. I'll push my draft anyway and try to address that. Thanks for the reminder.

…y defined

…ures

…tions

…t_loop

ngoldbaum · 2025-05-20T13:42:18Z

Just a head's up I haven't forgotten about this. I'm planning to spend some time later this week or next looking closely at this and the accompanying numpy-user-dtypes PR. Thanks so much for working on this having patience 🙂.,

MaanasArora · 2025-05-20T13:55:13Z

Thank you @ngoldbaum! I realize there's a lot here so appreciate you taking time on this. I'll try to address Sebastian's other comments and iterate soon.

ngoldbaum

This is getting there! Left a few comments inline.

ngoldbaum · 2025-05-22T19:52:17Z

numpy/_core/include/numpy/dtype_api.h

+                               NpyAuxData *);
+typedef int (PyArray_ArgSortFunc)(PyArrayMethod_Context *, 
+                                  void *, npy_intp *, npy_intp, 
+                                  NpyAuxData *);


These two need different names and you need to leave the original typedefs in ndarraytypes.h that had these names, since they're public API.

Thanks for reviewing! This is done.

ngoldbaum · 2025-05-22T19:54:18Z

numpy/_core/include/numpy/dtype_api.h

@@ -477,4 +481,18 @@ typedef PyArray_Descr *(PyArrayDTypeMeta_FinalizeDescriptor)(PyArray_Descr *dtyp
 typedef int(PyArrayDTypeMeta_SetItem)(PyArray_Descr *, PyObject *, char *);
 typedef PyObject *(PyArrayDTypeMeta_GetItem)(PyArray_Descr *, char *);

+typedef int (PyArray_CompareFuncWithDescr)(const void *, const void *,
+                                           PyArray_Descr *);


IMO this is fine, if only because it exists right now 😄

…lic API

github-actions bot added the 01 - Enhancement label Mar 14, 2025

MaanasArora commented Mar 27, 2025

View reviewed changes

ngoldbaum reviewed Mar 27, 2025

View reviewed changes

MaanasArora changed the title ~~ENH: Allocate lock only once in StringDType quicksort~~ ENH, API: New sorting mechanism for DType API Mar 28, 2025

ngoldbaum reviewed Mar 28, 2025

View reviewed changes

seberg reviewed Apr 1, 2025

View reviewed changes

MaanasArora commented Apr 5, 2025

View reviewed changes

MaanasArora force-pushed the enh/faster-string-sorting branch from d7cd9ed to 0e8b6a5 Compare April 11, 2025 21:29

seberg reviewed May 9, 2025

View reviewed changes

MaanasArora force-pushed the enh/faster-string-sorting branch from 39f5ef2 to 237b7f0 Compare May 13, 2025 22:39

MaanasArora mentioned this pull request May 14, 2025

ENH: Add usage of new sorting dtype slots numpy/numpy-user-dtypes#109

Draft

MaanasArora added 22 commits May 20, 2025 00:04

ENH: Add descending flag to internal sorting functions

b89accd

MAINT: Improve get dtype sort compare function name

a437eb9

MAINT: Fix doc typo

aa63d11

MAINT: Error out when non-legacy dtype has no sort_compare function

16e95a2

DOC: Add release notes for new dtype sorting API

42e76d6

DOC: Add doc for sort compare slot in release notes

88636cc

DOC: Add note for potential deprecation of sort arrfuncs in release note

9d14ec1

MAINT: Reorder dtype slots to prevent changing existing slot numbers

a556455

BUG: Error on missing sort_compare slot only when dtype is privatel…

3c0957e

…y defined

DOC: Add C-API documentation for new sorting slots

9506798

ENH: Replace array object with context and auxdata in sortfunc signat…

6ce5351

…ures

BUG: Fix unnecessarily private function call due to underscore typo

96a53b2

MAINT: Fix whitespace typos

9a2b100

ENH: Allow flexible sorting compare for arr or descr in npy_sort func…

8d4c75d

…tions

ENH: Add new sort func implementations and use in stringdtype

50988ba

DOC: Fix missing newline in ctype doc

ca5797e

DOC: Add sortfunc typedef docs

95cfd8f

DOC: Fix missing newline in ctype doc

6dd4f4c

ENH: Define SortCompareFunc type

4fa813c

Update dtype sorting signatures: move context, move out auxdata to ge…

894911e

…t_loop

MAINT: Check error in Get(Arg)SortFunc using return value

57687ac

DOC: Add missing newlines to c-types in array.rst

0edb4ea

MaanasArora force-pushed the enh/faster-string-sorting branch from be43aeb to 0edb4ea Compare May 20, 2025 04:04

ngoldbaum reviewed May 22, 2025

View reviewed changes

MaanasArora added 4 commits May 24, 2025 00:12

MAINT: Rename new sort funcs and restore older names for existing pub…

167301e

…lic API

MAINT: Rename start pointer in new sort func documentation to data

e6b8c1e

ENH: Add flags to new get_(arg)sort_function

579c351

DOC: Mention new sort func buffers to be contiguous

d854b00

	#define NPY_DT_MAX_ARRFUNCS_SLOT \
	NPY_NUM_DTYPE_PYARRAY_ARRFUNCS_SLOTS + _NPY_DT_ARRFUNCS_OFFSET

		@@ -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.

	beginning of the buffer containing num elements. A function of this
	beginning of the contiguous containing num elements. A function of this

		PyArrayMethod_Context context, NpyAuxData auxdata, \
		NpyAuxData **out_auxdata)

Uh oh!

ENH, API: New sorting mechanism for DType API #28516

Are you sure you want to change the base?

ENH, API: New sorting mechanism for DType API #28516

Conversation

MaanasArora commented Mar 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

MaanasArora commented Mar 14, 2025

Uh oh!

seberg commented Mar 25, 2025

Uh oh!

ngoldbaum commented Mar 25, 2025

Uh oh!

MaanasArora commented Mar 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ngoldbaum commented Mar 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

seberg commented Mar 25, 2025

Uh oh!

MaanasArora commented Mar 25, 2025

Uh oh!

MaanasArora commented Mar 26, 2025

Uh oh!

ngoldbaum commented Mar 26, 2025

Uh oh!

MaanasArora commented Mar 27, 2025

Uh oh!

MaanasArora Mar 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MaanasArora Mar 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ngoldbaum Mar 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ngoldbaum left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MaanasArora commented Mar 14, 2025 •

edited

Loading

MaanasArora commented Mar 25, 2025 •

edited

Loading

ngoldbaum commented Mar 25, 2025 •

edited

Loading

MaanasArora Mar 27, 2025 •

edited

Loading

MaanasArora Mar 27, 2025 •

edited

Loading

ngoldbaum Mar 27, 2025 •

edited

Loading

MaanasArora commented Apr 5, 2025 •

edited

Loading

MaanasArora commented Apr 11, 2025 •

edited

Loading

MaanasArora commented May 8, 2025 •

edited

Loading