Rename

Author: pimdh

python/ray/_private/arrow_serialization.py

@@ -235,17 +235,17 @@ def to_array(self) -> "pyarrow.Array":
         return _array_payload_to_array(self)
 
 
-custom_extension_array_deserializers = {}
+extension_array_deserializers = {}
 
 
-def register_custom_extension_array_deserializer(
+def register_extension_array_deserializer(
     type_cls: Type["pyarrow.DataType"],
     deserializer: Callable[[PicklableArrayPayload], "pyarrow.Array"],
 ):
-    """Add a custom deserializer for a specific Arrow Extension array type."""
-    if type_cls in custom_extension_array_deserializers:
+    """Add a deserializer for a specific Arrow Extension array type."""
+    if type_cls in extension_array_deserializers:
         raise ValueError(f"Deserializer for type {type_cls} is already registered")
-    custom_extension_array_deserializers[type_cls] = deserializer
+    extension_array_deserializers[type_cls] = deserializer
 
 
 def _array_payload_to_array(payload: "PicklableArrayPayload") -> "pyarrow.Array":
@@ -269,8 +269,8 @@ 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, tuple(custom_extension_array_deserializers.keys())):
-        for type_cls, deserializer in custom_extension_array_deserializers.items():
+    elif isinstance(payload.type, tuple(extension_array_deserializers.keys())):
+        for type_cls, deserializer in extension_array_deserializers.items():
             if isinstance(payload.type, type_cls):
                 return deserializer(payload)
         raise ValueError("Unreachable")

python/ray/data/tests/test_arrow_serialization.py

@@ -23,7 +23,7 @@
     _copy_normal_buffer_if_needed,
     _copy_offsets_buffer_if_needed,
     PicklableArrayPayload,
-    register_custom_extension_array_deserializer,
+    register_extension_array_deserializer,
 )
 from ray._private.arrow_utils import get_pyarrow_version
 from ray.data.extensions.object_extension import (
@@ -616,7 +616,7 @@ def deserialize_fixed_shape_tensor_array(
 
 
 def test_fixed_shape_tensor_array_serialization():
-    register_custom_extension_array_deserializer(
+    register_extension_array_deserializer(
         pa.FixedShapeTensorType, deserialize_fixed_shape_tensor_array
     )
     a = pa.FixedShapeTensorArray.from_numpy_ndarray(
@@ -667,7 +667,7 @@ def deserialize_variable_shape_tensor_array(
 
 
 def test_variable_shape_tensor_serialization():
-    register_custom_extension_array_deserializer(
+    register_extension_array_deserializer(
         _VariableShapeTensorType, deserialize_variable_shape_tensor_array
     )
     t = _VariableShapeTensorType(pa.float32(), 2)