Repository navigation
Add mergesort option for array.sort() - #10703
Conversation
`array.sort()` only supported `quicksort`, unlike `array.argsort()` which already accepted a `kind` argument. Handle `kind` in `resolve_sort()` the same way `argsort` does, so `array.sort(kind='mergesort')` works. Arrays of more than one dimension sort each last-axis slice, matching `quicksort`.
bfa6b5e to
902b12b
Compare
swap357
left a comment
swap357
left a comment
There was a problem hiding this comment.
Thank you @eyupcanakman, for the PR! I sense you're following the approach suggested here #9779 (comment) to follow array_argsort() reference. There is a gap however in value validation which is only caught much later with unclear error. I've pointed it out below, hope it helps.
| kwargs = dict(kws) | ||
| kind = kwargs.pop('kind', types.StringLiteral('quicksort')) | ||
| if not isinstance(kind, types.StringLiteral): | ||
| raise TypingError('"kind" must be a string literal') |
There was a problem hiding this comment.
kind is only checked for being a string literal here, so any spelling gets through to get_sort_func(). Now there is handling for quicksort and mergesort, but anything else valid or invalid just fails with -
UnboundLocalError: cannot access local variable 'func' where it is not associated with a value
>>> from numba import jit
>>> import numpy as np
>>> a = np.array([4,3,2,1])
>>> a.sort(kind='stable')
>>> a
array([1, 2, 3, 4])
>>> a = np.array([4,3,2,1])
>>> def f(a):
... a.sort(kind='stable')
... return a
...
>>> clear
>>> from numba import jit
>>> import numpy as np
>>> def f(a):
... a.sort(kind='stable')
... return a
...
>>> x = np.array([4,3,2,1])
>>> x
array([4, 3, 2, 1])
>>> f(x)
array([1, 2, 3, 4])
>>> x
array([1, 2, 3, 4])
>>> nb_f = jit(f)
>>> nb_f(x)
Traceback (most recent call last):
File "/Users/swap357/Documents/dev/numba/numba/np/arrayobj.py", line 6964, in get_sort_func
return _sorts[key]
~~~~~~^^^^^
KeyError: ('stable', 'default_lt', False)
During handling of the above exception, another exception occurred:
...
UnboundLocalError: cannot access local variable 'func' where it is not associated with a value
From the numpy docs, I gather that stable is mergesort alias -
Note that both ‘stable’ and ‘mergesort’ use timsort or radix sort under the covers and, in general, the actual implementation will vary with data type. The ‘mergesort’ option is retained for backwards compatibility.
So 'stable' is likely input and we should have proper handling for it. ('heapsort' too) We can reject the not implemented ones. But since we have 'mergesort', we might as well use 'stable' alias.
There was a problem hiding this comment.
maybe you can add value validation after string literal check ?
|
|
||
| expect = '"kind" must be a string literal' | ||
| self.assertIn(expect, str(raises.exception)) | ||
|
|
There was a problem hiding this comment.
can you add a test here that would check valid sort kinds and invalid kind string literals ?
There was a problem hiding this comment.
Added test_kinds for the three accepted spellings and a heapsort case in test_exceptions.
|
Could you also merge latest |
# Conflicts: # docs/source/reference/numpysupported.rst
Unknown kinds fell through `get_sort_func()` and raised `UnboundLocalError` for both `sort()` and `argsort()`. Reject them at typing time and treat `'stable'` as an alias for `'mergesort'`, as NumPy does. Assisted-by: Claude Code
|
Merged main and added the check in |
swap357
left a comment
swap357
left a comment
There was a problem hiding this comment.
The changes look good. Thank you, @eyupcanakman for all the efforts! Appreciate it
swap357
added
5 - Ready to merge
3 participants
array.sort()only supported quicksort, unlikearray.argsort()which already accepted akindargument. This adds the same handling toresolve_sort(), soarray.sort(kind='mergesort')now works, with quicksort still the default.Quicksort already sorts arrays of more than one dimension along the last axis, and mergesort now does the same by sorting each last-axis slice.
Only the sort method accepts
kind.np.sort()is unchanged, matching the issue.Closes #9779