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

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

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.

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.

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.

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!

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

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!

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

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!

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.

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!

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?

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

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

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!

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?

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!

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.

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.

It would be helpful to document how this apparatus is intended to operate. Maybe with some diagrams with arrows tracing a call.

I have a somewhat different concept as to how this should be handled. I would just add two slots that take function dispatchers that get passed the context and return the relevant sorting implementation, if there is one, null otherwise.

The context would also contain a stable variable, maybe an inplace variable as well, and leave out the compare function (an implementation detail) and maybe be versioned? The returned function would have the comparison function built in (there might be another level of indirection here), making it an implementation detail (sorting comparison, not the dtype comparison). Essentially, this keeps most of the sorting implementation private. The net effect would be to make it possible to extend sorting functionality by adding to the context structure.

I don't think this is very different from what is here at the lower level, but the upper level is somewhat different. It just adds a level of indirection. For reuse we probably need access to all the numpy builtin dtypes, so that their sorting functions can be called from user dtypes if needed. Maybe a toplevel function dispatcher? We currently have builtin introsort, heapsort, timsort, and radix sort, do we want to expose those to users? We should add to those decending sorts and the nan options.

The main idea here is the explicit use of function dispatchers, and maybe an explicit distinction between builtin and user. The goal should be simplicity for the average user, with detailed selection for experts who want it. I am reminded of our random module in that regard.

I think discussion should be about the top level of this. The original idea for making the builtin sorting functions a library was to make it easy to add them to the existing dtype slots.

I would just add two slots that take function dispatchers that get passed the context and return the relevant sorting implementation, if there is one, null otherwise.

@charris can you sketch that out slightly more? I think, you are basically saying to drop the element-wise compare function and just do the PyArrayDTypeMeta_GetSortFunction:
https://github.com/numpy/numpy/pull/28516/files#diff-136726bc7d1ef2257dd24ea927c31155ad76543e9fbd6eeb96c1a8b78da458faR3556-R3565
(which is the function the DType defines)

Then, rather than exposing an elementwise function, we have a public GetSortFunction() so that you can "inherit" the behavior from another DType (i.e. if I implement a unit dtype, I can call the one for float64).
Since we have a context that is specific to sort (and not shared with casts/ufuncs), I agree it makes sense to put a lot of information on it (that way the sort implementation that is returned is passed the context as well and has the information if desired).

For StringDType that may make the implementation slightly more complicated, but that is not a big deal, and we can probably figure out another path to make it easy.

FWIW, I am not too worried about versioning. We can add things to the context if backwards compatible. If not, it is possible to add a GetSortFunctionX slot and deprecate the old one.

Thanks @charris, that makes sense! I think we can definitely simplify the implementation. If I’m understanding correctly, you mean we could drop the sort-compare and expose the internal npy_sort functions instead, to be used in user implementations?

That sounds really nice, but the current sort functions read the comparison from the DType (a legacy array function rather than a slot, via PyArray_GetArrFuncs(dtype)→compare). We could change these to read the sort_compare from the context, but it seems you would like not to pass the sort_compare in the context?

Essentially I am wondering if the sort-compare user is actually more common than the sort user, mainly because most 'compound' types like StringDType or multi-precision float would need a new sort_compare or a complete new sort, and sort-compare is easier. Of course a dtype could go all the way and rewrite the functions completely, but then exposing the builtin functions may not be that useful.

Of course this doesn’t mean we need a full new slot! It actually seems to be an implementation detail, yes – when you say ‘indirection’ at the sort-compare level, would you like something like this (except we pass the sort-compare some other way if you don’t want it in the context)?

stringdtype_get_sort_function(SortContext *ctx)
{
  npy_quicksort_with_context(ctx, &stringdtype_sort_compare, ...);

  // or something like this, though probably not this?
  NpySortContext_SetCompare(ctx, &stringdtype_sort_compare);
  npy_quicksort_with_context(ctx, ...);
}

