ENH: Descending partitioning by MaanasArora · Pull Request #31511 · numpy/numpy

MaanasArora · 2026-05-26T10:43:23Z

Draft to add a descending kwarg to np.partition and np.argpartition, analogous to the sorting functions, thus addresses part of #31424. Not sure if the SELECTKIND changes constitute an ABI or C-API break, but analogous to the sorts as well. cc @seberg, thanks!

Note that this does not create new-style ArrayMethod registration analogous to the sorts, but rather updates the (apparently fully internal) partition_map and continues using the old-style PartitionFunc. (In fact, the PyArray_(Arg)Partition functions are unchanged!) Hopefully this is minimal enough and we can just add new-style ArrayMethods (and thus enable user-dtypes to use partitions) in follow-up? (This would then also unblock work for top_k.)

Still needs docs, tests, and a release note!

AI Disclosure

No AI used.

seberg

Note that this does not create new-style ArrayMethod registration analogous to the sorts, but rather updates the (apparently fully internal) partition_map and continues using the old-style PartitionFunc

I would like to introduce the new ArrayMethod approach, but currently this is all internal so I am very happy to start with this since it also splits things up into smaller PRs.
(Plus, may help to work on top-k in parallel.)

But yeah, needs tests and docs for sure :).

(Otherwise, I think kind+descending should be fine, except that kind is only a single option anyway, so it's a bit funny and we should maybe just have stable/unstable as for the sorts, even if that is even less interesting here probably.
)

seberg · 2026-05-28T09:24:56Z

+        NPY_INTROSELECT = 0,
+        // new style names
+        NPY_SELECT_DEFAULT = 0,
+        NPY_SELECT_DESCENDING = 1,


Hmmm, I wonder if we should just also make it 4, so that nsorts and nselects might just be identical (no point right now, I guess).

Thanks, yeah, right now the kind is selected by its position in the array... not great, but would go away when we put in ArrayMethods anyway!

Right, but this is public API, so we can't defer this if that is what you mean? (and if that means changing the other code, so be it).

Right, yeah, we should probably sort this out just now. I've made the enum effectively identical and changed the other code (I assume we're not actually adding the kwarg yet.) Not sure what to do with the NPY_NSELECTS...

MaanasArora · 2026-05-29T05:33:59Z

I would like to introduce the new ArrayMethod approach, but currently this is all internal so I am very happy to start with this since it also splits things up into smaller PRs.

Thanks for reviewing! Yes, this seems to be a good way to break the PRs, and I definitely will follow-up very soon with ArrayMethods.

I added tests, will add docs soon! I found a bug where dtypes were silently falling to (default) sort even with descending, so I disabled that (TypeErrors now). But I realized strings and dates even actually don't use introselect; they fell back to sorting, which can easily be replicated for descending with ArrayMethods. I'm trying to see if dates can actually be conveniently mapped to the existing introselect (given we have the nice new comparisons), but have just disabled the tests for this push. If that doesn't work, can keep them disabled, and strings should probably be disabled for this PR anyway (the introselect would have to duplicated for strings otherwise I think, which can be done in follow-up).

Edit: I was able to route datetime/timedelta to the existing introselect, and the tests are also enabled now!

we should maybe just have stable/unstable as for the sorts, even if that is even less interesting here probably.

Agreed having stable could be nice for consistency! Maybe would be nice to add another select algorithm but stable, from which we can fallback to stable sorts.

MaanasArora · 2026-05-30T02:48:06Z

Added docs, release note, and fixed up the ma stubs to align! I'll think a bit more about the angle for the docs, but this should be ready to review, thanks!

seberg

LGTM, thanks. A few smaller things, but most comments are just observations/thoughts really.
I might think about the docs a bit more. I wonder if it is clearer if we just keep the old text with smaller/greater even.
(And if the clarification at the end like "Sorting the partitions split by k-th individually would yield a fully sorted array." might be helpful.)

seberg · 2026-05-31T10:31:11Z

            import _testbuffer
        except ImportError:
-            raise pytest.skip("_testbuffer is not available")
+            raise pytest.mark.skip("_testbuffer is not available")


Suggested change

raise pytest.mark.skip("_testbuffer is not available")

pytest.skip("_testbuffer is not available")

Also above, pytest.skip() raises the SkipException. Although, in the first, we might as well use a @pytest.mark.skipif() decorator and in this one a _testbuffer = pytest.importskip("_testbuffer")

