Repository: pandas-dev/pandas

Issue:
ENH: Serialize view of ArrowStringArray
Currently Pandas serializes views of ArrowStringArrays by serailizing the whole thing, rather than a subset.  Here is an example:

```python
In [1]: import pandas as pd

In [2]: s = pd.Series([c * 1000 for c in "abcdefghijklmnopqrstuvwxyz"])

In [3]: s
Out[3]: 
0     aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa...
1     bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb...
2     cccccccccccccccccccccccccccccccccccccccccccccc...
3     dddddddddddddddddddddddddddddddddddddddddddddd...
4     eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee...
5     ffffffffffffffffffffffffffffffffffffffffffffff...
6     gggggggggggggggggggggggggggggggggggggggggggggg...
7     hhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhh...
8     iiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiii...
9     jjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjj...
10    kkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkk...
11    llllllllllllllllllllllllllllllllllllllllllllll...
12    mmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmm...
13    nnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnn...
14    oooooooooooooooooooooooooooooooooooooooooooooo...
15    pppppppppppppppppppppppppppppppppppppppppppppp...
16    qqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqq...
17    rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr...
18    ssssssssssssssssssssssssssssssssssssssssssssss...
19    tttttttttttttttttttttttttttttttttttttttttttttt...
20    uuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuu...
21    vvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvvv...
22    wwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwww...
23    xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx...
24    yyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyy...
25    zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz...
dtype: object

In [4]: import pickle

In [5]: len(pickle.dumps(s))
Out[5]: 26758

In [6]: len(pickle.dumps(s.astype("string[pyarrow]")))
Out[6]: 26891

In [7]: len(pickle.dumps(s.head(5)))
Out[7]: 5632

In [8]: len(pickle.dumps(s.astype("string[pyarrow]").head(5)))
Out[8]: 26891
```

This negatively affects dask dataframe operations that cut up pandas dataframes into small pieces, moves them around to different computers, and then pieces them back together again.  


Discussion:
This is also an upstream pyarrow issue:

```python
In [1]: import pyarrow as pa

In [2]: import pickle

In [3]: a = pa.array([c * 1000 for c in "abcdefghijklmnopqrstuvwxyz"])

In [4]: b = a[:5]

In [5]: len(pickle.dumps(a))
Out[5]: 26243

In [6]: len(pickle.dumps(a[:5]))
Out[6]: 26243
```

Seems that @maartenbreddels already noticed this and reported it upstream: https://issues.apache.org/jira/browse/ARROW-10739. It's nice sharing the same underlying technologies :)

I'm going to close this, since I don't think pandas could (easily) work around it and any fix would be better on the pyarrow side (cc @jorisvandenbossche, a bit more motivation if you have a chance to work on that issue).
Oh, I see that the workaround in Python isn't *too* bad. Something like

```python
pa.concat_arrays([self._data])
```

in the `__reduce__` for any arrow-backed array. We can consider implementing that.
Yes, that's the fastest version/workaround I've found.
I am actually using a different version right now: https://github.com/vaexio/vaex/blob/caed2cf106007c6a0141a02a5dfdf823fa38799e/packages/vaex-core/vaex/arrow/convert.py#L260


Patch:
diff --git a/doc/source/whatsnew/v2.0.0.rst b/doc/source/whatsnew/v2.0.0.rst
index e57ba92267855..0024ec0210696 100644
--- a/doc/source/whatsnew/v2.0.0.rst
+++ b/doc/source/whatsnew/v2.0.0.rst
@@ -265,6 +265,7 @@ MultiIndex
 I/O
 ^^^
 - Bug in :func:`read_sas` caused fragmentation of :class:`DataFrame` and raised :class:`.errors.PerformanceWarning` (:issue:`48595`)
+- Bug when a pickling a subset PyArrow-backed data that would serialize the entire data instead of the subset (:issue:`42600`)
 - Bug in :func:`read_csv` for a single-line csv with fewer columns than ``names`` raised :class:`.errors.ParserError` with ``engine="c"`` (:issue:`47566`)
 -
 
diff --git a/pandas/core/arrays/arrow/array.py b/pandas/core/arrays/arrow/array.py
index ea33bf58bacda..f18664915d015 100644
--- a/pandas/core/arrays/arrow/array.py
+++ b/pandas/core/arrays/arrow/array.py
@@ -351,6 +351,17 @@ def __pos__(self: ArrowExtensionArrayT) -> ArrowExtensionArrayT:
     def __abs__(self: ArrowExtensionArrayT) -> ArrowExtensionArrayT:
         return type(self)(pc.abs_checked(self._data))
 
+    # GH 42600: __getstate__/__setstate__ not necessary once
+    # https://issues.apache.org/jira/browse/ARROW-10739 is addressed
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state["_data"] = self._data.combine_chunks()
+        return state
+
+    def __setstate__(self, state) -> None:
+        state["_data"] = pa.chunked_array(state["_data"])
+        self.__dict__.update(state)
+
     def _cmp_method(self, other, op):
         from pandas.arrays import BooleanArray
 


Test patch:
diff --git a/pandas/tests/arrays/string_/test_string_arrow.py b/pandas/tests/arrays/string_/test_string_arrow.py
index 8a6c2b0586a0c..4f0c4daa3c64f 100644
--- a/pandas/tests/arrays/string_/test_string_arrow.py
+++ b/pandas/tests/arrays/string_/test_string_arrow.py
@@ -1,3 +1,4 @@
+import pickle
 import re
 
 import numpy as np
@@ -197,3 +198,20 @@ def test_setitem_invalid_indexer_raises():
 
     with pytest.raises(ValueError, match=None):
         arr[[0, 1]] = ["foo", "bar", "baz"]
+
+
+@skip_if_no_pyarrow
+def test_pickle_roundtrip():
+    # GH 42600
+    expected = pd.Series(range(10), dtype="string[pyarrow]")
+    expected_sliced = expected.head(2)
+    full_pickled = pickle.dumps(expected)
+    sliced_pickled = pickle.dumps(expected_sliced)
+
+    assert len(full_pickled) > len(sliced_pickled)
+
+    result = pickle.loads(full_pickled)
+    tm.assert_series_equal(result, expected)
+
+    result_sliced = pickle.loads(sliced_pickled)
+    tm.assert_series_equal(result_sliced, expected_sliced)
diff --git a/pandas/tests/extension/test_arrow.py b/pandas/tests/extension/test_arrow.py
index 9a6b24583c525..8979c145a223c 100644
--- a/pandas/tests/extension/test_arrow.py
+++ b/pandas/tests/extension/test_arrow.py
@@ -20,6 +20,7 @@
     BytesIO,
     StringIO,
 )
+import pickle
 
 import numpy as np
 import pytest
@@ -1347,3 +1348,19 @@ def test_is_bool_dtype():
     result = s[data]
     expected = s[np.asarray(data)]
     tm.assert_series_equal(result, expected)
+
+
+def test_pickle_roundtrip(data):
+    # GH 42600
+    expected = pd.Series(data)
+    expected_sliced = expected.head(2)
+    full_pickled = pickle.dumps(expected)
+    sliced_pickled = pickle.dumps(expected_sliced)
+
+    assert len(full_pickled) > len(sliced_pickled)
+
+    result = pickle.loads(full_pickled)
+    tm.assert_series_equal(result, expected)
+
+    result_sliced = pickle.loads(sliced_pickled)
+    tm.assert_series_equal(result_sliced, expected_sliced)