Numba signature for structured arrays

Question

Numba's documentation does not give any example of signatures for functions that take structured arrays. I have tried several ways, but all were rejected by Numba (and Pylance).

import numba as nb
import numpy as np

PairSpec = [("x", np.float32), ("y", np.float32)]
Pair = np.dtype(PairSpec)
NumbaPair = nb.from_dtype(Pair)

# BUG None of this works
# @nb.jit(np.float32(Pair[:]))
# @nb.jit(np.float32(NumbaPair[:]))
@nb.jit
def sum(pairs):
    pair = pairs[0]
    return pair.x + pair.y


pairs = np.array([(2, 3)], dtype=PairSpec)
print(sum(pairs))

How to give a signature to a function that takes structured arrays?

score 3 · Accepted Answer · answered Nov 15 '22 at 01:33

The correct signature is nb.float32(NumbaPair[:]). Note the use of nb.float32 and not np.float32. Also please note that arrays of structures (AoS) generally tend to be less efficient than structures of arrays (SoA). This is especially true for coordinates since most fields are generally read and AoS prevent any efficient vectorization (while modern x86-64 processors can typically compute ~16 float32 values per cycle and per core, as opposed to 2 for scalar values).

Numba signature for structured arrays

1 Answers1