[Data] (De)serialization of PyArrow Extension Arrays (#51972)

<!-- Thank you for your contribution! Please review
https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before
opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR.
If you don't have the access to it, we will shortly find a reviewer and
assign them to your PR. -->

## 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

- [X] I've signed off every commit(by using the -s flag, i.e., `git
commit -s`) in this PR.
- [X] 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.
- [X] 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
   - [X] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(

<!-- CURSOR_SUMMARY -->
---

> [!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.
> 
> <sup>Written by [Cursor
Bugbot](https://cursor.com/dashboard?tab=bugbot) for commit
4bbcdbe. This will update automatically
on new commits. Configure
[here](https://cursor.com/dashboard?tab=bugbot).</sup>
<!-- /CURSOR_SUMMARY -->

---------

Signed-off-by: Pim de Haan <pim@cusp.ai>

Author: pimdh

python/ray/_private/arrow_serialization.py

@@ -13,8 +13,6 @@
 if TYPE_CHECKING:
     import pyarrow
 
-    from ray.data.extensions import ArrowTensorArray
-
 RAY_DISABLE_CUSTOM_ARROW_JSON_OPTIONS_SERIALIZATION = (
     "RAY_DISABLE_CUSTOM_ARROW_JSON_OPTIONS_SERIALIZATION"
 )
@@ -240,12 +238,8 @@ def _array_payload_to_array(payload: "PicklableArrayPayload") -> "pyarrow.Array"
     """Reconstruct an Arrow Array from a possibly nested PicklableArrayPayload."""
     import pyarrow as pa
 
-    from ray.air.util.tensor_extensions.arrow import get_arrow_extension_tensor_types
-
     children = [child_payload.to_array() for child_payload in payload.children]
 
-    tensor_extension_types = get_arrow_extension_tensor_types()
-
     if pa.types.is_dictionary(payload.type):
         # Dedicated path for reconstructing a DictionaryArray, since
         # Array.from_buffers() doesn't work for DictionaryArrays.
@@ -258,16 +252,10 @@ def _array_payload_to_array(payload: "PicklableArrayPayload") -> "pyarrow.Array"
         assert len(children) == 3, len(children)
         offsets, keys, items = children
         return pa.MapArray.from_arrays(offsets, keys, items)
-    elif isinstance(
-        payload.type,
-        tensor_extension_types,
-    ):
-        # Dedicated path for reconstructing an ArrowTensorArray or
-        # ArrowVariableShapedTensorArray, both of which can't be reconstructed by the
-        # Array.from_buffers() API.
+    elif isinstance(payload.type, pa.BaseExtensionType):
         assert len(children) == 1, len(children)
         storage = children[0]
-        return pa.ExtensionArray.from_storage(payload.type, storage)
+        return payload.type.wrap_array(storage)
     else:
         # Common case: use Array.from_buffers() to construct an array of a certain type.
         return pa.Array.from_buffers(
@@ -288,10 +276,6 @@ def _array_to_array_payload(a: "pyarrow.Array") -> "PicklableArrayPayload":
     """
     import pyarrow as pa
 
-    from ray.air.util.tensor_extensions.arrow import get_arrow_extension_tensor_types
-
-    tensor_extension_types = get_arrow_extension_tensor_types()
-
     if _is_dense_union(a.type):
         # Dense unions are not supported.
         # TODO(Clark): Support dense unions.
@@ -319,9 +303,7 @@ def _array_to_array_payload(a: "pyarrow.Array") -> "PicklableArrayPayload":
         return _dictionary_array_to_array_payload(a)
     elif pa.types.is_map(a.type):
         return _map_array_to_array_payload(a)
-    elif isinstance(a.type, tensor_extension_types):
-        return _tensor_array_to_array_payload(a)
-    elif isinstance(a.type, pa.ExtensionType):
+    elif isinstance(a.type, pa.BaseExtensionType):
         return _extension_array_to_array_payload(a)
     else:
         raise ValueError("Unhandled Arrow array type:", a.type)
@@ -630,11 +612,9 @@ def _map_array_to_array_payload(a: "pyarrow.MapArray") -> "PicklableArrayPayload
     )
 
 
-def _tensor_array_to_array_payload(a: "ArrowTensorArray") -> "PicklableArrayPayload":
-    """Serialize tensor arrays to PicklableArrayPayload."""
-    # Offset is propagated to storage array, and the storage array items align with the
-    # tensor elements, so we only need to do the straightforward creation of the storage
-    # array payload.
+def _extension_array_to_array_payload(
+    a: "pyarrow.ExtensionArray",
+) -> "PicklableArrayPayload":
     storage_payload = _array_to_array_payload(a.storage)
     return PicklableArrayPayload(
         type=a.type,
@@ -646,16 +626,6 @@ def _tensor_array_to_array_payload(a: "ArrowTensorArray") -> "PicklableArrayPayl
     )
 
 
-def _extension_array_to_array_payload(
-    a: "pyarrow.ExtensionArray",
-) -> "PicklableArrayPayload":
-    payload = _array_to_array_payload(a.storage)
-    payload.type = a.type
-    payload.length = len(a)
-    payload.null_count = a.null_count
-    return payload
-
-
 def _copy_buffer_if_needed(
     buf: "pyarrow.Buffer",
     type_: Optional["pyarrow.DataType"],

python/ray/data/tests/test_arrow_serialization.py

@@ -16,6 +16,7 @@
 import ray.data
 import ray.train
 from ray._private.arrow_serialization import (
+    PicklableArrayPayload,
     _align_bit_offset,
     _bytes_for_bits,
     _copy_bitpacked_buffer_if_needed,
@@ -595,3 +596,60 @@ def test_custom_arrow_data_serializer_disable(shutdown_only):
     assert d_view["a"].chunk(0).buffers()[1].size == t["a"].chunk(0).buffers()[1].size
     # Check that the serialized slice view is large
     assert len(s_view) > 0.8 * len(s_t)
+
+
+def test_fixed_shape_tensor_array_serialization():
+    a = pa.FixedShapeTensorArray.from_numpy_ndarray(
+        np.arange(4 * 2 * 3).reshape(4, 2, 3)
+    )
+    payload = PicklableArrayPayload.from_array(a)
+    a2 = payload.to_array()
+    assert a == a2
+
+
+class _VariableShapeTensorType(pa.ExtensionType):
+    def __init__(
+        self,
+        value_type: pa.DataType,
+        ndim: int,
+    ) -> None:
+        self.value_type = value_type
+        self.ndim = ndim
+        super().__init__(
+            pa.struct(
+                [
+                    pa.field("data", pa.list_(value_type)),
+                    pa.field("shape", pa.list_(pa.int32(), ndim)),
+                ]
+            ),
+            "variable_shape_tensor",
+        )
+
+    def __arrow_ext_serialize__(self) -> bytes:
+        return b""
+
+    @classmethod
+    def __arrow_ext_deserialize__(cls, storage_type: pa.DataType, serialized: bytes):
+        ndim = storage_type[1].type.list_size
+        value_type = storage_type[0].type.value_type
+        return cls(value_type, ndim)
+
+
+def test_variable_shape_tensor_serialization():
+    t = _VariableShapeTensorType(pa.float32(), 2)
+    ar = pa.array(
+        [
+            {
+                "data": np.arange(2 * 3),
+                "shape": [2, 3],
+            },
+            {
+                "data": np.arange(4 * 5),
+                "shape": [4, 5],
+            },
+        ],
+        type=t,
+    )
+    payload = PicklableArrayPayload.from_array(ar)
+    ar2 = payload.to_array()
+    assert ar == ar2