ENH: support N-D arrays in np.searchsorted by searching along the last axis - #32346

Do we need the GIL for all types? Just ones with references, right, and these are built-in - so object?

That was indeed a regression. There is now a get_loop that builds the carrier once per call into NpyAuxData and clears NPY_METH_REQUIRES_PYAPI per descriptor. Only a reference check is not enough though: np.dtype([('x','i4'),('y','f8')]) has hasobject=False but NEEDS_PYAPI=True, because VOID_compare reaches into Python even without object fields. So it uses PyDataType_FLAGCHK(descr, NPY_NEEDS_PYAPI) now, the same test NPY_BEGIN_THREADS_DESCR makes, and it has to be per descriptor rather than per dtype since one void dtype covers both cases.

I had no idea how to solve this btw so the idea and implementation are all due to my AI agent. Check 0a59156 and let me know what you think. I personally felt like this is a bit hacky but at the same time could not think of something better.

Yeah, that's annoying! I'd say keep the dtypes separate: object obviously always needs the GIL, so we can scope this just to void. Also, don't string and unicode never need the GIL?

(Honestly, I'd be fine if we can even skip some (or all!) generic types here, and address them in follow-up. All the generic types except void have tag-level comparisons now, we're just not using them in binsearch.cpp.)

Should have done what you wanted in de9e24e

Well, this is very annoying. I am OK with the carrier hack, although if that is the only reason for auxdata, then there are a few places where I have been doing crazy things and just allocate a static array that is completely undefined except the dtype field.

The point is that there are exactly two reasons anyone could in theory use the array reasonably or how we use it:

• To fetch the dtype
• And in crazy rare cases for void/structured dtypes to know who owns the data.

I would be shocked if we need more than the first here. I am not proud of that hack and it's ugly, but so is allocating a full fresh array for nothing...

Alternatively, if that smells too much for you, we could skip it for builtin dtypes that aren't string/bytes/void.

My AI agent noticed and I confirmed that an override that raises another TypeError will also get swallowed here, as we're checking if array_ufunc is used at all:

self = <test_multiarray.TestMethods object at 0x7cd03e32c2d0>

    def test_searchsorted_array_ufunc_override(self):
        # an operand overriding ufuncs is consulted, and one that declines
        # falls back to the ordinary implementation
        class Handles(np.ndarray):
            def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
                return "handled"
    
        class Declines(np.ndarray):
            def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
                return NotImplemented
    
        class Raises(np.ndarray):
            def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
                raise TypeError("boom")
    
        a = np.arange(10.)
        v = np.array([2.5, 7.5])
        assert a.searchsorted(v.view(Handles)) == "handled"
        assert a.view(Handles).searchsorted(v) == "handled"
        assert_equal(a.searchsorted(v.view(Declines)), [3, 8])
        assert_equal(a.view(Declines).searchsorted(v), [3, 8])
>       with pytest.raises(TypeError, match="boom"):
E       Failed: DID NOT RAISE <class 'TypeError'>

We should fix and add a test for this too!

Indeed checking if any operand have has __array_ufunc__, which cannot tell apart a decline from a raise. It now tests whether the pending exception is the one the machinery raises when every override returned NotImplemented, so an override's own error propagates.

That also surfaced a second thing, __array_ufunc__ = None is an opt-out rather than an implementation, and searchsorted has always just converted such an operand. It started raising, so that is restored with a test too.

However, It's a bit hacky that the decline is identified by the message prefix from array_ufunc_errmsg_formatter, since there is no distinct exception type and PyUFunc_CheckOverride returns -1 for both cases. If you have a better idea of how to do this, please let me know.

See f1d2936

Right, I think this error-checking thing is difficult to do generally. Checking the message is definitely not the way to go though.

Could it be possible to check if __array_ufunc__ exists before you make the call to the gufunc?

Also: PyArray_SearchSorted is only for legacy dtypes? And we already have gufuncs for all of those. The only way a user-dtype can implement these gufuncs is by filling them in. So why do we need this fallback at all?

Edit: sorry, said it backwards. I meant, at least in the absence of a gufunc-registration mechanism, we can guard searchsorted_gufunc to only work with legacy dtypes. Still quite hacky, though.

Hmm, could searchsorted be overwritten by array ufunc before? I don't think so, right? In that case let's just not add that behavior! This is entirely private, and we'll provide ways to register anyway, after all.

Have no time to look right in detail right now, but ideally we do allow __array_ufunc__ to override all ufuncs! Will try to come back to this soon, but it is first week of classes and I'm teaching a fair lot this term, plus I have a visitor, so may have to be next week.

__array_ufunc__ never worked for searchsorted. It's not an exposed ufunc to the user. The internals are here being turned into 4 gufuncs so __array_ufunc__ can work on those if the user overrides them. But at the top level, searchsorted is not a (g)ufunc. I'd argue that it's not needed to support __array_ufunc__ here at all. I don't think a single person is going to use it. Nobody is going to override four ufuncs for this and since it never worked it's fine to not enable such a thing here.

__array_function__ keeps working fine as it always has for searchsorted.

The internals are here being turned into 4 gufuncs so __array_ufunc__ can work on those if the user overrides them.

