I think this needs to be documented in the C API docs for the dtype API:

https://numpy.org/devdocs/reference/c-api/array.html#slot-ids-and-api-function-typedefs

Yes thanks, missed that. This is done! Will add a release note also.

Awesome, thanks! Can you also update your PR over at numpy-user-dtypes to make use of the updated API (if any updates are needed)? numpy/numpy-user-dtypes#109

The CI over there is slow at the moment so hopefully that's not annoying. We need to add some caching for the quaddtype build.

Sure! Currently that PR relies on the sort_compare feature as well, so the changes are incompatible, but we can write the sort functions themselves. So I'm looking into specialized sorting algorithms for these dtypes which will allow for more comprehensive documentation as well (instead of just sort_compare)!

Actually, do you think we should get this PR to parity first instead? It doesn't seem those dtypes have too many considerations for overall sorting itself, and I think a lot of examples including StringDType rely on the comparison feature. I think a comparison system for using generic sorts is about non-negotiable at this point, so I might as well add it now.

Sure, whichever you think would work best. Just wanted to make sure it's on your radar.

Note that sortkind has the property that sortkind & 0x7 >> 1 yields an index ranging from 0 to 3 that can be used to pick the sorting functions if they are in an array. That trick is not currently used, but things were deliberately designed so that it would work.

A question that I also raised in #29642 - is the design of having a struct table with pointers to different sort functions for different properties set in stone? To me, it feels illogical if we are having flags to determine sort function behaviour -- it is less extensible than just passing in flags to a general function, and does not allow later adding "request" flags later (like, optimize for low memory usage or for performance).

It seems to me more sensible to have just two functions (sort, argsort), both of which takes flags (perhaps via the context struct?), and have some way to expose which flags they can or cannot deal with (either as a mask exposed in a struct, or via a return code).

Implementation-wise, it would also seem easy enough for a developer to implement the table if they want to go that route, while the other way around is just annoying (I can easily imagine that ascending/descending is easily done in one function).

But it may well be that I'm missing the bigger picture here!

ascending/descending is easily done in one function).

Only for the generic sorts. The type specific sorts gain speed by having < built in. The idea is to avoid choice/function calls in the innermost loops. That said, it is easy to build functions with alternative definitions of <, but that would give you two functions, not one. Even for the generic loops, that might be the way to go, instead of looking for -1 returned by the compare functions, look for +1. Note the absence of 0.

Once this is settled, I plan to generate the needed descending sort functions.

EDIT: It is also easy to add entries to a private tables without causing incompatibilities. I don't see the added complexities of hints as having much use at the moment, they can be added later if needed. The only real use I can see would be for out-of-memory sorts, and that would probably qualify as a requirement rather than a hint.

In #28516 we actually had just two functions that mirrored ufuncs, get_(arg)sort_function that returned a pointer to the appropriate function given the parameters, but Chuck had the good point that (as I understand) we effectively have that many functions anyway if we want to be explicit about efficiency. I can see a user dtype implementing deeply nested code or such in a single function, which may not be that good for performance, so now actually I think it makes sense to be explicit that we're looking for multiple functions in practice, unless I'm missing something!

@charris - yes, my example was perhaps poorly chosen. My question is really more general: isn't it better to define the API as just having only one sort and one argsort function that take flags? Much easier to extend in the future. (Though perhaps I read too many LWN articles where people bemoan the fact that old kernel function do not have flags, making the creation of a new one necessary...)

The only downside I see is that this does not automatically provide a way to tell that a given dtype can sort descending, etc., hence the idea of a mask or specific return code. Though it probably is fine to generically just error if a sort function doesn't support it (same as function=NULL in the present proposal) and fall back to a standard sort using the comparison.

@MaanasArora - sorry, missed your answer, which appeared while I wrote mine.

I can see the wish for performance but it seems to me calling a function with flags that internally just has a switch right up front (with each option a separate code block, say in-lined sort functions) would not impact performance much. Though I also wonder if any of this is going to matter: sorting generally is not a fast process.

Anyway, I don't want to make too much of this...

@mhvk no problem! Yes, it would probably not impact performance terribly, and if user dtypes want to define a simple sort, even if slightly inefficient, I guess we could let them :)

