Passing layout nodes into Numba-JITted functions

[luk-f-a]

Hi! I took an indexed array from the tests and tried to apply a jitted function.

content = awkward1.layout.NumpyArray(numpy.array([0.0, 1.1, 2.2, 3.3, 4.4]))

ind = numpy.array([0,0,0,0,1,2,3,4,4,4,4,4,], dtype=numpy.int64)
index = awkward1.layout.Index64(ind)
array = awkward1.layout.IndexedArray64(index, content)


@njit
def foo(x):
    return x[1]

foo.py_func(array) #works
foo(array) #fails

with error

Invalid use of Function(<built-in function getitem>) with argument(s) of type(s): (awkward1.IndexedArrayType(array(int64, 1d, C), awkward1.NumpyArrayType(array(float64, 1d, A), none, {}), none, {}), Literal[int](1))

This is using version 0.2.19, and numba 0.49.1rc1

[jpivarski]

The reason is that only high-level ak.Array objects can be passed into (and out of) a Numba-compiled function. (By design: the high-level arrays are the ones that users are supposed to interact with primarily; the layout objects you're creating, such as IndexedArray, are part of the "internals". Also, behaviors such as "this array of variable-length lists is actually an array of strings" are defined at the ak.Array level and not the layout level.)

To make your example work, you only need to wrap your array as an ak.Array:

>>> array2 = awkward1.Array(array)
>>> foo(array2) #works
0.0

The limitations of what is supported inside the lowered code need to be listed prominently, as it is for Numba itself (Supported Python features and Supported NumPy features). PR #268 adds subs for where this would go, just under "Converting arrays" and "Creating arrays", visible without scrolling on a regular-sized screen.