How to parallelize for-loop over an array with Numba

Question

I'm trying to write a Numba function that loops over a list (not a range) with @numba.njit(parallel=True). For example,

import numpy as np
import numba

arr = np.ones(10)
idx = np.array([4, 2, 5])

@numba.njit(parallel=True)
def foo(arr, idx):
    for i in idx:
        arr[i] = 0
    return arr

foo(arr, idx)

The warning I get is NumbaPerformanceWarning: The keyword argument 'parallel=True' was specified but no transformation for parallel execution was possible. I am aware that one can explicitly parallelize similar loops using numba.prange(), but I need the loop to loop over an array of indices. Is this possible and how?

Answer 1

for i in idx: ... is semantically an equivalent of:

for j in range(len(idx)):
    i = idx[j]
    [...]

The performance of the second version may be slightly different. Based on the second version, you can replace range by prange and get a parallel version.

However, note that this will cause a race condition if two value of idx are the same. Race conditions results in undefined behaviour (basically, anything can happen).