Another benefit of having the older architecture that @seberg suggested would be that it mirrors ufuncs, and allows for more flexible configuration with the context as you said.

The reason (as far as I understand) we went with structs, was that I think tables are more explicit in a way — we have 4 ways of doing sorts, so declaring them outright is simpler in its own way as we have straightforward and performant routines to call without needing to consider configuration.

Personally, I think the config doesn't change the code in the routines much at all so it might be simpler to just pass it into two functions. I guess it depends on whether the performance gains are good to have too.

I'm happy to go with whichever approach everyone thinks is better for the project, I am just clarifying what I understand so far!

we have 4 ways of doing sorts

Except that @charris also introduced an option to sort nans first or last... It is currently commented out, but it would make it eight options... Which is partially why I like passing on flags... 😺

Note that I'd be absolutely fine with going with a mechanism mirroring the ufuncs -- in the end, sorting could in principle easily be done with a gufunc if that would have a flags option.

p.s. Sorry to come to this discussion late!

Right, so basically we have 4 typical implementations. With NaNs, it would be 8, although we wouldn't need all. I am now leaning towards not worrying about NaNs, though. I just doubt that sorting NaNs to the front is niche enough to make it hard.

I can see the wish for performance but it seems to me calling a function with flags that internally just has a switch right up front

I think optimization is actually a reason to expose the get_* style API. The overhead is tiny, but:

• It makes it much easier to reason about how to build an optimized version for structured dtypes. This could have a huge effect on
• If we add sort-hints (or other metadata to the context), it might allow future further optimized versions to select the best algorithm once. E.g. a "this is a very small sort" hint.

Not sure this matters, but I think optimization is a reason for the get_ style API. The reason for allowing to just pass 4 functions here is that it is less verbose to implement (i.e. user, including NumPy convenience).

UFuncs actually support both styles. The closest design to what ufuncs actually be having both!
But, I am happy to just do either one, we can always add a NPY_DT_sort_getter slot and one should only define one of them. We could also add some void *reserved = NULL at the end of the struct, but since NPY_DT_sort_getter and NPY_DT_sort_functions would be mutually exclusive, I actually don't think I care about that.

