Numba - string types

Question

I think this is a simple question, but I'm finding the numba documentation lacking on how to use string types with numpy arrays and dictionaries. I have a function I want to use numba which takes a list of zip codes, and then a dictionary which maps zip codes -> value. I know I could cast zips to an int datatype, but I want to be able to do this with other values that cannot coerce to int types.

from numba.typed import Dict
from numba.core import types
@jit(nopython=True)
def zip_func(zip_array,zip_dict):
    return zip_dict[zip_array[0]]
#This would be a longer array using real data
myzips = np.array(['12345'],dtype='<U5')
mydict = {"12345":5}
numba_dict = Dict.empty(key_type=types.unicode_type,value_type=types.int64)
for key,value in mydict.items():
    numba_dict[key] = value
zip_func(myzips,numba_dict)

This gives the following error, which if I understand right it's basically not able to use a fixed-length string of length 5 as a key when the key is typed as types.unicode_type. How do I align these types?

---------------------------------------------------------------------------
TypingError                               Traceback (most recent call last)
<ipython-input-4-93b4a32d59d4> in <module>
     12 for key,value in mydict.items():
     13     numba_dict[key] = value
---> 14 zip_func(myzips,numba_dict)

/sas/python/app/miniconda3/envs/py3lu/lib/python3.6/site-packages/numba/core/dispatcher.py in _compile_for_args(self, *args, **kws)
    399                 e.patch_message(msg)
    400 
--> 401             error_rewrite(e, 'typing')
    402         except errors.UnsupportedError as e:
    403             # Something unsupported is present in the user code, add help info

/sas/python/app/miniconda3/envs/py3lu/lib/python3.6/site-packages/numba/core/dispatcher.py in error_rewrite(e, issue_type)
    342                 raise e
    343             else:
--> 344                 reraise(type(e), e, None)
    345 
    346         argtypes = []

/sas/python/app/miniconda3/envs/py3lu/lib/python3.6/site-packages/numba/core/utils.py in reraise(tp, value, tb)
     78         value = tp()
     79     if value.__traceback__ is not tb:
---> 80         raise value.with_traceback(tb)
     81     raise value
     82 

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Internal error at <numba.core.typeinfer.IntrinsicCallConstraint object at 0x7ff4562e2518>.
Failed in nopython mode pipeline (step: nopython mode backend)
Cannot cast [unichr x 5] to unicode_type: %".21" = load [5 x i32], [5 x i32]* %"key"

File "../../../python/app/miniconda3/envs/py3lu/lib/python3.6/site-packages/numba/typed/dictobject.py", line 730:
    def impl(d, key):
        castedkey = _cast(key, keyty)
        ^

[1] During: lowering "$0.4 = call $0.1(key, $0.3, func=$0.1, args=[Var(key, dictobject.py:730), Var($0.3, dictobject.py:730)], kws=(), vararg=None)" at /sas/python/app/miniconda3/envs/py3lu/lib/python3.6/site-packages/numba/typed/dictobject.py (730)
[2] During: typing of intrinsic-call at <ipython-input-4-93b4a32d59d4> (7)
Enable logging at debug level for details.

File "<ipython-input-4-93b4a32d59d4>", line 7:
def zip_func(zip_array,zip_dict):
    return zip_dict[zip_array[0]]

Answer 1

Instead of using numba.types.unicode_type as the key_type for your dictionary, you could create a custom Numba type from the unichar dtype.

np_unichar = numpy.dtype('<U5')
UnicharType = numba.from_dtype(np_unichar)

Now, you can use UnicharType to tell Numba you're expecting values of type [unichr x 5].

Alternatively, I found that simply adding an empty string to a [unichr x N] will force a type conversion. Which means, zip_array[0] + '' would be processed as numba.types.unicode_type.