Skip to content

Add mergesort option for array.sort() #10703

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
swap357 merged 3 commits into from
Sep 3, 2026
Merged

Add mergesort option for array.sort() #10703

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
902b12b
Add mergesort option for array.sort()
eyupcanakman Jul 13, 2026
e68b9cb
Merge remote-tracking branch 'origin/main' into feat/array-sort-kind
eyupcanakman Sep 1, 2026
6750266
Validate the sort kind and accept 'stable'
eyupcanakman Sep 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Validate the sort kind and accept 'stable' 
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
  • Loading branch information
@eyupcanakman
eyupcanakman committed Sep 1, 2026
commit 6750266745b65a9b71b3854491e4a1d8ba60f461
6 changes: 3 additions & 3 deletions docs/source/reference/numpysupported.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -128,9 +128,9 @@ Item selection and manipulation:
* :meth:`~numpy.ndarray.take`
* :meth:`~numpy.ndarray.repeat` (no axis argument)
* :meth:`~numpy.ndarray.sort` (``kind`` key word argument supported for
values ``'quicksort'`` and ``'mergesort'``)
values ``'quicksort'``, ``'mergesort'`` and ``'stable'``)
* :meth:`~numpy.ndarray.argsort` (``kind`` key word argument supported for
values ``'quicksort'`` and ``'mergesort'``)
values ``'quicksort'``, ``'mergesort'`` and ``'stable'``)
* :meth:`~numpy.ndarray.nonzero`

Where applicable, the corresponding top-level NumPy functions (such as
Expand Down Expand Up @@ -1107,7 +1107,7 @@ Sorting, searching, and counting
* :func:`numpy.sort` (no optional arguments, quicksort accepts
multi-dimensional array and sorts its last axis).
* :func:`numpy.argsort` (``kind`` key word argument supported for values
``'quicksort'`` and ``'mergesort'``)
``'quicksort'``, ``'mergesort'`` and ``'stable'``)
* :func:`numpy.argwhere`
* :func:`numpy.count_nonzero` (axis only supports scalar values)
* :func:`numpy.extract`
Expand Down
5 changes: 3 additions & 2 deletions docs/upcoming_changes/10703.np_support.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,5 +3,6 @@ Added ``kind`` argument to ``ndarray.sort``

``ndarray.sort`` now accepts the ``kind`` argument in nopython mode and
supports ``"mergesort"`` in addition to the default ``"quicksort"``, matching
the ``kind`` support already available in ``ndarray.argsort``. The default
value keeps the existing behaviour.
the ``kind`` support already available in ``ndarray.argsort``. Both methods
also accept ``"stable"`` as an alias for ``"mergesort"`` and reject any other
``kind`` at compile time. The default value keeps the existing behaviour.
6 changes: 6 additions & 0 deletions numba/core/typing/arraydecl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -466,6 +466,9 @@ def resolve_sort(self, ary, args, kws):
kind = kwargs.pop('kind', types.StringLiteral('quicksort'))
if not isinstance(kind, types.StringLiteral):
raise TypingError('"kind" must be a string literal')

@swap357 swap357 Aug 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@swap357 swap357 Aug 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe you can add value validation after string literal check ?

if kind.literal_value not in ('quicksort', 'mergesort', 'stable'):
msg = 'Unsupported "kind": {!r}'
raise TypingError(msg.format(kind.literal_value))
if kwargs:
msg = "Unsupported keywords: {!r}"
raise TypingError(msg.format([k for k in kwargs.keys()]))
Expand All @@ -481,6 +484,9 @@ def resolve_argsort(self, ary, args, kws):
kind = kwargs.pop('kind', types.StringLiteral('quicksort'))
if not isinstance(kind, types.StringLiteral):
raise TypingError('"kind" must be a string literal')
if kind.literal_value not in ('quicksort', 'mergesort', 'stable'):
msg = 'Unsupported "kind": {!r}'
raise TypingError(msg.format(kind.literal_value))
if kwargs:
msg = "Unsupported keywords: {!r}"
raise TypingError(msg.format([k for k in kwargs.keys()]))
Expand Down
8 changes: 6 additions & 2 deletions numba/np/arrayobj.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6982,6 +6982,9 @@ def get_sort_func(kind, lt_impl, is_argsort=False):
"""
Get a sort implementation of the given kind.
"""
if kind == 'stable':
# NumPy treats 'stable' as an alias for 'mergesort'
kind = 'mergesort'
key = kind, lt_impl.__name__, is_argsort

