Using Awkward Arrays in Numba's jitclass

[kevinxu-git]

Version of Awkward Array

2.5.0

Description and code to reproduce

Context

In order to pass a ragged array in a class wrapped using @jitclass of numba. I want to use ak.Array to store some information that will be used inside a method of the class.

Example of usage

import awkward as ak
import numpy as np
from numba.experimental.jitclass import jitclass
from unittest import TestCase, main

@jitclass([("arr", ak.Array([[[0]]]).numba_type)])
class A:
    def __init__(self, arr):
        self.arr = arr

    def do_smth(self):
        # some access to information stored in self.arr
        self.arr[0][0]
        print(len(self.arr))

def create_a():
    arr = ak.Array([
         [[0]], 
         [[0]],
         [[0]]
    ])
    return A(arr=arr)

class TestA(TestCase):
    def test_a(self):
        a = create_a()
        
        a.do_smth()
        x = np.zeros(4)
        a.do_smth()

main(argv=[''], verbosity=2, exit=False)

Issue

This code raises :
ValueError: slice index out of bounds

It seems like only the first call to do_smth can access to self.arr, but the second call can't find anymore the values in self.arr

What's more surprising is that

• When I init arr with more or less elements, it works ... It's only when there are exactly 3 elements that breaks
• If I remove x = np.zeros(4), it works as well
• If I use x = np.zeros(3), it works as well ...

So, if anymore could help me to explain this bug and how to solve it 🥲

[agoose77]

@ianna is this something that you'd be the person to ping for? :)

[jpivarski]

This is the sort of thing that @ianna can look into, but we haven't thought about what would happen if someone tries to use Awkward Arrays in jitclass; we should probably figure out if they can be expected to work and confirm that they're supported or else declare them unsupported, preferably with an error message.

The symptoms sound a lot like looking at uninitialized memory—to have something work with anything but 3 elements suspiciously sounds like the "working" cases are just lucky that nothing has overwritten that memory. If so, it's very unlikely that the "working" cases will continue to work on another computer or another operating system.

Awkward Arrays in Numba are borrowed references—the Lookup object on the Python side holds a reference to keep the array in memory while the Numba-compiled function runs. If it goes out of scope on the Python side and gets deleted by the Python garbage collector (or just when the reference count goes to zero), then the Numba-compiled function will be looking at memory that has been released, which might or might not still have the array's data in it.

Understanding what's happening here would involve understanding how Numba's jitclass holds things in memory. (Also, we should verify with the Numba team whether jitclass is an interface that will continue to be supported. There's another way to make class-like objects that seems more robust and I'm wondering if it's intended as a replacement. I can't remember its name and I can't find it right now. I think it started with "Struct...".)

[kevinxu-git]

The symptoms sound a lot like looking at uninitialized memory—to have something work with anything but 3 elements suspiciously sounds like the "working" cases are just lucky that nothing has overwritten that memory. If so, it's very unlikely that the "working" cases will continue to work on another computer or another operating system.

Awkward Arrays in Numba are borrowed references—the Lookup object on the Python side holds a reference to keep the array in memory while the Numba-compiled function runs. If it goes out of scope on the Python side and gets deleted by the Python garbage collector (or just when the reference count goes to zero), then the Numba-compiled function will be looking at memory that has been released, which might or might not still have the array's data in it.

I see, it would explain it !

Understanding what's happening here would involve understanding how Numba's jitclass holds things in memory. (Also, we should verify with the Numba team whether jitclass is an interface that will continue to be supported. There's another way to make class-like objects that seems more robust and I'm wondering if it's intended as a replacement. I can't remember its name and I can't find it right now. I think it started with "Struct...".)

In the current version of the awkward library, it's not supposed to be used with jitclass then 😅

[jpivarski]

We need to have more comprehensive documentation on using Awkward Array in Numba, and that documentation would specify which Numba features are applicable and which are not. For now, though, I'll make this a Discussion, since it can act as a public service announcement to not use Awkward Arrays in nb.jitclass.