apacheGH-38607: [Python] Disable PyExtensionType autoload (apache#38608)

### Rationale for this change

PyExtensionType autoload is really a misfeature. It creates PyArrow-specific extension types, though using ExtensionType is almost the same complexity while allowing deserialization from non-PyArrow software.

### What changes are included in this PR?

* Disable PyExtensionType autoloading and deprecate PyExtensionType instantiation.
* Update the docs to emphasize ExtensionType.

### Are these changes tested?

Yes.

### Are there any user-facing changes?

Yes.

* Closes: apache#38607

Authored-by: Antoine Pitrou <antoine@python.org>
Signed-off-by: Raúl Cumplido <raulcumplido@gmail.com>

docs/source/python/extending_types.rst

@@ -68,34 +68,43 @@ message).
 See the :ref:`format_metadata_extension_types` section of the metadata
 specification for more details.
 
-Pyarrow allows you to define such extension types from Python.
-
-There are currently two ways:
-
-* Subclassing :class:`PyExtensionType`: the (de)serialization is based on pickle.
- This is a good option for an extension type that is only used from Python.
-* Subclassing :class:`ExtensionType`: this allows to give a custom
- Python-independent name and serialized metadata, that can potentially be
- recognized by other (non-Python) Arrow implementations such as PySpark.
+Pyarrow allows you to define such extension types from Python by subclassing
+:class:`ExtensionType` and giving the derived class its own extension name
+and serialization mechanism. The extension name and serialized metadata
+can potentially be recognized by other (non-Python) Arrow implementations
+such as PySpark.
 
 For example, we could define a custom UUID type for 128-bit numbers which can
-be represented as ``FixedSizeBinary`` type with 16 bytes.
-Using the first approach, we create a ``UuidType`` subclass, and implement the
-``__reduce__`` method to ensure the class can be properly pickled::
+be represented as ``FixedSizeBinary`` type with 16 bytes::
 
- class UuidType(pa.PyExtensionType):
+ class UuidType(pa.ExtensionType):
 
  def __init__(self):
- pa.PyExtensionType.__init__(self, pa.binary(16))
+ super().__init__(pa.binary(16), "my_package.uuid")
+
+ def __arrow_ext_serialize__(self):
+ # Since we don't have a parameterized type, we don't need extra
+ # metadata to be deserialized
+ return b''
 
- def __reduce__(self):
- return UuidType, ()
+ @classmethod
+ def __arrow_ext_deserialize__(cls, storage_type, serialized):
+ # Sanity checks, not required but illustrate the method signature.
+ assert storage_type == pa.binary(16)
+ assert serialized == b''
+ # Return an instance of this subclass given the serialized
+ # metadata.
+ return UuidType()
+
+The special methods ``__arrow_ext_serialize__`` and ``__arrow_ext_deserialize__``
+define the serialization of an extension type instance. For non-parametric
+types such as the above, the serialization payload can be left empty.
 
 This can now be used to create arrays and tables holding the extension type::
 
  >>> uuid_type = UuidType()
  >>> uuid_type.extension_name
- 'arrow.py_extension_type'
+ 'my_package.uuid'
  >>> uuid_type.storage_type
  FixedSizeBinaryType(fixed_size_binary[16])
 
@@ -112,8 +121,11 @@ This can now be used to create arrays and tables holding the extension type::
  ]
 
 This array can be included in RecordBatches, sent over IPC and received in
-another Python process. The custom UUID type will be preserved there, as long
-as the definition of the class is available (the type can be unpickled).
+another Python process. The receiving process must explicitly register the
+extension type for deserialization, otherwise it will fall back to the
+storage type::
+
+ >>> pa.register_extension_type(UuidType())
 
 For example, creating a RecordBatch and writing it to a stream using the
 IPC protocol::
@@ -129,43 +141,12 @@ and then reading it back yields the proper type::
  >>> with pa.ipc.open_stream(buf) as reader:
  ... result = reader.read_all()
  >>> result.column('ext').type
