feat: add support for custom picklers

[agoose77]

This PR adds support for customisation of the pickling __reduce__ implementation by third-party libraries. The intention is to provide a mechanism for dask-awkward to intercept serialisation in a form-preserving manner c.f. dask-contrib/dask-awkward#344.

The problem in dask-contrib/dask-awkward#344 is that whilst regular users of Awkward Array are likely not that concerned about the form of an array, Dask requires guarantees that the form does not change during transfer between workers. As such, we need an environment-sensitive serialisation mechanism.

The simplest/decoupled way of doing this is to let dask-awkward handle the serialisation, rather than internally making Awkward aware of dask-awkward directly.

With this PR it is possible to define an awkward.pickle.reduce entrypoint that exposes a __reduce_ex__ implementation, e.g. with hatchling:

[project.entry-points."awkward.pickle.reduce"]
dask-awkward-pickler = "dask_awkward.pickle:plugin"

where

def plugin(obj, protocol: int):
    if isinstance(obj, (ak.Array, ak.Record)):
        ...
    else:
        # Allow future extensions of this mechanism
        return NotImplemented

[codecov]

Codecov Report

Merging #2682 (035c891) into main (59d4235) will increase coverage by 0.00%.
The diff coverage is 85.41%.

Additional details and impacted files

Files Changed	Coverage Δ
src/awkward/_pickle.py	85.36% <85.36%> (ø)
src/awkward/highlevel.py	76.38% <85.71%> (+0.09%)	⬆️

[jpivarski]

This is good, letting dask-awkward have the responsibility of implementing the serialization. Awkward can't know about all of its dependent projects (though dask-awkward is certainly one that we do know about!), but dependent projects by definition all know about Awkward.

So yes, I think this is the right way to go.

If we eventually do have a Form-preserving packing algorithm, we can provide that centrally in Awkward so that dependent projects that do override serialization can choose that as an option. It would be a performance improvement, but not a correctness one, so not the highest priority.

[jpivarski]

Oh, and this is ready to merge!