[Data] (De)serialization of PyArrow Extension Arrays

[pimdh]

Why are these changes needed?

This feature adds the ability to (de)serialize arbitrary PyArrow extension arrays. This is needed to use Ray in code bases that use extension arrays.

The serialization already seemed sufficiently general, but as far as I can tell, the deserialization can not be done in generality. Hence, this setup allows registration of custom deserializers for extension types.

For serialization, the selector has been changed from ExtensionType to BaseExtensionType to accommodate for non-Python ExtensionArrays, like pyarrow.FixedShapeTensorArray.

This is at the moment a proof-of-concept. If you like the idea, I suppose the registration function may need to move to a better place, and docs need adding.

The implementation now works without registration on any extension type.

Related issue number

Closes #51959

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

---

Note

Generalizes Arrow array (de)serialization to any pyarrow.BaseExtensionType, removing tensor-specific handling and adding tests for fixed/variable-shape tensors.

• Arrow (De)serialization:
  • Switch from tensor-specific checks to generic pyarrow.BaseExtensionType handling.
  • Reconstruct extension arrays via type.wrap_array(storage); serialize via storage payload wrapped with extension metadata.
  • Remove ray.air.util.tensor_extensions.arrow dependencies and special-casing.
• Tests:
  • Add roundtrip tests for pa.FixedShapeTensorArray and a custom variable-shape ExtensionType.
  • Import PicklableArrayPayload in tests for constructing payloads.

^{Written by Cursor Bugbot for commit 4bbcdbe. This will update automatically on new commits. Configure here.}

[richardliaw]

hmm, would we need to expose this as a user API?

[pimdh]

I think so, if we'd want to allow a user to use their own ExtensionTypes. If the public api should live elsewhere/look different, please let me know and I'll change it.

[israbbani]

@pimdh thank you for your contribution. We have an API for registering custom serializers and deserializers (I think it's what is used internally to register the arrow serializer/deserializer too).

Can that be used to add serializers and deserializers for PyArrow Extension Arrays? See Example 2 in https://docs.ray.io/en/latest/ray-core/objects/serialization.html#customized-serialization for more details.

[pimdh]

Hi,
As I understand it, no that doesn't work, because this needs a modification of the built-in Arrow Table serializer. Perhaps one could replace that entire serializer with a custom one, but that isn't particularly user-friendly.
Thanks!

[israbbani]

@pimdh I see your point. I'm adding @alexeykudinkin from the Ray Data team to shepherd this.

[alexeykudinkin]

Thank you for your contribution @pimdh!

Can you please help me understand what kind of extensions you folks are using?

[alexeykudinkin]

Let's actually lift this support to happen outside of the extension mechanism -- Ray should just support FixedShapeTensorArray not requiring any extensions for it

[pimdh]

I'm using the built-in FixedShapeTensorArray, and a custom variable shape tensor array, which seems equivalent to the Ray one (and the one in the PR in the arrow repo).
I'm reading and writing with Lance, so the array type affects the storage format. Hence I'm not so keen to migrate all my code to Ray's internal extension type.

[pimdh]

Another representative example is this one for encoding bfloat16 arrays: https://github.com/lancedb/lance/blob/main/python/python/lance/_arrow/bf16.py

[alexeykudinkin]

Is there a particular reason you're adding support for FixedShapeTensorType?

It should just work out of the box.

[alexeykudinkin]

I've just tested it and it works out of the box:

import ray
import numpy as np
import pyarrow as pa
from pyarrow import FixedShapeTensorType


@ray.remote
def _gen_tensor():
    shape = (16, 16, 8)  # Shape per tensor
    num_tensors = 128  # Number of tensors
    tensors = [np.ones(shape, dtype=np.int8).flatten() for _ in range(num_tensors)]

    # Create array of fixed shape tensors
    return pa.array(tensors, type=pa.fixed_shape_tensor(pa.int8(), shape=shape))


tensors = ray.get(_gen_tensor.remote())

print(tensors)
print(tensors.type)

[pimdh]

Anyone able to help me land this? Thanks!

[github-actions]

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

[iamjustinhsu]

Hi @pimdh, looks fine to me. I'll await @alexeykudinkin response since he's the expert on this

[iamjustinhsu]

_tensor_array_to_array_payload is this still being used?

[pimdh]

No it isn't, thanks. Removed

[alexeykudinkin]

Thanks for making the changes @pimdh!

And special thanks to @mdekstrand for chiming in and getting everyone aligned!

[pimdh]

It looks like the CI failed due to an unrelated reason. @alexeykudinkin , would you mind restarting CI? Thanks

[Sparks0219]

LGTM