Merge branch 'main' into agoose77/refactor-highlevel-typetracer

pyproject.toml

@@ -27,7 +27,7 @@ classifiers = [
   "Topic :: Scientific/Engineering",
 ]
 dependencies = [
-  "awkward >=2.2.4",
+  "awkward >=2.4.0",
   "dask >=2023.04.0",
 ]
 dynamic = ["version"]
@@ -66,6 +66,9 @@ test = [
 [project.entry-points."dask.sizeof"]
 awkward = "dask_awkward.sizeof:register"
 
+[project.entry-points."awkward.pickle.reduce"]
+dask_awkward = "dask_awkward.pickle:plugin"
+
 [tool.hatch.version]
 source = "vcs"
 path = "src/dask_awkward/__init__.py"

src/dask_awkward/pickle.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+__all__ = ("plugin",)
+
+import pickle
+
+import awkward as ak
+
+
+def pickle_record(record: ak.Record, protocol: int) -> tuple:
+    layout = ak.to_layout(record, allow_record=True)
+    form, length, container = ak.operations.to_buffers(
+        layout.array,
+        buffer_key="{form_key}-{attribute}",
+        form_key="node{id}",
+        byteorder="<",
+    )
+
+    # For pickle >= 5, we can avoid copying the buffers
+    if protocol >= 5:
+        container = {k: pickle.PickleBuffer(v) for k, v in container.items()}
+
+    if record.behavior is ak.behavior:
+        behavior = None
+    else:
+        behavior = record.behavior
+
+    return (
+        object.__new__,
+        (ak.Record,),
+        (form.to_dict(), length, container, behavior, layout.at),
+    )
+
+
+def pickle_array(array: ak.Array, protocol: int) -> tuple:
+    layout = ak.to_layout(array, allow_record=False)
+    form, length, container = ak.operations.to_buffers(
+        layout,
+        buffer_key="{form_key}-{attribute}",
+        form_key="node{id}",
+        byteorder="<",
+    )
+
+    # For pickle >= 5, we can avoid copying the buffers
+    if protocol >= 5:
+        container = {k: pickle.PickleBuffer(v) for k, v in container.items()}
+
+    if array.behavior is ak.behavior:
+        behavior = None
+    else:
+        behavior = array.behavior
+
+    return (
+        object.__new__,
+        (ak.Array,),
+        (form.to_dict(), length, container, behavior),
+    )
+
+
+def plugin(obj, protocol: int) -> tuple | NotImplemented:
+    if isinstance(obj, ak.Record):
+        return pickle_record(obj, protocol)
+    elif isinstance(obj, ak.Array):
+        return pickle_array(obj, protocol)
+    else:
+        return NotImplemented

tests/test_pickle.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import pickle
+
+import awkward as ak
+import numpy as np
+
+
+def test_pickle_ak_array():
+    buffers = []
+    array = ak.Array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])[[0, 2]]
+    next_array = pickle.loads(
+        pickle.dumps(array, protocol=5, buffer_callback=buffers.append), buffers=buffers
+    )
+    assert ak.almost_equal(array, next_array)
+    assert array.layout.form == next_array.layout.form
+    assert buffers
+    assert np.shares_memory(
+        array.layout.content.data,
+        next_array.layout.content.data,
+    )
+    assert np.shares_memory(
+        array.layout.starts.data,
+        next_array.layout.starts.data,
+    )
+    assert np.shares_memory(
+        array.layout.stops.data,
+        next_array.layout.stops.data,
+    )
+
+
+def test_pickle_ak_record():
+    buffers = []
+    record = ak.zip(
+        {"x": [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]}, depth_limit=1
+    )[2]
+    next_record = pickle.loads(
+        pickle.dumps(record, protocol=5, buffer_callback=buffers.append),
+        buffers=buffers,
+    )
+    assert record.layout.at == next_record.layout.at
+
+    array = ak.Array(record.layout.array)
+    next_array = ak.Array(next_record.layout.array)
+
+    assert buffers
+    assert np.shares_memory(
+        array.layout.content("x").content.data,
+        next_array.layout.content("x").content.data,
+    )
+    assert np.shares_memory(
+        array.layout.content("x").starts.data,
+        next_array.layout.content("x").starts.data,
+    )
+    assert np.shares_memory(
+        array.layout.content("x").stops.data,
+        next_array.layout.content("x").stops.data,
+    )