Skip to content

Commit 5834dd0

Browse files
committed
fix: preserve series name in ListAccessor
1 parent 29d7e08 commit 5834dd0

File tree

2 files changed

+30
-5
lines changed

2 files changed

+30
-5
lines changed

pandas/core/arrays/arrow/accessors.py

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,10 @@ def len(self) -> Series:
117117

118118
value_lengths = pc.list_value_length(self._pa_array)
119119
return Series(
120-
value_lengths, dtype=ArrowDtype(value_lengths.type), index=self._data.index
120+
value_lengths,
121+
dtype=ArrowDtype(value_lengths.type),
122+
index=self._data.index,
123+
name=self._data.name,
121124
)
122125

123126
def __getitem__(self, key: int | slice) -> Series:
@@ -181,7 +184,12 @@ def __getitem__(self, key: int | slice) -> Series:
181184
if step is None:
182185
step = 1
183186
sliced = pc.list_slice(self._pa_array, start, stop, step)
184-
return Series(sliced, dtype=ArrowDtype(sliced.type), index=self._data.index)
187+
return Series(
188+
sliced,
189+
dtype=ArrowDtype(sliced.type),
190+
index=self._data.index,
191+
name=self._data.name,
192+
)
185193
else:
186194
raise ValueError(f"key must be an int or slice, got {type(key).__name__}")
187195

@@ -223,7 +231,12 @@ def flatten(self) -> Series:
223231
counts = pa.compute.list_value_length(self._pa_array)
224232
flattened = pa.compute.list_flatten(self._pa_array)
225233
index = self._data.index.repeat(counts.fill_null(pa.scalar(0, counts.type)))
226-
return Series(flattened, dtype=ArrowDtype(flattened.type), index=index)
234+
return Series(
235+
flattened,
236+
dtype=ArrowDtype(flattened.type),
237+
index=index,
238+
name=self._data.name,
239+
)
227240

228241

229242
class StructAccessor(ArrowAccessor):

pandas/tests/series/accessors/test_list_accessor.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ def test_list_getitem(list_dtype):
2525
ser = Series(
2626
[[1, 2, 3], [4, None, 5], None],
2727
dtype=ArrowDtype(list_dtype),
28+
name="a",
2829
)
2930
actual = ser.list[1]
3031
expected = Series([2, None, None], dtype="int64[pyarrow]")
@@ -37,9 +38,15 @@ def test_list_getitem_index():
3738
[[1, 2, 3], [4, None, 5], None],
3839
dtype=ArrowDtype(pa.list_(pa.int64())),
3940
index=[1, 3, 7],
41+
name="a",
4042
)
4143
actual = ser.list[1]
42-
expected = Series([2, None, None], dtype="int64[pyarrow]", index=[1, 3, 7])
44+
expected = Series(
45+
[2, None, None],
46+
dtype="int64[pyarrow]",
47+
index=[1, 3, 7],
48+
name="a",
49+
)
4350
tm.assert_series_equal(actual, expected)
4451

4552

@@ -48,6 +55,7 @@ def test_list_getitem_slice():
4855
[[1, 2, 3], [4, None, 5], None],
4956
dtype=ArrowDtype(pa.list_(pa.int64())),
5057
index=[1, 3, 7],
58+
name="a",
5159
)
5260
if pa_version_under11p0:
5361
with pytest.raises(
@@ -60,6 +68,7 @@ def test_list_getitem_slice():
6068
[[2, 3], [None, 5], None],
6169
dtype=ArrowDtype(pa.list_(pa.int64())),
6270
index=[1, 3, 7],
71+
name="a",
6372
)
6473
tm.assert_series_equal(actual, expected)
6574

@@ -68,22 +77,25 @@ def test_list_len():
6877
ser = Series(
6978
[[1, 2, 3], [4, None], None],
7079
dtype=ArrowDtype(pa.list_(pa.int64())),
80+
name="a",
7181
)
7282
actual = ser.list.len()
73-
expected = Series([3, 2, None], dtype=ArrowDtype(pa.int32()))
83+
expected = Series([3, 2, None], dtype=ArrowDtype(pa.int32()), name="a")
7484
tm.assert_series_equal(actual, expected)
7585

7686

7787
def test_list_flatten():
7888
ser = Series(
7989
[[1, 2, 3], None, [4, None], [], [7, 8]],
8090
dtype=ArrowDtype(pa.list_(pa.int64())),
91+
name="a",
8192
)
8293
actual = ser.list.flatten()
8394
expected = Series(
8495
[1, 2, 3, 4, None, 7, 8],
8596
dtype=ArrowDtype(pa.int64()),
8697
index=[0, 0, 0, 2, 2, 4, 4],
98+
name="a",
8799
)
88100
tm.assert_series_equal(actual, expected)
89101

0 commit comments

Comments
 (0)