I think this would actually be quite nice! Users can choose to do this, or reimplement the full function.

Sorry, this got long, but just wanted to clarify what you were thinking. Thanks again for your feedback; it was very helpful!

but the current sort functions read the comparison from the DType'

They don't, except for generic types. The comparison functions are defined as inline functions in npysort_common.h. They tend to be in the inner loop, so it is best to avoid a function call there.

What I am thinking is that for public use there are four attributes to select sorting shim functions. The following information needs to be passed, either in context or by arguments:

• dtype (we could save one level of indirection by using a separate dispatcher for each dtype)
• stable=True (array api). If False, the implementation might still be stable (radix sort).
• decending=False (array api)
• nan position.

Note that with included dtype, only one dispatcher is needed for all dtypes. Or we could omit the dtype and attach different dispatchers to each dtype.

The returned sorting shim function allocates/releases working memory, makes contiguous copies if needed, and calls the lowlevel sorting implementation. Arguments something like

• array (for data and shape and flags)
• axis
• output array (if we want that option)

This is close to what the current implementation has evolved to. In the beginning we called by algorithm, but here call by functionality.

The question if we want to expose the lowest level functions, or maybe the shims, in addition to the dispatchers.

I don't intend to this to be the final proposal, but I think this design should start by answering the question of what we want the new interface to accomplish. Feedback and alternate proposals would be good to hear.

They don't, except for generic types. The comparison functions are defined as inline functions in npysort_common.h. They tend to be in the inner loop, so it is best to avoid a function call there.

Ah, right thanks, those are for most builtin types indeed. Only dtypes for which explicit dispatchers aren't defined actually take the npy_quick/heap/timsort. I suppose my question was moreso if every user dtype would have to necessarily write a sort function - currently, that is not needed, as we pass the compare arrfunc and use the generic sorts. But if we move to new slots as in this PR we need to somehow inherit the compare or define it, unless we want every dtype to write their own sort.

That aside, the API seems good! What I understand is that we pass stable instead of the sort_kind? The interface in this PR does contain:

• dtype: one dispatcher for each, so not passed
• sort_kind: this is the algorithm, rather than stable (this is not in the context, but can be moved there)
• descending: in the context
• nan_position: in the context

I'm not sure I fully understand the dispatcher idea but it seems useful to add a layer on top of the low-level sort. The lower-level generic implementation, as far as I can see, does not contain axis and output array arguments.

So, to be clear if I understand this correctly, if we want to implement these changes, I think we first need to decide:

1. if / how to handle comparisons for user dtypes if a custom sort function is not defined
2. what level of internal sort functionality we are exposing

And then maybe:

1. Add a shim between the sort function, and make dispatching explicit
2. Handle all memory handling (including axes and output) in the underlying shim
3. Expose parts of the internal sorting library to be reused by users
4. (if needed) change how sort comparisons are defined

Sorry if this is off base. I'll think about this further, thanks!

Edit: If all of this won't be done in this PR, perhaps just the user API with some dispatch architecture (such that it fixes the external API) would be a good place to start. Once it's fixed, we can change the internals.

I would be happy to say that we only need stable/unstable sort kind. The main comment I have is that

The returned sorting shim function allocates/releases working memory, makes contiguous copies if needed, and calls the lowlevel sorting implementation. [...]
Arguments something like array (for data and shape and flags)

We may be thinking of much the same, but I would like NumPy sort to deal with contiguous copies, result allocation and outer iteration, so that the (user) DType only needs to provide the lowlevel sorting implementation. (I.e. the same pattern that ufuncs/casts use)

I would like to achieve that anything that a user DType may provide doesn't take arrays. NumPy DTypes need no knowledge of arrays (except for some old API, one of which we are replacing here).