(I am not clear why this is suddenly showing up, but so be it.)

Thanks, this is done! (For the first one, not sure if the decorator would work, since there an argument is part of the condition)

seberg · 2026-05-31T10:31:24Z

        axis: SupportsIndex = -1,
-        kind: _PartitionKind = "introselect",
-        order: None = None,
+        kind: _PartitionKind | None = None,        order: None = None,


Sorry, fixed!

seberg · 2026-05-31T10:32:59Z

+
+    if (sortkind == _NPY_SELECT_UNDEFINED) {
+        // keywords only if sortkind not passed
+        sortkind = descending > 0? NPY_SELECT_DESCENDING: NPY_SELECT_DEFAULT;


Suggested change

sortkind = descending > 0? NPY_SELECT_DESCENDING: NPY_SELECT_DEFAULT;

sortkind = descending > 0 ? NPY_SELECT_DESCENDING : NPY_SELECT_DEFAULT;

(but just silly nit)

Fixed, thanks!

seberg · 2026-06-01T07:29:20Z

            {"|kind", &PyArray_SelectkindConverter, &sortkind},
-            {"|order", NULL, &order}) < 0) {
+            {"|order", NULL, &order},
+            {"|descending", &PyArray_OptionalBoolConverter, &descending}) < 0) {


Suggested change

{"|descending", &PyArray_OptionalBoolConverter, &descending}) < 0) {

{"$descending", &PyArray_OptionalBoolConverter, &descending}) < 0) {

Probably might as well make it kwonly here as well.

seberg · 2026-06-01T07:48:08Z

        specified, but unspecified fields will still be used, in the
        order in which they come up in the dtype, to break ties.
+    descending : bool, optional
+        Partition order. If ``True``, the left partition would contain


I wonder if we should just start with "Sort order" (feels like it may be simpler).

Yeah, that's better, this is done - thanks!

seberg · 2026-06-01T07:54:58Z

+        indices of elements less than or equal to the k-th element and the
+        right partition will contain indices of elements greater than or
+        equal to the k-th element. If ``False`` or ``None``, the left partition
+        will contain indices of elements less than or equal to the k-th element


Whoops, this is reversed. But I am wondering if we can make it a bit shorter/closer to the sort docs and maybe even just say "partitioned in ascending/descending order" and rely on the introduction documentation to be clear enough on it?

Sorry, the reversal is fixed! Shortening the docs would be nice. I think the issue is really that this is moved from the kth docs to here, and it's not natural. Perhaps we can move this to k-th, but it is not really relevant there. Maybe just like the original, a single smaller/greater wording is better, but it's not really correct for descending?

seberg · 2026-06-01T07:58:55Z

-    are located to the left of this element and all equal or greater are
-    located to its right. The ordering of the elements in the two partitions
-    on the either side of the k-th element in the output array is undefined.
+    Partially sorts the array in such a way that the value of the element in k-th


Suggested change

Partially sorts the array in such a way that the value of the element in k-th

Partially sorts the array in such a way that the value of the element in the k-th

I wonder if we should remove "the value of". Since I think we can just say k-th element and the other docs do (but this was there before).

Article fixed, thanks! Yeah, I think we should remove it (missed in this push sorry). Let me do that.

MaanasArora · 2026-06-01T08:51:07Z

Sorry, silly errors, thanks for the thorough review! Aside from the docs, I guess the big thing is checking astropy and pandas are fine with the default change for kind, I'll get started on that! (A complex nan mixing test would also be nice.)

MaanasArora · 2026-06-02T01:47:56Z

Pushed another iteration of the docs, not sure if it turned out much better with the parentheses, but the descending entry is simpler at least.

For the kind interop issue: Astropy refers to partitioning in its MaskedNDArray class, but doesn't actually seem to invoke np.(arg)partition, just full sorts, so the argument isn't being passed I think. Pandas seems to only use a partition method in the context of string partitioning, so perhaps not relevant at all there either? (I'm a bit surprised, though.) Test suites passed for both projects with dev NumPy installation on this branch.

seberg · 2026-06-02T13:23:07Z

Cool, I am happy to not worry too much about it then. Although, I think it may still make sense to just go with np._NoValue to not pass it on to the method at all (rather than None) -- will probably need to build a kwarg dict, but so be it...
Not sure if that is slightly wrong for the descending sorts, but should't the types also use NoValue?

MaanasArora · 2026-06-03T00:09:46Z

Thanks, I switched to np._noValue, and used a dict for the kwargs! Also, added a test for complex with nan mixing.

Not sure if that is slightly wrong for the descending sorts, but should't the types also use NoValue?

Yes, that seems right! Going to make that change, thanks. (np._NoValue itself doesn't seem to occur anywhere in the types, but _NoValueType does!)

seberg

One small round of nits, but we should get this done. I think the new docs read quite well so happy there.
The things around None default, etc. mirrors sorts, so I don't want to think about the best approach _NoValue use is good now.

My agent had a few tiny test suggestions, but I don't think any of them need to be fixed particularly.

seberg · 2026-06-03T13:41:22Z

+        NPY_INTROSELECT = 0,
+        // new style names
+        NPY_SELECT_DEFAULT = 0,
+        NPY_SELECT_DESCENDING = 1,


Right, but this is public API, so we can't defer this if that is what you mean? (and if that means changing the other code, so be it).

seberg · 2026-06-03T13:58:50Z

-        indexed by kth of them into their sorted position at once.
+        Element index to partition by. The k-th value of the array will
+        be in the position it would be in a sorted array, all elements that
+        sorted array, all elements that are less than this element (or


Docs got a little bit garbled here. "... that sorted array".

Fixed, thanks!

seberg · 2026-06-03T14:08:24Z

+
+        no_nans = arange + 1j * arange
+        im_nans = arange + 1j * nans
+        re_nans = nans + 1j * arange


My agent actually noticed that this doesn't do what it says it does (maybe not new), this was the reason why I mentioned complex(..., ...) in a sense.

Just replacing 1j * nans with complex(0, np.nan) should do the trick (and the last nans just with np.nan or complex(np.nan, 0) just to be obvious)

Ah yeah, that's right, thanks - I fixed this, but this actually unearthed a bad bug with the complex comparators (they were always a bit broken, but descending exposed it more)... I'm going to make a PR to fix that (as it probably should be backported), so I didn't update the sorting test here, but doing so in the other PR. Sorry it got missed in the descending PR.

seberg

Thanks, one more comment on the kind. I am not quite sure if we should write descending=None or NoValue (in the docs), but we support None now, and it reads maybe better so I am happy with it.

I think the rest is all good. It would be nice to change it to just use the correct full sort (rather than the cmp function), but I think that is a good follow-up.

So with that one narrowing down, I think we can go ahead here.

seberg · 2026-06-09T16:59:05Z

@@ -0,0 +1,9 @@
+New `descending` keyword argument for ``numpy.partition`` and ``numpy.argpartition``


Suggested change

New `descending` keyword argument for ``numpy.partition`` and ``numpy.argpartition``

New ``descending`` keyword argument for ``numpy.partition`` and ``numpy.argpartition``

also below string, etc. single ticks mean that it should be a link. numpy.partition can be single ticks though, because they can link to the actual functions :).

Ah sorry, I kept switching these! Thank you, this is done.

seberg · 2026-06-09T17:08:50Z

+        NPY_INTROSELECT = 0,
+        // new style names
+        NPY_SELECT_DEFAULT = 0,
+        NPY_SELECT_STABLE = 2,


Sorry, maybe this whole "make it identical" wasn't the best idea. But it's also fine. However, we should make sure that we reject unsupported values in get_partition_func which I don't think is currently the case (could do both and just delete the stable for now, but I do think we need to reject things).

Thanks, yeah, I just switched to a switch statement like we have for sorts, seems more explicit and clearer hopefully! I'm happy to delete stable if you like.

… available for decsending

…scending with unsupported dtypes

…boffsets

Co-authored-by: Sebastian Berg <sebastian@sipsolutions.net>

…eters

MaanasArora added the 01 - Enhancement label May 26, 2026

MaanasArora marked this pull request as draft May 26, 2026 10:44

seberg reviewed May 28, 2026

View reviewed changes

MaanasArora force-pushed the desc-partitions branch from 3cdc80b to 400ed71 Compare May 30, 2026 02:46

MaanasArora marked this pull request as ready for review May 30, 2026 02:48

seberg approved these changes Jun 1, 2026

View reviewed changes

seberg reviewed Jun 3, 2026

View reviewed changes

MaanasArora force-pushed the desc-partitions branch from 847f77e to 715bac8 Compare June 3, 2026 17:40

seberg reviewed Jun 9, 2026

View reviewed changes

MaanasArora added 15 commits June 9, 2026 22:42

ENH: Descending partitioning

8575c5d

TST, BUG: Descending partition tests and raise when sort fallback not…

e30f1ab

… available for decsending

TST: Add tests for partition and argpartition raising TypeError on de…

0a3273c

…scending with unsupported dtypes

DOC: Improve error message

2f013ec

ENH: Enable datetime/timedelta introselect

3f519d3

DOC, BUG: Add descending partition docs, release note, fix ma stubs

2b642c8

DOC: Update argpartition and partition docstrings

93504a0

DOC: Fix release note title line

ed9adc0

DOC, STYLE: Fix silly errors

2ed0639

BUG: Make descending keyword only in argpartition

bfc6c99

BUG: Fix pytest skip usage in test_multiarray.py

a4839ea

DOC: Remove incorrect parameter

c4ce1e2

DOC: Sort order instead of partition order

6bcfd4e

BUG: Replace pytest.importskip with pytest.importorskip in test_no_su…

05e5dae

…boffsets

DOC: Fix wording in docstrings for partition function

e23346c

MaanasArora and others added 13 commits June 9, 2026 22:43

REF: Apply suggestion from Sebastian for complex initialization

eaf0da8

Co-authored-by: Sebastian Berg <sebastian@sipsolutions.net>

DOC: Update partition and argpartition docstrings for clarity

668d2cc

TST: Test for complex partition lexicographical order

49a6442

BUG: Default kind=np._NoValue for (arg)partitions

9c92b6f

REF: Update kind method definitions to accept _NoValueType

c5a9638

REF: Refactor partition function to use kwargs for additional param…

b01ed97

…eters

DOC: Fix formatting of examples section in argpartition docstring

57aed8d

DOC: Fix garbled doc

2bcbce8

TST, BUG: Fix nan creation logic in complex lexorder tests

658ceaa

ENH: Make select kinds identical to sort kinds

e06618f

DOC: Fix release note quotes

5af9e96

ENH: Reject invalid values in _get_partition_func

206a30d

TST, BUG: Fix signature for helper in object nan sort tests

9f05d17

MaanasArora force-pushed the desc-partitions branch from 715bac8 to 9f05d17 Compare June 10, 2026 02:47

	raise pytest.mark.skip("_testbuffer is not available")
	pytest.skip("_testbuffer is not available")

	sortkind = descending > 0? NPY_SELECT_DESCENDING: NPY_SELECT_DEFAULT;
	sortkind = descending > 0 ? NPY_SELECT_DESCENDING : NPY_SELECT_DEFAULT;

	{"\|descending", &PyArray_OptionalBoolConverter, &descending}) < 0) {
	{"$descending", &PyArray_OptionalBoolConverter, &descending}) < 0) {

	Partially sorts the array in such a way that the value of the element in k-th
	Partially sorts the array in such a way that the value of the element in the k-th

		@@ -0,0 +1,9 @@
		New `descending` keyword argument for ``numpy.partition`` and ``numpy.argpartition``

	New `descending` keyword argument for ``numpy.partition`` and ``numpy.argpartition``
	New ``descending`` keyword argument for ``numpy.partition`` and ``numpy.argpartition``

Uh oh!

Conversation

MaanasArora commented May 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

AI Disclosure

Uh oh!

seberg left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MaanasArora May 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MaanasArora commented May 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

MaanasArora commented May 30, 2026

Uh oh!

seberg left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MaanasArora Jun 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MaanasArora Jun 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MaanasArora commented Jun 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

MaanasArora commented Jun 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

seberg commented Jun 2, 2026

Uh oh!

MaanasArora commented Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

seberg left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

MaanasArora commented May 26, 2026 •

edited

Loading

seberg left a comment •

edited

Loading

MaanasArora May 29, 2026 •

edited

Loading

MaanasArora commented May 29, 2026 •

edited

Loading

MaanasArora Jun 1, 2026 •

edited

Loading

MaanasArora Jun 1, 2026 •

edited

Loading

MaanasArora commented Jun 1, 2026 •

edited

Loading

MaanasArora commented Jun 2, 2026 •

edited

Loading

MaanasArora commented Jun 3, 2026 •

edited

Loading

MaanasArora Jun 10, 2026 •

edited

Loading