Skip to content

Commit 23c3dc2

Browse files
authored
BUG: pickling ArrowExtensionArray with _pa_data rename (#52027)
1 parent 4ef638a commit 23c3dc2

File tree

2 files changed

+23
-2
lines changed

2 files changed

+23
-2
lines changed

pandas/core/arrays/arrow/array.py

+6-2
Original file line numberDiff line numberDiff line change
@@ -433,11 +433,15 @@ def __abs__(self) -> Self:
433433
# https://issues.apache.org/jira/browse/ARROW-10739 is addressed
434434
def __getstate__(self):
435435
state = self.__dict__.copy()
436-
state["_data"] = self._pa_array.combine_chunks()
436+
state["_pa_array"] = self._pa_array.combine_chunks()
437437
return state
438438

439439
def __setstate__(self, state) -> None:
440-
state["_pa_array"] = pa.chunked_array(state["_data"])
440+
if "_data" in state:
441+
data = state.pop("_data")
442+
else:
443+
data = state["_pa_array"]
444+
state["_pa_array"] = pa.chunked_array(data)
441445
self.__dict__.update(state)
442446

443447
def _cmp_method(self, other, op):

pandas/tests/extension/test_arrow.py

+17
Original file line numberDiff line numberDiff line change
@@ -2351,3 +2351,20 @@ def test_concat_empty_arrow_backed_series(dtype):
23512351
expected = ser.copy()
23522352
result = pd.concat([ser[np.array([], dtype=np.bool_)]])
23532353
tm.assert_series_equal(result, expected)
2354+
2355+
2356+
# _data was renamed to _pa_data
2357+
class OldArrowExtensionArray(ArrowExtensionArray):
2358+
def __getstate__(self):
2359+
state = super().__getstate__()
2360+
state["_data"] = state.pop("_pa_array")
2361+
return state
2362+
2363+
2364+
def test_pickle_old_arrowextensionarray():
2365+
data = pa.array([1])
2366+
expected = OldArrowExtensionArray(data)
2367+
result = pickle.loads(pickle.dumps(expected))
2368+
tm.assert_extension_array_equal(result, expected)
2369+
assert result._pa_array == pa.chunked_array(data)
2370+
assert not hasattr(result, "_data")

0 commit comments

Comments
 (0)