Repository: numba/numba

Issue:
`side` kwarg for numpy.searchsorted
Currently, `searchsorted()` in nopython mode only supports the default `side="left"` behavior, but it would be useful to also be able to set `side="right"`.


Discussion:
It would be nice to have this since I might need it down the line in pydata/sparse. If you point me to the code I might be able to take a look.
I found the code, I'll maybe send a PR soon, but first, a question: I'm assuming if I do the exception inside the scalar version, it'll be very slow. I guess it would be best to define Two separate functions for the different side variations and dispatch to those?

How would I do this? A closure?

Resolving patch:
diff --git a/numba/targets/arraymath.py b/numba/targets/arraymath.py
--- a/numba/targets/arraymath.py
+++ b/numba/targets/arraymath.py
@@ -1158,59 +1158,67 @@ def bincount_impl(a, weights=None):
 
     return bincount_impl
 
+def _searchsorted(func):
+    def searchsorted_inner(a, v):
+        n = len(a)
+        if np.isnan(v):
+            # Find the first nan (i.e. the last from the end of a,
+            # since there shouldn't be many of them in practice)
+            for i in range(n, 0, -1):
+                if not np.isnan(a[i - 1]):
+                    return i
+            return 0
+        lo = 0
+        hi = n
+        while hi > lo:
+            mid = (lo + hi) >> 1
+            if func(a[mid], (v)):
+                # mid is too low => go up
+                lo = mid + 1
+            else:
+                # mid is too high, or is a NaN => go down
+                hi = mid
+        return lo
+    return searchsorted_inner
+
+_lt = register_jitable(lambda x, y: x < y)
+_le = register_jitable(lambda x, y: x <= y)
+_searchsorted_left = register_jitable(_searchsorted(_lt))
+_searchsorted_right = register_jitable(_searchsorted(_le))
 
 @overload(np.searchsorted)
-def searchsorted(a, v):
+def searchsorted(a, v, side='left'):
+    side_val = getattr(side, 'value', side)
+    if side_val == 'left':
+        loop_impl = _searchsorted_left
+    elif side_val == 'right':
+        loop_impl = _searchsorted_right
+    else:
+        raise ValueError("Invalid value given for 'side': %s" % side_val)
+
     if isinstance(v, types.Array):
         # N-d array and output
-
-        def searchsorted_impl(a, v):
+        def searchsorted_impl(a, v, side='left'):
             out = np.empty(v.shape, np.intp)
             for view, outview in np.nditer((v, out)):
-                index = np.searchsorted(a, view.item())
+                index = loop_impl(a, view.item())
                 outview.itemset(index)
             return out
 
-        return searchsorted_impl
-
     elif isinstance(v, types.Sequence):
         # 1-d sequence and output
-
-        def searchsorted_impl(a, v):
+        def searchsorted_impl(a, v, side='left'):
             out = np.empty(len(v), np.intp)
             for i in range(len(v)):
-                out[i] = np.searchsorted(a, v[i])
+                out[i] = loop_impl(a, v[i])
             return out
-
-        return searchsorted_impl
-
     else:
         # Scalar value and output
         # Note: NaNs come last in Numpy-sorted arrays
+        def searchsorted_impl(a, v, side='left'):
+            return loop_impl(a, v)
 
-        def searchsorted_impl(a, v):
-            n = len(a)
-            if np.isnan(v):
-                # Find the first nan (i.e. the last from the end of a,
-                # since there shouldn't be many of them in practice)
-                for i in range(n, 0, -1):
-                    if not np.isnan(a[i - 1]):
-                        return i
-                return 0
-            lo = 0
-            hi = n
-            while hi > lo:
-                mid = (lo + hi) >> 1
-                if a[mid] < v:
-                    # mid is too low => go up
-                    lo = mid + 1
-                else:
-                    # mid is too high, or is a NaN => go down
-                    hi = mid
-            return lo
-
-        return searchsorted_impl
-
+    return searchsorted_impl
 
 @overload(np.digitize)
 def np_digitize(x, bins, right=False):


Test patch:
diff --git a/numba/tests/test_np_functions.py b/numba/tests/test_np_functions.py
--- a/numba/tests/test_np_functions.py
+++ b/numba/tests/test_np_functions.py
@@ -42,6 +42,12 @@ def bincount2(a, w):
 def searchsorted(a, v):
     return np.searchsorted(a, v)
 
+def searchsorted_left(a, v):
+    return np.searchsorted(a, v, side='left')
+
+def searchsorted_right(a, v):
+    return np.searchsorted(a, v, side='right')
+
 def digitize(*args):
     return np.digitize(*args)
 
@@ -312,11 +318,25 @@ def test_searchsorted(self):
         pyfunc = searchsorted
         cfunc = jit(nopython=True)(pyfunc)
 
+        pyfunc_left = searchsorted_left
+        cfunc_left = jit(nopython=True)(pyfunc_left)
+
+        pyfunc_right = searchsorted_right
+        cfunc_right = jit(nopython=True)(pyfunc_right)
+
         def check(a, v):
             expected = pyfunc(a, v)
             got = cfunc(a, v)
             self.assertPreciseEqual(expected, got)
 
+            expected = pyfunc_left(a, v)
+            got = cfunc_left(a, v)
+            self.assertPreciseEqual(expected, got)
+
+            expected = pyfunc_right(a, v)
+            got = cfunc_right(a, v)
+            self.assertPreciseEqual(expected, got)
+
         # First with integer values (no NaNs)
         bins = np.arange(5) ** 2
         values = np.arange(20) - 1
@@ -345,6 +365,20 @@ def check(a, v):
             # Sequence values
             check(a, list(values))
 
+        # nonsense value for 'side' raises TypingError
+        def bad_side(a, v):
+            return np.searchsorted(a, v, side='nonsense')
+        cfunc = jit(nopython=True)(bad_side)
+        with self.assertTypingError():
+            cfunc([1,2], 1)
+
+        # non-constant value for 'side' raises TypingError
+        def nonconst_side(a, v, side='left'):
+            return np.searchsorted(a, v, side=side)
+        cfunc = jit(nopython=True)(nonconst_side)
+        with self.assertTypingError():
+            cfunc([1,2], 1, side='right')
+
     def test_digitize(self):
         pyfunc = digitize
         cfunc = jit(nopython=True)(pyfunc)