- UuidType(extension<arrow.py_extension_type>)
-
-We can define the same type using the other option::
-
- class UuidType(pa.ExtensionType):
-
- def __init__(self):
- pa.ExtensionType.__init__(self, pa.binary(16), "my_package.uuid")
-
- def __arrow_ext_serialize__(self):
- # since we don't have a parameterized type, we don't need extra
- # metadata to be deserialized
- return b''
-
- @classmethod
- def __arrow_ext_deserialize__(self, storage_type, serialized):
- # return an instance of this subclass given the serialized
- # metadata.
- return UuidType()
-
-This is a slightly longer implementation (you need to implement the special
-methods ``__arrow_ext_serialize__`` and ``__arrow_ext_deserialize__``), and the
-extension type needs to be registered to be received through IPC (using
-:func:`register_extension_type`), but it has
-now a unique name::
-
- >>> uuid_type = UuidType()
- >>> uuid_type.extension_name
- 'my_package.uuid'
-
- >>> pa.register_extension_type(uuid_type)
+ UuidType(FixedSizeBinaryType(fixed_size_binary[16]))
 
 The receiving application doesn't need to be Python but can still recognize
-the extension type as a "uuid" type, if it has implemented its own extension
-type to receive it.
-If the type is not registered in the receiving application, it will fall back
-to the storage type.
+the extension type as a "my_package.uuid" type, if it has implemented its own
+extension type to receive it. If the type is not registered in the receiving
+application, it will fall back to the storage type.
 
 Parameterized extension type
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -187,7 +168,7 @@ of the given frequency since 1970.
  # attributes need to be set first before calling
  # super init (as that calls serialize)
  self._freq = freq
- pa.ExtensionType.__init__(self, pa.int64(), 'my_package.period')
+ super().__init__(pa.int64(), 'my_package.period')
 
  @property
  def freq(self):
@@ -198,7 +179,7 @@ of the given frequency since 1970.
 
  @classmethod
  def __arrow_ext_deserialize__(cls, storage_type, serialized):
- # return an instance of this subclass given the serialized
+ # Return an instance of this subclass given the serialized
  # metadata.
  serialized = serialized.decode()
  assert serialized.startswith("freq=")
@@ -209,31 +190,10 @@ Here, we ensure to store all information in the serialized metadata that is
 needed to reconstruct the instance (in the ``__arrow_ext_deserialize__`` class
 method), in this case the frequency string.
 
-Note that, once created, the data type instance is considered immutable. If,
-in the example above, the ``freq`` parameter would change after instantiation,
-the reconstruction of the type instance after IPC will be incorrect.
+Note that, once created, the data type instance is considered immutable.
 In the example above, the ``freq`` parameter is therefore stored in a private
 attribute with a public read-only property to access it.
 
-Parameterized extension types are also possible using the pickle-based type
-subclassing :class:`PyExtensionType`. The equivalent example for the period
-data type from above would look like::
-
- class PeriodType(pa.PyExtensionType):
-
- def __init__(self, freq):
- self._freq = freq
- pa.PyExtensionType.__init__(self, pa.int64())
-
- @property
- def freq(self):
- return self._freq
-
- def __reduce__(self):
- return PeriodType, (self.freq,)
-
-Also the storage type does not need to be fixed but can be parameterized.
-
 Custom extension array class
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -252,12 +212,16 @@ the data as a 2-D Numpy array ``(N, 3)`` without any copy::
  return self.storage.flatten().to_numpy().reshape((-1, 3))
 
 
- class Point3DType(pa.PyExtensionType):
+ class Point3DType(pa.ExtensionType):
  def __init__(self):
- pa.PyExtensionType.__init__(self, pa.list_(pa.float32(), 3))
+ super().__init__(pa.list_(pa.float32(), 3), "my_package.Point3DType")
 
- def __reduce__(self):
- return Point3DType, ()
+ def __arrow_ext_serialize__(self):
+ return b''
+
+ @classmethod
+ def __arrow_ext_deserialize__(cls, storage_type, serialized):
+ return Point3DType()
 
  def __arrow_ext_class__(self):
  return Point3DArray
@@ -289,11 +253,8 @@ The additional methods in the extension class are then available to the user::
 
 
 This array can be sent over IPC, received in another Python process, and the custom
-extension array class will be preserved (as long as the definitions of the classes above
-are available).
-
-The same ``__arrow_ext_class__`` specialization can be used with custom types defined
-by subclassing :class:`ExtensionType`.
+extension array class will be preserved (as long as the receiving process registers
+the extension type using :func:`register_extension_type` before reading the IPC data).
 
 Custom scalar conversion
 ~~~~~~~~~~~~~~~~~~~~~~~~
@@ -304,18 +265,24 @@ If you want scalars of your custom extension type to convert to a custom type wh
 For example, if we wanted the above example 3D point type to return a custom
 3D point class instead of a list, we would implement::
 