I suppose I am just repeating that DType->get_(arg)sortfunction() is what we should add (and this adds it, with the signature PyArrayDTypeMeta_Get(Arg)SortFunction).
I.e. I think this is just the right abstraction and mirrors ufuncs/casts, that is the DType (class/type) provides the lowlevel sorting implementation (and we give the DType the chance to do setup/have multiple, which isn't often needed but is useful for structured dtypes for example).

Although, I would be OK to just provide a DType->lowlevel_sort_implementation instead, because we can always add the above to when we want (and most DTypes don't need the get_ version. Right now no dtype needs it).

But beyond that, e.g. whether we should have a CompareFunction, I really don't have a strong opinion.

Note that with included dtype, only one dispatcher is needed for all dtypes. Or we could omit the dtype and attach different dispatchers to each dtype.

We should always include the dtype (i.e. PyArray_Descr), for example StringDType needs it to work with its elements and structured dtypes (e.g. in casts) also use it for specialization.

In a sense I would love if we could implement sorting as (simple) gufunc even if that is overcomplicated for sorting... Unfortunately, we don't have a defined way to pass the parameters like descending.

@seberg Your proposal sounds like the current implementation, but with more slots for the low level functions. As I see it, we would like eight slots rather than the current three, that is, three booleans:

• ascending/decending
• stable/possibly not stable
• nan last/first

We currently only really use two: stable/possibly not stable. Note that we already have shims for functions that require workspace, the shims allocate/release working memory.

The question then becomes how to allow extensions without making disruptive changes. I suggest that the easiest way is to have a pointer to a structure dedicated to holding function pointers rather than a function that returns the wanted pointer. The only complication is that the function signature is different between the generic/builtin cases, but I suppose that can be handled by a NULL in the comparison function arguments.

If we are not worried about future extensions, we could just add eight slots, sixteen total with sort/argsort.

What is it you would like with dtype inheritance?

I am interested in what users might like, @mhvk and pandas might have some ideas.

As I see it, we would like eight slots rather

I would like to just add one slot actually (please forget about the point that we don't have to use that). I.e. the function (a second one for argsort):

def PyArrayDTypeMeta_GetSortFunction(sort_context, ...):
    return lowlevel_sort_implementation

Now, internally, the dtype may or may not have 8 different versions. Because the new lowlevel_sort_implementation will also get the sort_context passed, so in principle you could choose to merge the functions into a single version if it doesn't matter speed wise.

Why the GetSortFunction pattern?

1. It mirrors ufuncs and casts.
2. It allows us to have the 8 versions without actually adding 8 slots explicitly.
3. For some dtypes it can make sense to specialize things and that works well here. You could of course do that on the first call in principle (but, this mirrors ufuncs and casts).

Now, passing sort_context rather than stable, descending, nan_order is totally fine by me, I guess for PyArrayDTypeMeta_GetSortFunction it seemed fine to not bunch it up for convenience (and because I didn't expect it to grow any time soon).

If we are not worried about future extensions, we could just add eight slots, sixteen total with sort/argsort.

We are not (overly) worried about the future, because these are not the same kind of slots as the old ones. These are slots on the DType class the PyArray_DTypeMeta and we can add and deprecate these slots basically as freely as in Python!

To stress this: Even if we add PyArrayDTypeMeta_GetSortFunction today, we can easily add a PyArrayDTypeMeta_GetSortFunctionWithMoreOptions in the next version and keep both around or deprecate the other one slowly.

So we should be comfortable with the choice for the mid-term, but we don't have to worry too much in case someone comes around and wants sorts that e.g. work with non-contiguous data (if we don't really expect that to happen).

Adding nan last/first and descending now, is because there is clear interest for it in the near or mid term (e.g. because of top_k, although NaN we could also just define "always to end", I like the fact that having it around also documents the expectation).

There are going to effectively be 8 slots even with the function + context, it is a matter of how they are accessed. I'd rather they be in a table instead of in a switch. That said, I'll note that the improvement in speed between the original quicksort implementation and the type specific one was about 5x. That was due to several things, first, builtin comparisons, and second, a cleaner and simpler version of quicksort.

There are going to effectively be 8 slots even with the function + context, it is a matter of how they are accessed

True, although some of them may be duplicate entries (if perfect specialization isn't worth it or unnecessary e.g. for integers which can ignore NaNs).

I'd rather they be in a table instead of in a switch

I don't mind that too much the slot could be a struct with 8 entries. Although, I prefer this because of the reasons I said above:

1. It is how ufuncs/casts are designed and I think it gives more flexibility to user DTypes (this is for the public API). Even if a bit more verbose.
2. It is much more obvious to do setup work there, which could be useful e.g. if anyone ever wanted to optimize structured dtype comparisons.
3. (In the extreme case a future user like numba could even JIT a function.)

But yeah, it may be that that we repeat the pattern to dispatch to the specialized version a few more times with my preference, because it could be repeated for multiple DTypes rather than just once.
(I guess that may depend a bit on how we actually do it internally, we could use the same function for all our Numerical dtypes after all and use a static struct internally in principle.)

That said, I guess we could always add my GetSortFunction slot later if I want to, so in that sense, I shouldn't care about what approach we use, and if you prefer slots a lot (and nobody chimes in), I am very happy to go that way @charris.

Slots may have the advantage of being slightly easier to reason about for user DType authors as well. (i.e. for ufuncs, the user doesn't need to use the get_loop slot either, it's optional)

For the public API it seems easier and clearer (in intent) to me to use a function rather than a struct. I'm not too familiar, but I would think it is possible for compilers to optimize, say, NaN order and descending, if the user writes obvious fast paths; in such cases it may be simpler to allow the user to merge functions than write 8 wrappers or versions (and as @seberg said this is similar to ufuncs too). And not all dtypes would care (equally) about all options (especially nan). Also, this seems more consistent with the slot design we have so far.

Also I think it would be nice (if @charris meant that as well) to expose some of the internal sort machinery we have and in that case user dtypes can reuse it more easily with a single function, presumably, if we wrap things (efficiently) (though this ties into the discussion whether we need to generalize compare, which I'm not sure about anymore...)

To be clear I'm not against going with either approach if there are strong opinions, just thinking out loud! Happy to go with whatever everyone thinks makes sense :)

To be clear, I wasn't recommending actually adding 8 (16) slots to the dtype, that was just to illustrate the effective functionality. What I would suggest is adding a pointer to a structure instead of a function. Yes, some of the entries could be duplicates or NULL. For instance the only reason to maybe prefer quicksort over radix sort for short integers is that radix sort requires extra memory. The only difference from a function is using an index rather than a context. I suspect a function might want to use a table for implementation, converting the context to an index. So I guess it is just a matter of the easiest way to pass the selection information.

I would suggest either new PyArray_SortExt and PyArray_ArgSortExt functions that take more/different arguments, context maybe, or simply extending NPY_SORTKIND. We will want new sorting functions for the existing types that do not use the generic sorting functions, but rather builtin comparison functions. Such functions are easy to add, and it looks like the SIMD sorting functions can do descending sorts, not sure what they do for nans.

For descending sorts, we could just reverse the output of existing sorts as a quick, if dirty, solution. That would also work for floats if we added the option of sorting nans to the beginning.

For descending sorts, we could just reverse the output of existing sorts as a quick, if dirty, solution

Yeah, that's a good point. If we provide the family of functions (yes, via a single slot as you said), then we can auto-fill half the slots via reversing others.

So in the end, I think where we are heading is basically adding a sorting-slots struct. That was maybe always the right thing to do (even for the fewer ones in the PR, I guess I just didn't need it as much before; the one indirection should be meaningless).

Since I do like the idea of potentially adding a GetSortFunction(), I think we could do the following single slot:

struct {
    sortfunc *sort_ascending_stable;
    sortfunc *sort_ascending_unstable;
    sortfunc *sort_descending_stable;
    sortfunc *sort_descending_unstable;
    ...
    /* and argsorts */
    ...
    void *reserved1 = NULL;  // for a GetSortFunction probably
    void *reserved2 = NULL;  // GetArgsort
    /* could reserve more if we want a compare function, but no need to worry */
}

I don't have an opinion on how many sort functions we include. 4+4 we need, but maybe all 8+8 (as the NaN versions can be implemented via reversing if descending sorts to end).
I would personally be OK with relying on GetSortFunction for the full zoo, but I think you prefer to avoid it which is fine by me.

Been thinking about this some more, The current logic is:

• Look for function with builtin compare (in table)
• If no function (table entry is NULL, use generic sort with descr compare function.

That is pretty much what this PR does AFAICT.

What I think we need is to have more sorting options, i.e., decending, maybe nan position. To make those accessible, we need to either extend NPY_SORTKIND or add a new C-API function with appropriate arguments. I think that is the first decision to be made. Once that is decided, we should implement that for all current sort types and use it internally. That will allow the current tests to be used to check the infrastructure. That is steo one. The second step (new PR) would add descending builtin sorts, with extended tests. That could be broken down into several steps, maybe doing only generic/stable/etc at each point.

Thinking about descending sorts, stable sorts cannot be implemented by reversal. Indeed, I think we need to explicitly define what "stable" means for descending sorts, which should mean the equal elements are sorted in index order.

Come to think of it, reversal before sorting followed by reversal after should work in the stable case.

Table pointer or function is an implementation detail, I could go with either.

I think we should be able to get descending sorts added for 2.4.

Yes, that is the current logic! In order of priority:

1. If a get_(arg)sort_function is found, that is used,
2. If it is not found and a sort_compare is found, that is used with the new built-in functions,
3. If that is not found, and the older arrfunc slot ((arg)sort), that is used (this includes SIMD implementations currently, it seems),
4. Failing all that, we rely on the compare arrfunc and use the older built-in functions.

What I think we need is to have more sorting options

Just to be clear, for 1 and 2, the sort methods are restructured and nan_position and descending are actually functional! Those are retrieved from the context. It isn't completely redone, so it's patchy, but there is the basic infrastructure. This is done by patching in a compare_with_context function that replaces the compare arrfunc - it acts more like a 'should swap' function than a compare function, and reacts accordingly to these two options e.g., if the sort is descending, the effective comparison is flipped. Only if we get to 3 and 4 those options aren't available anymore, and context goes unused.

What is not done yet is the sort implementations for built-in dtypes, like numpy integers or floats - those are also implemented in SIMD AFAIK. We only have public API so far; none of this is in use internally, so yes, that should be next so we can test it here. This basically amounts to replacing options 3 and 4 with 1 and 2 for all dtypes. This should be possible by initializing the internal dtypes with the get_(arg)sort_function and (if needed) sort_compare slots.

How do you access those from the public C-API, PyArray_Sort doesn't have the arguments ...

Should have been precise, sorry, functional in terms of the actual sorting routines, the levels down from _new_(arg)sortlike, not at the level of PyArray_(Arg)Sort. By public API I meant the dtype author side; public user dtype authors can define sorts for their dtypes using the arguments in the context. I think functional is not the right term then, just wanted to clarify a default context is being created and used:

https://github.com/numpy/numpy/pull/28516/files#diff-e469e991ad3c24ac159bcc0469153ebb85fbed6f81c937c07e126f001a1007f7R1219

We do need to add the actual options as you said.

As I understand things, the first commit 87f410b is all that is needed to actually solve the slow string sorting. Perhaps it would be worthwhile to have that as a separate PR to solve the immediate problem, together with some benchmarks, until the dust settles around the wider refactor?

Moved this PR to #29737 given it would need major changes after #29642. Closing this - thanks everyone.