(Implementation wise, NPY_DT_sort_getter could be set to a default that looks up what's in NPY_DT_sort_functions.)

There was the question of making it easy to use the generic implementation by only providing a compare function. But it seems fair to defer that to a follow-up, it is also again be somewhat mutually exclusive (since the compare function could fill in all of partitioning, sorting, argsorting, and potentially even compare ufuncs!).

I am quite happy to leave the design choice to @seberg here. My first instinct would be to go with the simplest implementation, but agree that a getter approach offers more flexibility for future changes. That said, the future is hard to predict, so we shouldn't spend too much time on it. I think the most important thing in the short term is to make alternative selections easy, SIMD for instance. I think Sebastian is thinking of run time dispatch as a possibility.

@seberg - thanks for the details! I still should at some point look at how the ufuncs actually do it, but perhaps you can help by pointing out in the code or docs where things are defined... If ufuncs do both, then I think we should pick one of those for now, and mimic it as close as possible. (To me, the getter still seems the more scaleable solution, so I'd prefer that...)

Ah, right, I would love to move the runtime dispatching into the getter. I didn't think of it TBH, because that isn't where it is for ufuncs right now (it may be there for casts, not sure).

I think I convinced myself in the previous comment, that both is useful (for simplicity) but that it is perfectly reasonable to use two distinct slots for both (because you only need to use one of them).
Given that, I would argue to push on with the current version and then follow up, even if I would like to follow up probably.

(I do have to go through the API here, although I suspect it is pretty much settled from the last PR.)

If ufuncs do both, then I think we should pick one of those for now, and mimic it as close as possible. (To me, the getter still seems the more scaleable solution, so I'd prefer that...)

Hehe, fine, we can also start with the getter. I was considering starting with this just to make the step of changing our internals slightly easier :). Honestly, I don't care, so if you like a getter maybe that is the way to go.

@mhvk I think it's easier to look at the public API here. Not great layout, but you can see: https://numpy.org/doc/2.1/reference/c-api/array.html#c.NPY_METH_strided_loop the slots (of which most ufuncs only need the first strided-loop one, but casts need more). And further down the get_loop one with the corresponding signature.
The implementation for ufunc just fills in a default get_loop that checks what is in the other slots.
(But ArrayMethods have their own full set of slots, while here it would be nice to integrate it into the DType slots, so if we add functions I think a struct is nice.)

Thanks, that's useful!

One question before proceeding is whether it might be possible to treat the sorting code as an array method. Since effectively these are gufuncs, I think it should be, but it is not clear to me how one would pass on the sorting flags. Though perhaps the analogy is that sorting ascending and descending (e.g.) is like different gufuncs. Maybe that is actually the way to go?

Maybe that is actually the way to go?

Maybe, but that means 4 gufuncs here if we don't figure out parameters and I think that is too much for the immediate need (if interesting). And sorting may be a bit special? We simplified to contiguous loops since it sounded like supporting strided ones is unlikely to be useful (but of course one could likely wrangle that into gufuncs to focus on that).
Sorting is also naturally in-place (not argsorting of course), which is also slightly awkward for ufuncs probably.

I suppose I agree that gufuncs might be the better thing, but without figuring out some loose ends that seems hard. And figuring out the loose ends seems maybe a bit much for this.

I guess the question is whether we can get away with defining a get_loop that has the same signature?

If ufuncs do both, then I think we should pick one of those for now, and mimic it as close as possible. (To me, the getter still seems the more scaleable solution, so I'd prefer that...)

I do suspect sorting is a complex enough operation that some consideration for future extensions is important; the getter solution seems to target many of them well enough and it seems more natural to what this API actually does (allow users to define custom sorts). If optimization is actually easier with it, that's great and seems more reason to go that way!

There was the question of making it easy to use the generic implementation by only providing a compare function. But it seems fair to defer that to a follow-up, it is also again be somewhat mutually exclusive (since the compare function could fill in all of partitioning, sorting, argsorting, and potentially even compare ufuncs!).

Makes sense! That should make the size of this PR much more palatable :) In practice we're still accessing the compare arrayfunc, so we can probably fit the stringdtype fix into this PR without needing to move to a slot for that. Though I'm happy to do it in a small follow-up too if you prefer.

I suppose I agree that gufuncs might be the better thing, but without figuring out some loose ends that seems hard. And figuring out the loose ends seems maybe a bit much for this.

Would it be possible to transition to gufuncs using this API down the line? We could presumably allow the same API that we have now and internally insert it into a gufunc?

Edit: I mean the getter API, not the one implemented right now, since it seems we're going that way!

I guess the question is whether we can get away with defining a get_loop that has the same signature?

Would it be possible to transition to gufuncs using this API down the line?

I suppose. In the end, the missing piece of gufuncs is the fact that their context never holds parameters (and parameters here are restricted to things that are scalar and do not take part in promotion, probably -- e.g. that doesn't work for quantiles where the parameter changes promotion).
The sort context would have to for the sort options flags. So the thing to solve would be two-fold, I think:

1. You would have to modify argument parsing to create a gufunc specific parameter context structure. (Maybe, that doesn't need to be a struct usually, because if you have a void * slot, you can cast it the flags here). That also needs cleanup (maybe as NpyAuxData *, not sure!).
   • No, I don't think any of this is particularly hard, but not sure it is easy either and there may be other design approaches. (I do agree that it would be nice to make context be identical to such a future thing here, but I am not sure I want to worry about it.)
2. Our sort loops don't support strides, that could be solved via a small sort specific copier that wraps that calls the contiguous version in principle. Or more generic tooling.
3. Sort loops work in-place. They effectively have zero inputs and a single output (that we need to read from). I suspect this part is largely fine, but might need a clearer "input and output" marker within the ufunc machinery.

We could presumably allow the same API that we have now and internally insert it into a gufunc?

If we do create a gufunc with the above capabilities, then yes, that gufunc could use this mechanism to implement itself for all dtypes that support it via this (while also allowing to do it directly if you want to support e.g. strided versions).

(The biggest annoyance/dance would be that if the context is structured differently, the translater/wrapper needs to modify it additionally to dropping the strides argument.)

I am not sure that all of this should matter. I like the idea of doing the above, but I don't like the idea of sitting on this change for it indefinitely.

@mhvk what does this mean to you? To me it doesn't change much right now to be honest (I wish it did, but I don't want to stall things). Unless this is an argument for doing just the explicit 4 loops for now, because a get_loop() could be done if we manage to wranlge it into a gufunc.

Yes, we do need to move forward and concreteness would be good...

But maybe first back to basics: looking again at the PR, there are two things I'd really like to change.

Use the standard array method loop signature

I'd suggest we ensure that the (arg)sort loops have the standard array method loop signature:

typedef int (PyArrayMethod_StridedLoop)(
    PyArrayMethod_Context *context,
    char *const *data,
    const npy_intp *dimensions,
    const npy_intp *strides,
    NpyAuxData *auxdata
)

This is obviously needed if we move to (arg)sort being gufuncs, but I think it is good regardless, just not to expand special cases. It is fine to state that for sort this will only be called with data[0] == data[1] (though perhaps it is not a problem to require that a sort implementation just copies the data if that is not the case; could eventually be an implementation flag).

Use loop registration, not slots on the dtype

It feels like a mistake to me to define things on the dtype, in a way moving us back to the previous API, where the dtype defined a set of "handy" functions in PyArray_ArrFuncs. I think it had become clear this was a bad idea and instead we now have array methods, where a dtype registers itself as needed. I think we should go via registration here too.

For registration, using the ufunc registration mechanism seems by far the easiest -- any custom one would mostly duplicate it. Making (arg)sort gufuncs may be more work than warranted but perhaps one could at first have "fake" ones that just work as storage and expose those via names. In that case, registration would be via,

PyUFunc_AddLoopFromSpec(NPY_[ARG]SORT_UFUNC, spec);

and inside the current sort/argsort code, the relevant loop would be looked up in those fake gufuncs.

Alternatively, one could expose a registration function specifically for (arg)sort that wraps this.

Decide on how to deal with the flags

The way the array method system is currently written, there is no straightforward way to pass on runtime information to a ufuncs other than the dtypes. Thus, by far the easiest way to deal with flags is to imagine different gufunc for the different sort types. In that case, registration would become,

PyUFunc_AddLoopFromSpec(NPY_[ARG]SORT_STABLE_ASCENDING, spec);
PyUFunc_AddLoopFromSpec(NPY_[ARG]SORT_STABLE_DESCENDING, spec);
...

(which would argue even more for wrapping this in a PyArray_RegisterSortLoopsFromSpecs(....))

That said, it would seem relatively easy to add a new void *runtime_data to PyArrayMethod_Context which in the [arg]sort main implementation could be used to pass on the flags. It may not be that long before that will be needed for other cases too...

p.s. The above is mostly because this PR introduces new public API, which I feel we need to get somewhat right. A private implementation for StringDType could be substantially hackier...

Thanks @mhvk for reviewing! I've handled two reference counting comments and will look at the rest (such as the flags for GIL) after @seberg's review. Glad to hear things are looking good.

Edit: Wow, that PR for ufuncs on scalars has a huge speedup indeed!

The test failure is unrelated.

p.s. @seberg - it would be nice to expose the ability to get loops from a ufunc. Partially inspired by this PR, I wrote a little class that multiplies whatever is input by a constant, and ended up getting the legacy loop via ->types and ->functions just because it seemed impossible to get the loop directly using the regular API. Or did I miss something?

We need to add such public API, there is the Python side "unstable" API as np.add._resolve_dtypes_and_context (check it's docs).
Yes, I would be happy to expose this as more convenient API, it's that, because for testing/prototyping that seemed like the right level at the time.

OK, I think it's best to just put this in, thanks @MaanasArora, sorry that this is so difficult with the API design and ufunc back and forth. Right now, it changes (practically) nothing public API wise anyway. The added but unused *parameters is fine, and honestly a good idea anyway.

Yay!! A really nice contribution, looking forward to the next steps, of allowing registration!

Thank you both for reviewing; I learned a lot here! Looking forward to the next steps as well.

Same here, I learned a lot (thanks @seberg! what you implemented is really nice, but it is non-trivial to understand without a simple example - which is why I like the scaled float case so much...)