How to check for field in a structured array?

Numba Support: How do I do ...?
1

I have a structured array and some of its fields may or may not exist. I would like to check whether it exists and do something based on it, like the following:

import numpy as np
from numba import njit
from numba.extending import overload


def get_keys(recarr):
    return tuple(recarr.dtype.fields.keys())


@overload(get_keys)
def ol_get_keys(recarr):
    keys = get_keys(recarr)
    return lambda recarr: keys


INFO = np.void(
    (np.nan,),
    dtype=np.dtype(
        [('a', np.float64)],
        align=True,
    ),
)
info = np.full(1, INFO)


@njit
def test(info):
    n = 0
    if 'b' in get_keys(info):
        n += info[0]['b']
    return n


test(info)

But this is not working:

KeyError: "Field 'b' was not found in record with fields ('a',)"

Anyone has ideas? Thanks in advance!

2

While compiling

n += info[0]['b']

numba will encounter an error, since for the given typed signature info does not have the field b. Same can be seen with the simpler conditional

if 'b' in ['a']:
    n += info[0]['b']

which will raise the same error, indicating that the issue is not only in checking whether the given field exists in the structure.

Perhaps try overloading the entire calculation n += info[0]['b'] itself, dispatching different functions depending on the type of the structure? Something like

@overload(add_b)
def ol_add_b(recarr, n):
    if 'b' in recarr.dtype.fields.keys():
        return lambda recarr, n: recarr[0]['b'] + n
    return lambda recarr, n: n
3

Yes, it works! Thanks a lot!