numba @njit to update a big dict

Question

I try to use numba for a function that need to do a search on a very big (10e6) dict with (int, int) tuple as key.

import numpy as np
from numba import njit

myarray = np.array([[0, 0],  # 0, 1
                    [0, 1],
                    [1, 1],  # 1, 2
                    [1, 2],  # 1, 3
                    [2, 2],
                    [1, 3]]
) # a lot of this with shape~(10e6, 2)

dict_with_tuples_key = {(0, 1): 1,
                        (3, 7): 1} # ~10e6 keys 

A simplified version look like this

# @njit
def update_dict(dict_with_tuples_key, myarray):
    for line in myarray:
        i, j = line
        if (i, j) in dict_with_tuples_key:
            dict_with_tuples_key[(i, j)] += 1
        else:
            dict_with_tuples_key[(i, j)] = 1
    return dict_with_tuples_key

new_dict = update_dict(dict_with_tuples_key, myarray)
print new_dict

new_dict = update_dict2(dict_with_tuples_key, myarray)
# print new_dict
# {(0, 1): 2,   # +1 already in dict_with_tuples_key
#  (0, 0): 1,   # diag
#  (1, 1): 1,   # diag
#  (2, 2): 1,   # diag
#  (1, 2): 1,   # new from myarray
#  (1, 3): 1,   # new from myarray
#  (3, 7): 1 }

It would appear that @njit does not accept dict as function arg ?

I'm wondering how to rewrite this, specially the if (i, j) in dict_with_tuples_key part that do the search.

Answer 1

njit means that the function is compiled in nopython mode. A dict, list and tuple are python objects and therefore not supported. Not as arguments and not inside the function.

If your dict keys are all different I would consider using a 2D numpy array where the first axis represents the first index of the dict-key-tuple and the second axis the second index. Then you could rewrite it as:

from numba import njit
import numpy as np

@njit
def update_array(array, myarray):
    elements = myarray.shape[0]
    for i in range(elements):
        array[myarray[i][0]][myarray[i][1]] += 1 
    return array


myarray = np.array([[0, 0], [0, 1], [1, 1],
                    [1, 2], [2, 2], [1, 3]])

# Calculate the size of the numpy array that replaces the dict:
lens = np.max(myarray, axis=0) # Maximum values
array = np.zeros((lens[0]+1, lens[1]+1)) # Create an empty array to hold all indexes in myarray
update_array(array, myarray)

Since you already indexed your dictionary with tuples the transition problems to indexing an array will not be great.