Repository: scikit-hep/awkward

Issue:
Pickle of array includes `array.__class__`
### Version of Awkward Array

main

### Description and code to reproduce

This might not be ideal; we'd prefer e.g. that behaviours defined in Awkward aren't included in the pickle explicitly.


Discussion:
I think we want to use `__reduce__` for this, rather than `__getstate__`.
@jpivarski can we be sure to hold the Python release until this issue has been fixed? Can do it today, later.

Resolving patch:
diff --git a/src/awkward/highlevel.py b/src/awkward/highlevel.py
--- a/src/awkward/highlevel.py
+++ b/src/awkward/highlevel.py
@@ -1430,7 +1430,7 @@ def numba_type(self):
 
         return numba.typeof(self._numbaview)
 
-    def __getstate__(self):
+    def __reduce__(self):
         packed = ak.operations.to_packed(self._layout, highlevel=False)
         form, length, container = ak.operations.to_buffers(
             packed,
@@ -1442,10 +1442,10 @@ def __getstate__(self):
             behavior = None
         else:
             behavior = self._behavior
-        return form.to_dict(), length, container, behavior
+        return object.__new__, (Array,), (form.to_dict(), length, container, behavior)
 
     def __setstate__(self, state):
-        form, length, container, behavior = state
+        form, length, container, behavior, *_ = state
         layout = ak.operations.from_buffers(
             form,
             length,
@@ -2068,7 +2068,7 @@ def numba_type(self):
 
         return numba.typeof(self._numbaview)
 
-    def __getstate__(self):
+    def __reduce__(self):
         packed = ak.operations.to_packed(self._layout, highlevel=False)
         form, length, container = ak.operations.to_buffers(
             packed.array,
@@ -2080,10 +2080,14 @@ def __getstate__(self):
             behavior = None
         else:
             behavior = self._behavior
-        return form.to_dict(), length, container, behavior, packed.at
+        return (
+            object.__new__,
+            (Record,),
+            (form.to_dict(), length, container, behavior, packed.at),
+        )
 
     def __setstate__(self, state):
-        form, length, container, behavior, at = state
+        form, length, container, behavior, at, *_ = state
         layout = ak.operations.from_buffers(
             form,
             length,


Test patch:
diff --git a/tests/test_2106_pickle_class.py b/tests/test_2106_pickle_class.py
new file mode 100644
--- /dev/null
+++ b/tests/test_2106_pickle_class.py
@@ -0,0 +1,100 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+import contextlib
+import pickle
+import sys
+import types
+import uuid
+
+import pytest  # noqa: F401
+
+import awkward as ak
+
+
+@contextlib.contextmanager
+def temporary_module():
+    name = str(uuid.uuid1()).replace("-", "_")
+    module = types.ModuleType(name)
+    sys.modules[name] = module
+    yield module
+    del sys.modules[name]
+
+
+def impl():
+    with temporary_module() as module:
+        exec(
+            """
+import awkward as ak
+
+behavior = {}
+@ak.mixin_class(behavior, "my_array")
+class MyArray:
+    @property
+    def meaning_of_life(self):
+        return 42
+        """,
+            module.__dict__,
+        )
+
+        array = ak.Array(
+            [None, [{"x": [None, 1, "hi"]}]],
+            with_name="my_array",
+            behavior=module.behavior,
+        )
+        record = ak.Array(
+            [None, [{"x": [None, 1, "hi"]}]],
+            with_name="my_array",
+            behavior=module.behavior,
+        )[1, 0]
+        return pickle.dumps(array), pickle.dumps(record)
+
+
+def global_impl():
+    with temporary_module() as module:
+        exec(
+            """
+import awkward as ak
+
+@ak.mixin_class(ak.behavior, "my_array")
+class MyArray:
+    @property
+    def meaning_of_life(self):
+        return 42
+        """,
+            module.__dict__,
+        )
+
+        array = ak.Array([None, [{"x": [None, 1, "hi"]}]], with_name="my_array")
+        record = ak.Array([None, [{"x": [None, 1, "hi"]}]], with_name="my_array")[1, 0]
+        return pickle.dumps(array), pickle.dumps(record)
+
+
+def test():
+    array_data, record_data = impl()
+
+    # Using a custom behavior dictionary will always break unpickling if
+    # the objects in the dictionary can't be resolved
+    with pytest.raises(ModuleNotFoundError):
+        pickle.loads(array_data)
+
+    with pytest.raises(ModuleNotFoundError):
+        pickle.loads(record_data)
+
+
+def test_global(monkeypatch):
+    with monkeypatch.context() as m:
+        m.setattr(ak, "behavior", {})
+        array_data, record_data = global_impl()
+
+    # The global ak.behavior is not written to the pickle
+    # Awkward can create arrays with missing behavior classes
+    other_array = pickle.loads(array_data)
+    assert other_array.to_list() == [None, [{"x": [None, 1, "hi"]}]]
+    # But it won't have their methods / properties when asked for
+    with pytest.raises(AttributeError):
+        assert other_array.meaning_of_life == 42
+
+    other_record = pickle.loads(record_data)
+    assert other_record.to_list() == {"x": [None, 1, "hi"]}
+    # But it won't have their methods / properties when asked for
+    with pytest.raises(AttributeError):
+        assert other_record.meaning_of_life == 42