Use pickle protocol 5 with NumPy object arrays

[jakirkham]

Depends on PR ( #3868 )

Builds off the work in PR ( #3784 ) to support pickle protocol 5. Supports out-of-band pickling when serializing NumPy object arrays. This can be useful when serializing ragged arrays for example or really any NumPy object arrays containing data that supports out-of-band pickling.

[mrocklin]

Why are we converting into a memoryview here?

[jakirkham]

The buffer_callback is called on values of type PickleBuffer. Here's an example:

In [1]: import pickle                                                           

In [2]: import numpy as np                                                      

In [3]: a = np.array([np.arange(3), np.arange(4, 6)], dtype=object)             

In [4]: l = []; d = pickle.dumps(a, protocol=5, buffer_callback=l.append)       

In [5]: l                                                                       
Out[5]: [<pickle.PickleBuffer at 0x115a64740>, <pickle.PickleBuffer at 0x115a647c0>]

Though PickleBuffer is a wrapper that expects to be passed some bytes-like object, which it wraps zero-copy. In the example above these are the nested NumPy arrays.

However as the rest of the serialization code expects frames to be bytes or memoryviews. We would either need to retool it to support PickleBuffer (with special casing when it is not available) or we can simply convert it to a memoryview, which is also zero-copy.

We've gone for the latter path of converting to a memoryview similar to what was done in PR ( #3784 ). Note this is done without loss of generality.

[jakirkham]

This should be ready on my end. Please let me know if anything else is needed 🙂

[mrocklin]

One more dumb question (sorry for drawing this out): how does this work if the user doesn't have a more modern version of Python? If pickle5 isn't installed and we're using an older version of Python then what happens here?

[jakirkham]

Not at all. The discussion is welcome. 🙂 Was more noting this is not in draft status any more.

In the case where pickle protocol 5 is not supported, we would skip collecting buffers when calling pickle.dumps or cloudpickle.dumps and would skip calling the buffer_callback (as there would be no buffers). As a result we would have only one frame here, which results from calling pickle.dumps or cloudpickle.dumps as before. When it comes to unpickle the data, we would have the first frame as an argument to loads and an empty list/tuple for buffers so would just calling pickle.loads without buffers.

TL;DR it behaves the same as the code does before this change.

[mrocklin]

Ah, ok I was thrown by seeing pickle.dumps(...) without realizing that this was our wrapping around pickle. Grand. Thanks for the detailed explanation @jakirkham .

[jakirkham]

Yeah can see how that would be confusing. Yep, we are importing it here.

Of course! Thanks for the review 😄