Repository: scikit-hep/awkward

Issue:
Numba's `njit` doesn't support non-default behaviors
### Version of Awkward Array

main

### Description and code to reproduce

```python
import numba
import awkward as ak

SOME_ATTRS = {"FOO": "BAR"}
builder = ak.ArrayBuilder(behavior=SOME_ATTRS)


@numba.njit
def func(array):
    return array


assert func(builder).behavior is SOME_ATTRS

```


Discussion:
Fairly recently, the internal Numba type string for an array with behaviors was changed from including a string of the entire behavior dict (with `repr`) to just a hash (much faster).

On the way out, is it reconstituting the array by trying to turn that string back into a Python object? Because the hash is not reversible. Could that be related to the issue?

(I think round-tripping behaviors through Numba used to work, and I'm surprised that there aren't tests for it.)
I haven't looked at the code in detail (I opened the issue to keep an eye on it!), but I think the issue is that dicts aren't hashable, so the hashing fails.

We should definitely add a test for this as part of the fix.
> Fairly recently, the internal Numba type string for an array with behaviors was changed from including a string of the entire behavior dict (with `repr`) to just a hash (much faster).
> 
> On the way out, is it reconstituting the array by trying to turn that string back into a Python object? Because the hash is not reversible. Could that be related to the issue?
> 
> (I think round-tripping behaviors through Numba used to work, and I'm surprised that there aren't tests for it.)

Yes, that's what I observe. The Python 'dict' is unhashable. The TypeError occures while unboxing the ArrayBuilder. Specifically in:
```
  File "/Users/yana/Projects/PR2763/awkward/src/awkward/_connect/numba/builder.py", line 78, in box_ArrayBuilder
    c.pyapi.serialize_object(arraybuildertype.behavior)
  File "/Users/yana/opt/anaconda3/envs/numba_py311/lib/python3.11/site-packages/numba/core/pythonapi.py", line 1410, in serialize_object
    gv = self.module.__serialized[obj]
         ~~~~~~~~~~~~~~~~~~~~~~~~^^^^^
TypeError: unhashable type: 'dict'

```
I think, we need to consider using the Numba Dict type here.  
Numba dicts are for lowering; the problem here is hashing.

```python
>>> d = nb.typed.Dict()
>>> d["one"] = 1
>>> d["two"] = 2
>>> d["three"] = 3
>>> d
DictType[unicode_type,int64]<iv=None>({one: 1, two: 2, three: 3})
>>> hash(d)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: unhashable type: 'Dict'
```

Assuming that all of the values are hashable, we could convert them to and from tuples (since the key order in modern dicts—that is, after Python 3.5—is stable).

```python
>>> d = {"one": 1, "two": 2, "three": 3}
>>> tuple_d = tuple(d.items())
>>> hash(tuple_d)
-1150126290131628349
>>> dict(tuple_d)
{'one': 1, 'two': 2, 'three': 3}
```

Resolving patch:
diff --git a/src/awkward/_connect/numba/builder.py b/src/awkward/_connect/numba/builder.py
--- a/src/awkward/_connect/numba/builder.py
+++ b/src/awkward/_connect/numba/builder.py
@@ -74,8 +74,16 @@ def box_ArrayBuilder(arraybuildertype, arraybuilderval, c):
     ArrayBuilder_obj = c.pyapi.unserialize(
         c.pyapi.serialize_object(ak.highlevel.ArrayBuilder)
     )
-    behavior_obj = c.pyapi.unserialize(
-        c.pyapi.serialize_object(arraybuildertype.behavior)
+    serializable2dict_obj = c.pyapi.unserialize(
+        c.pyapi.serialize_object(ak._connect.numba.arrayview.serializable2dict)
+    )
+    behavior2_obj = c.pyapi.unserialize(
+        c.pyapi.serialize_object(
+            ak._connect.numba.arrayview.dict2serializable(arraybuildertype.behavior)
+        )
+    )
+    behavior_obj = c.pyapi.call_function_objargs(
+        serializable2dict_obj, (behavior2_obj,)
     )
 
     proxyin = c.context.make_helper(c.builder, arraybuildertype, arraybuilderval)
diff --git a/src/awkward/_connect/numba/layoutbuilder.py b/src/awkward/_connect/numba/layoutbuilder.py
--- a/src/awkward/_connect/numba/layoutbuilder.py
+++ b/src/awkward/_connect/numba/layoutbuilder.py
@@ -302,6 +302,16 @@ def getter(builder):
     return getter
 
 
+@numba.extending.overload_method(EmptyType, "append")
+def Empty_append(builder, datum):
+    if isinstance(builder, EmptyType):
+
+        def append(builder, datum):
+            raise NumbaTypeError("Empty cannot append data")
+
+        return append
+
+
 ########## ListOffset #########################################################
 
 


Test patch:
diff --git a/tests/test_2408_layoutbuilder_in_numba.py b/tests/test_2408_layoutbuilder_in_numba.py
--- a/tests/test_2408_layoutbuilder_in_numba.py
+++ b/tests/test_2408_layoutbuilder_in_numba.py
@@ -7,6 +7,8 @@
 
 numba = pytest.importorskip("numba")
 
+from numba.core.errors import NumbaTypeError  # noqa: E402
+
 import awkward.numba.layoutbuilder as lb  # noqa: E402
 
 ak.numba.register_and_check()
@@ -593,7 +595,7 @@ def f2(x):
 
     builder = lb.Empty()
     # Unknown attribute 'append' of type lb.Empty
-    with pytest.raises(numba.core.errors.TypingError):
+    with pytest.raises(NumbaTypeError):
         f2(builder)
 
 
diff --git a/tests/test_2770_serialize_and_deserialize_behaviour_for_numba.py b/tests/test_2770_serialize_and_deserialize_behaviour_for_numba.py
new file mode 100644
--- /dev/null
+++ b/tests/test_2770_serialize_and_deserialize_behaviour_for_numba.py
@@ -0,0 +1,43 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+import numpy as np
+import pytest
+
+import awkward as ak
+
+numba = pytest.importorskip("numba")
+
+
+@numba.njit
+def func(array):
+    return array
+
+
+def test_ArrayBuilder_behavior():
+    SOME_ATTRS = {"FOO": "BAR"}
+    builder = ak.ArrayBuilder(behavior=SOME_ATTRS)
+
+    assert builder.behavior is SOME_ATTRS
+    assert func(builder).behavior == SOME_ATTRS
+
+
+def test_ArrayBuilder_non_picklable_behavior():
+    def make_add_xyr():
+        def add_xyr(left, right):
+            x = left.x + right.x
+            y = left.y + right.y
+            return ak.zip(
+                {
+                    "x": x,
+                    "y": y,
+                    "r": np.sqrt(x**2 + y**2),
+                },
+                with_name="xyr",
+            )
+
+        return add_xyr
+
+    behavior = {(np.add, "xyr", "xyr"): make_add_xyr()}
+    builder = ak.ArrayBuilder(behavior=behavior)
+    result = func(builder)
+    assert result.behavior.keys() == behavior.keys()