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 all commits
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
7 changes: 4 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 @@ -127,9 +127,10 @@ Item selection and manipulation:

* :meth:`~numpy.ndarray.take`
* :meth:`~numpy.ndarray.repeat` (no axis argument)
* :meth:`~numpy.ndarray.sort` (without arguments)
* :meth:`~numpy.ndarray.sort` (``kind`` key word argument supported for
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 @@ -1106,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
8 changes: 8 additions & 0 deletions docs/upcoming_changes/10703.np_support.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
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``. 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.
19 changes: 17 additions & 2 deletions numba/core/typing/arraydecl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -462,8 +462,20 @@ def sentry_shape_scalar(ty):
@bound_function("array.sort")
def resolve_sort(self, ary, args, kws):
assert not args
assert not kws
return signature(types.none)
kwargs = dict(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()]))
def sort_stub(kind='quicksort'):
pass
pysig = utils.pysignature(sort_stub)
return signature(types.none, kind).replace(pysig=pysig)

@bound_function("array.argsort")
def resolve_argsort(self, ary, args, kws):
Expand All @@ -472,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
28 changes: 21 additions & 7 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 @@ -7011,18 +7014,29 @@ def lt_implementation(dtype):
return default_lt


@lower_builtin("array.sort", types.Array)
@lower_builtin("array.sort", types.Array, types.StringLiteral)
def array_sort(context, builder, sig, args):
arytype = sig.args[0]
arytype, kind = sig.args
kind = kind.literal_value

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

def array_sort_impl(arr):
# Note we clobber the return value
sort_func(arr)
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
for idx in np.ndindex(arr.shape[:-1]):
sort_func(arr[idx])
else:
def array_sort_impl(arr):
# Note we clobber the return value
sort_func(arr)

return context.compile_internal(builder, array_sort_impl, sig, args)
innersig = sig.replace(args=sig.args[:1])
innerargs = args[:1]
return context.compile_internal(builder, array_sort_impl,
innersig, innerargs)


@overload(np.sort)
Expand Down
145 changes: 142 additions & 3 deletions numba/tests/test_sort.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,12 @@ def make_temp_array(keys, n):
def sort_usecase(val):
val.sort()

def sort_kind_usecase(val, is_stable=False):
if is_stable:
val.sort(kind='mergesort')
else:
val.sort(kind='quicksort')

def argsort_usecase(val):
return val.argsort()

Expand Down Expand Up @@ -846,11 +852,13 @@ def has_duplicates(self, arr):
return True
return False

def check_sort_inplace(self, pyfunc, cfunc, val):
def check_sort_inplace(self, pyfunc, cfunc, val, kwargs=None):
if kwargs is None:
kwargs = {}
expected = copy.copy(val)
got = copy.copy(val)
pyfunc(expected)
cfunc(got)
pyfunc(expected, **kwargs)
cfunc(got, **kwargs)
self.assertPreciseEqual(got, expected)

def check_sort_copy(self, pyfunc, cfunc, val):
Expand Down Expand Up @@ -900,6 +908,50 @@ def test_array_sort_complex(self):
orig = np.array([complex(*x) for x in zip(real, imag)])
self.check_sort_inplace(pyfunc, cfunc, orig)

def test_array_sort_kind_int(self):
def check(pyfunc, is_stable):
cfunc = jit(nopython=True)(pyfunc)
for orig in self.int_arrays():
self.check_sort_inplace(pyfunc, cfunc, orig,
dict(is_stable=is_stable))

check(sort_kind_usecase, is_stable=True)
check(sort_kind_usecase, is_stable=False)

def test_array_sort_kind_float(self):
def check(pyfunc, is_stable):
cfunc = jit(nopython=True)(pyfunc)
for orig in self.float_arrays():
self.check_sort_inplace(pyfunc, cfunc, orig,
dict(is_stable=is_stable))

check(sort_kind_usecase, is_stable=True)
check(sort_kind_usecase, is_stable=False)

def test_array_sort_kind_complex(self):
def check(pyfunc, is_stable):
cfunc = jit(nopython=True)(pyfunc)
for real in self.float_arrays():
imag = real[::]
np.random.shuffle(imag)
orig = np.array([complex(*x) for x in zip(real, imag)])
self.check_sort_inplace(pyfunc, cfunc, orig,
dict(is_stable=is_stable))

check(sort_kind_usecase, is_stable=True)
check(sort_kind_usecase, is_stable=False)

def test_array_sort_kind_multidim(self):
def check(pyfunc, is_stable):
cfunc = jit(nopython=True)(pyfunc)
for shape in ((4, 5), (3, 4, 5), (2, 30)):
orig = np.random.random(shape) * 100
self.check_sort_inplace(pyfunc, cfunc, orig,
dict(is_stable=is_stable))

check(sort_kind_usecase, is_stable=True)
check(sort_kind_usecase, is_stable=False)

def test_np_sort_int(self):
pyfunc = np_sort_usecase
cfunc = jit(nopython=True)(pyfunc)
Expand Down Expand Up @@ -1228,6 +1280,56 @@ def external_key(z):
self.assertIn(expect, str(raises.exception))


class TestArraySort(MemoryLeakMixin, TestCase):
"""Tests specific to array.sort"""

def test_exceptions(self):

@njit
def nonliteral_kind(kind):
np.arange(5).sort(kind=kind)

# check non-literal kind
with self.assertRaises(errors.TypingError) as raises:
# valid spelling but not literal
nonliteral_kind('quicksort')

expect = '"kind" must be a string literal'
self.assertIn(expect, str(raises.exception))

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

can you add a test here that would check valid sort kinds and invalid kind string literals ?

@eyupcanakman eyupcanakman Sep 1, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Added test_kinds for the three accepted spellings and a heapsort case in test_exceptions.

@njit
def unsupported_kwarg():
np.arange(5).sort(foo='')

with self.assertRaises(errors.TypingError) as raises:
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 @@ -1255,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