That is all I meant; if that works, I'm OK, we definitely do not have to do more (it is a bit like the string functions, where one can override the uderlying ufuncs too but then has to deal oneself with the prerequisites).

EDIT: apologies for commenting while not actually having time to look at the code.

Yeah the problem is that because the top level entry point is in C and the gufuncs are being called from C, checking for the override and handling it here becomes complicated which is why I'm saying it's fine to not enable such a thing here.

Yeah, that's annoying! I'd say keep the dtypes separate: object obviously always needs the GIL, so we can scope this just to void. Also, don't string and unicode never need the GIL?

(Honestly, I'd be fine if we can even skip some (or all!) generic types here, and address them in follow-up. All the generic types except void have tag-level comparisons now, we're just not using them in binsearch.cpp.)

Right, I think this error-checking thing is difficult to do generally. Checking the message is definitely not the way to go though.

Could it be possible to check if __array_ufunc__ exists before you make the call to the gufunc?

Also: PyArray_SearchSorted is only for legacy dtypes? And we already have gufuncs for all of those. The only way a user-dtype can implement these gufuncs is by filling them in. So why do we need this fallback at all?

Edit: sorry, said it backwards. I meant, at least in the absence of a gufunc-registration mechanism, we can guard searchsorted_gufunc to only work with legacy dtypes. Still quite hacky, though.

Sorry, missed this: can user-defined types be legacy? If not, is this reachable? Either way, I don't think this should be done in the promoter, if at all!

It seems this function is actually "legacy" and roundabout. We don't need it anyway:

(or Py_NewRef if you prefer)

In hindsight, shouldn't this be part of the promoter rather than resolve_descriptors?

(Sorry if this feels like it adds weird comments, or if you already discussed it. Just let me know, the caveat means I am not 100% sure that doing this falls into place well enough, although I somewhat suspect so.)

There is an important point to this, but this looks right to me. We have two problems:

1. np.searchsorted(int64_arr, float64_arr) currently works by using the common dtype of both. This is as of now still handled in the wrapper, though. The promoter assumes this is the case.
2. If you have fixed-size string, bytes, or structured dtypes, then those actually do define a compare but need both sides to have the identical dtype instance.
   (The caller really already ensures this right now, but it's still right.)

So this is correct and necessary to resolve point 2 (i.e. while calling "PromoteTypes", this is the correct type-resolution step).

However, it loops back a bit to the other part and I think I now see what you meant yesterday with things not being sure about the design around using the compare functions.

This is a thought, one caveat that may make it only semi useful noted in-inline.

What I might prefer is if we can move things around a little bit so that:

1. Change the promoter to effectively have two paths (yes, we'll call it twice typically!):
   • The current: If all DTypes are as expected, check if we can add a loop for the DTypes based on the compare function. (Promoter returns 1 if it successfully added the loop.) ¹
   • If DTypes are not OK for the above, it calls CommonDType on the first two and sets the second to Intp.
     There is a real problem with this! It'll break np.searchsorted(np.array([1, 2, 3, 4], dtype="uint8"), 500). There are ways around this (see integer comparison ufuncs) but they are awkward... We could consider just converting incoming Python integers -- still changing promotion for floats subtly, but to the better...
   • It makes sense to have a "third" step and start with only modifying the last Intp operand if it isn't already. That will help with broadening things up to mixed DType loops.)
   • (Searchsorted is one of those functions where it feels usually wrong to promote the first array to the second, but ...)
2. You are currently relying on the setup call to already return two identical arrays. I think it would be nice if we can not do that.

I think this (with subok=False) should allow us to remove the early array conversion (with caveat!) in the searchsorted wrapper function which additionally means we push promotion into the gufuncs.
Right now that would make things like searchsorted(uint8_arr, pyint) more awkward, but would also create a path for us to fix searchsorted(uint64_arr, int64_arr) to not incorrectly promote to float64.

Another small question: Is there a BC reason to not just change PyArray_SearchSorted() more drastically to call the gufunc as well? (Fine if there is, I'd prefer to deprecate many of these rather than trying hard to keep them aligned.)

Is this to ensure we accommodate both types of ufuncs? Would *dtypes[4] not work? IIRC we avoid use of NPY_MAXARGS (it is so many anyway!)

nit: maybe add a simple comment that this is for legacy compatibility?

Should maybe add a release note on the fact that we now may call the search functions with non-contiguous loops, I suspect.

(oops, sorry, wanted to just store it as a review for myself, because I think it drifted very subtly.)

Effectively always!

Well, this is very annoying. I am OK with the carrier hack, although if that is the only reason for auxdata, then there are a few places where I have been doing crazy things and just allocate a static array that is completely undefined except the dtype field.

The point is that there are exactly two reasons anyone could in theory use the array reasonably or how we use it:

• To fetch the dtype
• And in crazy rare cases for void/structured dtypes to know who owns the data.

I would be shocked if we need more than the first here. I am not proud of that hack and it's ugly, but so is allocating a full fresh array for nothing...

Alternatively, if that smells too much for you, we could skip it for builtin dtypes that aren't string/bytes/void.