fix: handle unhashable behaviour type (#2770)

* fix: handle unhashable behaviour type * fix: add checks for None * fix: use arrayview functions - Jim's suggestion * fix: cleanup * test: add a test for a non-pickleable behaviour * test: rename the test * fix: clenup * test: compare behaviours * test: add more checks * test: just check keys --------- Co-authored-by: Angus Hollands <[email protected]>
scikit-hep · Oct 26, 2023 · 65dc991 · 65dc991
1 parent e8072ce
commit 65dc991
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 3 deletions.
diff --git 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
@@ -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 #########################################################
 
 

diff --git 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
@@ -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()