Parallelize np.searchsorted

user_18335909 0
Viewed 134

Is there a way to parallelize the implementation of np.searchsorted()?

I have a situation where the base array a and value array v are of the same order of size. From what I understand of the search sorted algorithm, it does the operation for each element in v in turn. I'd like to parallelize it so it performs this sorting on multiple elements of v at the same time. (If I have 32 cores, it should be able to sort 32 elements at once, right?) Is there a way to implement this?

I tried to use Numba @jit(nopython=True, nogil=True, parallel=True) but it shows no improvement in speed and no increase in CPU usage.

For reference, a and v are lists of integers with length on the order of 10^7 elements.

numba numpy python
1 Answers

A trivially parallelized np.searchsorted works for me. Between 3.4x and 3.9x speed up on a 2-core x 2 threads colab instance with a, b length 10**7 (e.g. 9.94 s/2.62 s) using numba 0.55.1, omp threading layer.

import numba as nb

@nb.njit(parallel=True)
def searchsorted_parallel(a, b):
    res = np.empty(len(b), np.intp)
    for i in nb.prange(len(b)):
        res[i] = np.searchsorted(a, b[i])
    return res

Running a micro-benchmark

import numpy as np

a = np.random.randint(200000, size=10**7)
b = np.random.randint(200000, size=10**7)

a.sort()

r = [0,0]
%timeit -r1 -n1 r[0] = np.searchsorted(a,b)
#1 loop, best of 1: 9.31 s per loop
%timeit -r1 -n1 r[1] = searchsorted_parallel(a,b)
#1 loop, best of 1: 2.36 s per loop

np.testing.assert_array_equal(r[0], r[1])
user_142777220
Related