try:
Expand Down Expand Up @@ -7014,11 +7017,12 @@ def lt_implementation(dtype):
@lower_builtin("array.sort", types.Array, types.StringLiteral)
def array_sort(context, builder, sig, args):
arytype, kind = sig.args
kind = kind.literal_value

sort_func = get_sort_func(kind=kind.literal_value,
sort_func = get_sort_func(kind=kind,
lt_impl=lt_implementation(arytype.dtype))

if kind.literal_value == 'mergesort' and arytype.ndim > 1:
if kind in ('mergesort', 'stable') and arytype.ndim > 1:
def array_sort_impl(arr):
# mergesort only sorts 1D arrays, so sort each slice along
# the last axis as quicksort does for multidimensional arrays
Expand Down
59 changes: 59 additions & 0 deletions numba/tests/test_sort.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1307,6 +1307,28 @@ def unsupported_kwarg():
expect = "Unsupported keywords: ['foo']"
self.assertIn(expect, str(raises.exception))

@jit
def unsupported_kind():
np.arange(5).sort(kind='heapsort')

with self.assertRaises(errors.TypingError) as raises:
unsupported_kind()

expect = "Unsupported \"kind\": 'heapsort'"
self.assertIn(expect, str(raises.exception))

def test_kinds(self):
for kind in ('quicksort', 'mergesort', 'stable'):
@jit
def sort_kind(arr):
arr.sort(kind=kind)

for arr in (np.array([3, 1, 4, 1, 5, 9, 2, 6]),
np.array([[3, 1, 4, 1], [5, 9, 2, 6]])):
expected = np.sort(arr, kind=kind)
sort_kind(arr)
self.assertPreciseEqual(arr, expected)


class TestArrayArgsort(MemoryLeakMixin, TestCase):
"""Tests specific to array.argsort"""
Expand Down Expand Up @@ -1335,6 +1357,43 @@ def unsupported_kwarg():
expect = "Unsupported keywords: ['foo']"
self.assertIn(expect, str(raises.exception))

@jit
def unsupported_kind():
np.arange(5).argsort(kind='heapsort')

with self.assertRaises(errors.TypingError) as raises:
unsupported_kind()

expect = "Unsupported \"kind\": 'heapsort'"
self.assertIn(expect, str(raises.exception))

def test_kinds(self):
arr = np.array([3, 1, 4, 7, 5, 9, 2, 6])
for kind in ('quicksort', 'mergesort', 'stable'):
@jit
def argsort_kind(arr):
return arr.argsort(kind=kind)

expected = np.argsort(arr, kind=kind)
self.assertPreciseEqual(argsort_kind(arr), expected)

# ties and more than SMALL_QUICKSORT elements, so only a stable
# sort reproduces numpy's index order
arr = np.array([3, 2, 2, 1, 1, 0, 0, 0, 0, 3,
2, 3, 2, 2, 3, 2, 2, 2, 2, 3])
expected = np.argsort(arr, kind='stable')
for kind in ('mergesort', 'stable'):
@jit
def argsort_stable(arr):
return arr.argsort(kind=kind)

@jit
def np_argsort_stable(arr):
return np.argsort(arr, kind=kind)

self.assertPreciseEqual(argsort_stable(arr), expected)
self.assertPreciseEqual(np_argsort_stable(arr), expected)


if __name__ == '__main__':
unittest.main()
Loading