+ from collections import namedtuple
+
  Point3D = namedtuple("Point3D", ["x", "y", "z"])
 
  class Point3DScalar(pa.ExtensionScalar):
  def as_py(self) -> Point3D:
  return Point3D(*self.value.as_py())
 
- class Point3DType(pa.PyExtensionType):
+ class Point3DType(pa.ExtensionType):
  def __init__(self):
- pa.PyExtensionType.__init__(self, pa.list_(pa.float32(), 3))
+ super().__init__(pa.list_(pa.float32(), 3), "my_package.Point3DType")
 
- def __reduce__(self):
- return Point3DType, ()
+ def __arrow_ext_serialize__(self):
+ return b''
+
+ @classmethod
+ def __arrow_ext_deserialize__(cls, storage_type, serialized):
+ return Point3DType()
 
  def __arrow_ext_scalar_class__(self):
  return Point3DScalar

python/pyarrow/tests/test_cffi.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import contextlib
 import ctypes
 import gc
 
@@ -51,18 +52,33 @@ def PyCapsule_IsValid(capsule, name):
  return ctypes.pythonapi.PyCapsule_IsValid(ctypes.py_object(capsule), name) == 1
 
 
-class ParamExtType(pa.PyExtensionType):
+@contextlib.contextmanager
+def registered_extension_type(ext_type):
+ pa.register_extension_type(ext_type)
+ try:
+ yield
+ finally:
+ pa.unregister_extension_type(ext_type.extension_name)
+
+
+class ParamExtType(pa.ExtensionType):
 
  def __init__(self, width):
  self._width = width
- pa.PyExtensionType.__init__(self, pa.binary(width))
+ super().__init__(pa.binary(width),
+ "pyarrow.tests.test_cffi.ParamExtType")
 
  @property
  def width(self):
  return self._width
 
- def __reduce__(self):
- return ParamExtType, (self.width,)
+ def __arrow_ext_serialize__(self):
+ return str(self.width).encode()
+
+ @classmethod
+ def __arrow_ext_deserialize__(cls, storage_type, serialized):
+ width = int(serialized.decode())
+ return cls(width)
 
 
 def make_schema():
@@ -75,6 +91,12 @@ def make_extension_schema():
  metadata={b'key1': b'value1'})
 
 
+def make_extension_storage_schema():
+ # Should be kept in sync with make_extension_schema
+ return pa.schema([('ext', ParamExtType(3).storage_type)],
+ metadata={b'key1': b'value1'})
+
+
 def make_batch():
  return pa.record_batch([[[1], [2, 42]]], make_schema())
 
@@ -204,7 +226,10 @@ def test_export_import_array():
  pa.Array._import_from_c(ptr_array, ptr_schema)
 
 
-def check_export_import_schema(schema_factory):
+def check_export_import_schema(schema_factory, expected_schema_factory=None):
+ if expected_schema_factory is None:
+ expected_schema_factory = schema_factory
+
  c_schema = ffi.new("struct ArrowSchema*")
  ptr_schema = int(ffi.cast("uintptr_t", c_schema))
 
@@ -215,7 +240,7 @@ def check_export_import_schema(schema_factory):
  assert pa.total_allocated_bytes() > old_allocated
  # Delete and recreate C++ object from exported pointer
  schema_new = pa.Schema._import_from_c(ptr_schema)
- assert schema_new == schema_factory()
+ assert schema_new == expected_schema_factory()
  assert pa.total_allocated_bytes() == old_allocated
  del schema_new
  assert pa.total_allocated_bytes() == old_allocated
@@ -240,7 +265,13 @@ def test_export_import_schema():
 
 @needs_cffi
 def test_export_import_schema_with_extension():
- check_export_import_schema(make_extension_schema)
+ # Extension type is unregistered => the storage type is imported
+ check_export_import_schema(make_extension_schema,
+ make_extension_storage_schema)
+
+ # Extension type is registered => the extension type is imported
+ with registered_extension_type(ParamExtType(1)):
+ check_export_import_schema(make_extension_schema)
 
 
 @needs_cffi
@@ -319,7 +350,8 @@ def test_export_import_batch():
 
 @needs_cffi
 def test_export_import_batch_with_extension():
- check_export_import_batch(make_extension_batch)
+ with registered_extension_type(ParamExtType(1)):
+ check_export_import_batch(make_extension_batch)
 
 
 def _export_import_batch_reader(ptr_stream, reader_factory):