Skip to content

Commit f865fbd

Browse files
authored
GH-36038: [Python] Implement __reduce__ on ExtensionType class (#36170)
### Rationale for this change `ExtensionType` subclasses can't be pickled if `__reduce__` method isn't implemented separately. ### What changes are included in this PR? Add `__reduce__` method to `ExtensionType` class. ### Are these changes tested? Yes, test is added to python/pyarrow/tests/test_extension_type.py. ### Are there any user-facing changes? No. * Closes: #36038 Authored-by: AlenkaF <frim.alenka@gmail.com> Signed-off-by: Joris Van den Bossche <jorisvandenbossche@gmail.com>
1 parent 1e00e7a commit f865fbd

File tree

2 files changed

+28
-0
lines changed

2 files changed

+28
-0
lines changed

python/pyarrow/tests/test_extension_type.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -867,6 +867,31 @@ def test_generic_ext_type_equality():
867867
assert not period_type == period_type3
868868

869869

870+
def test_generic_ext_type_pickling(registered_period_type):
871+
# GH-36038
872+
for proto in range(0, pickle.HIGHEST_PROTOCOL + 1):
873+
period_type, _ = registered_period_type
874+
ser = pickle.dumps(period_type, protocol=proto)
875+
period_type_pickled = pickle.loads(ser)
876+
assert period_type == period_type_pickled
877+
878+
879+
def test_generic_ext_array_pickling(registered_period_type):
880+
for proto in range(0, pickle.HIGHEST_PROTOCOL + 1):
881+
period_type, _ = registered_period_type
882+
storage = pa.array([1, 2, 3, 4], pa.int64())
883+
arr = pa.ExtensionArray.from_storage(period_type, storage)
884+
ser = pickle.dumps(arr, protocol=proto)
885+
del storage, arr
886+
arr = pickle.loads(ser)
887+
arr.validate()
888+
assert isinstance(arr, pa.ExtensionArray)
889+
assert arr.type == period_type
890+
assert arr.type.storage_type == pa.int64()
891+
assert arr.storage.type == pa.int64()
892+
assert arr.storage.to_pylist() == [1, 2, 3, 4]
893+
894+
870895
def test_generic_ext_type_register(registered_period_type):
871896
# test that trying to register other type does not segfault
872897
with pytest.raises(TypeError):

python/pyarrow/types.pxi

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1487,6 +1487,9 @@ cdef class ExtensionType(BaseExtensionType):
14871487
"""
14881488
return NotImplementedError
14891489

1490+
def __reduce__(self):
1491+
return self.__arrow_ext_deserialize__, (self.storage_type, self.__arrow_ext_serialize__())
1492+
14901493
def __arrow_ext_class__(self):
14911494
"""Return an extension array class to be used for building or
14921495
deserializing arrays with this extension type.

0 commit comments

Comments
 (0)