Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""Accessors for arrow-backed data."""
3from __future__ import annotations
5from abc import (
6 ABCMeta,
7 abstractmethod,
8)
9from typing import (
10 TYPE_CHECKING,
11 cast,
12)
14from pandas.compat import HAS_PYARROW
16from pandas.core.dtypes.common import is_list_like
18if HAS_PYARROW:
19 import pyarrow as pa
20 import pyarrow.compute as pc
22 from pandas.core.dtypes.dtypes import ArrowDtype
24if TYPE_CHECKING:
25 from collections.abc import Iterator
27 from pandas import (
28 DataFrame,
29 Series,
30 )
33class ArrowAccessor(metaclass=ABCMeta):
34 @abstractmethod
35 def __init__(self, data, validation_msg: str) -> None:
36 self._data = data
37 self._validation_msg = validation_msg
38 self._validate(data)
40 @abstractmethod
41 def _is_valid_pyarrow_dtype(self, pyarrow_dtype) -> bool:
42 pass
44 def _validate(self, data) -> None:
45 dtype = data.dtype
46 if not HAS_PYARROW or not isinstance(dtype, ArrowDtype):
47 # Raise AttributeError so that inspect can handle non-struct Series.
48 raise AttributeError(self._validation_msg.format(dtype=dtype))
50 if not self._is_valid_pyarrow_dtype(dtype.pyarrow_dtype):
51 # Raise AttributeError so that inspect can handle invalid Series.
52 raise AttributeError(self._validation_msg.format(dtype=dtype))
54 @property
55 def _pa_array(self):
56 return self._data.array._pa_array
59class ListAccessor(ArrowAccessor):
60 """
61 Accessor object for list data properties of the Series values.
63 Parameters
64 ----------
65 data : Series
66 Series containing Arrow list data.
67 """
69 def __init__(self, data=None) -> None:
70 super().__init__(
71 data,
72 validation_msg="Can only use the '.list' accessor with "
73 "'list[pyarrow]' dtype, not {dtype}.",
74 )
76 def _is_valid_pyarrow_dtype(self, pyarrow_dtype) -> bool:
77 return (
78 pa.types.is_list(pyarrow_dtype)
79 or pa.types.is_fixed_size_list(pyarrow_dtype)
80 or pa.types.is_large_list(pyarrow_dtype)
81 )
83 def len(self) -> Series:
84 """
85 Return the length of each list in the Series.
87 Returns
88 -------
89 pandas.Series
90 The length of each list.
92 See Also
93 --------
94 str.len : Python built-in function returning the length of an object.
95 Series.size : Returns the length of the Series.
96 StringMethods.len : Compute the length of each element in the Series/Index.
98 Examples
99 --------
100 >>> import pyarrow as pa
101 >>> s = pd.Series(
102 ... [
103 ... [1, 2, 3],
104 ... [3],
105 ... ],
106 ... dtype=pd.ArrowDtype(pa.list_(pa.int64())),
107 ... )
108 >>> s.list.len()
109 0 3
110 1 1
111 dtype: int32[pyarrow]
112 """
113 from pandas import Series
115 value_lengths = pc.list_value_length(self._pa_array)
116 return Series(
117 value_lengths,
118 dtype=ArrowDtype(value_lengths.type),
119 index=self._data.index,
120 name=self._data.name,
121 )
123 def __getitem__(self, key: int | slice) -> Series:
124 """
125 Index or slice lists in the Series.
127 Parameters
128 ----------
129 key : int | slice
130 Index or slice of indices to access from each list.
132 Returns
133 -------
134 pandas.Series
135 The list at requested index.
137 See Also
138 --------
139 ListAccessor.flatten : Flatten list values.
141 Examples
142 --------
143 >>> import pyarrow as pa
144 >>> s = pd.Series(
145 ... [
146 ... [1, 2, 3],
147 ... [3],
148 ... ],
149 ... dtype=pd.ArrowDtype(pa.list_(pa.int64())),
150 ... )
151 >>> s.list[0]
152 0 1
153 1 3
154 dtype: int64[pyarrow]
155 """
156 from pandas import Series
158 if isinstance(key, int):
159 # TODO: Support negative key but pyarrow does not allow
160 # element index to be an array.
161 # if key < 0:
162 # key = pc.add(key, pc.list_value_length(self._pa_array))
163 element = pc.list_element(self._pa_array, key)
164 return Series(
165 element,
166 dtype=ArrowDtype(element.type),
167 index=self._data.index,
168 name=self._data.name,
169 )
170 elif isinstance(key, slice):
171 # TODO: Support negative start/stop/step, ideally this would be added
172 # upstream in pyarrow.
173 start, stop, step = key.start, key.stop, key.step
174 if start is None:
175 # TODO: When adding negative step support
176 # this should be setto last element of array
177 # when step is negative.
178 start = 0
179 if step is None:
180 step = 1
181 sliced = pc.list_slice(self._pa_array, start, stop, step)
182 return Series(
183 sliced,
184 dtype=ArrowDtype(sliced.type),
185 index=self._data.index,
186 name=self._data.name,
187 )
188 else:
189 raise ValueError(f"key must be an int or slice, got {type(key).__name__}")
191 def __iter__(self) -> Iterator:
192 raise TypeError(f"'{type(self).__name__}' object is not iterable")
194 def flatten(self) -> Series:
195 """
196 Flatten list values.
198 Returns
199 -------
200 pandas.Series
201 The data from all lists in the series flattened.
203 See Also
204 --------
205 ListAccessor.__getitem__ : Index or slice values in the Series.
207 Examples
208 --------
209 >>> import pyarrow as pa
210 >>> s = pd.Series(
211 ... [
212 ... [1, 2, 3],
213 ... [3],
214 ... ],
215 ... dtype=pd.ArrowDtype(pa.list_(pa.int64())),
216 ... )
217 >>> s.list.flatten()
218 0 1
219 0 2
220 0 3
221 1 3
222 dtype: int64[pyarrow]
223 """
224 from pandas import Series
226 counts = pa.compute.list_value_length(self._pa_array)
227 flattened = pa.compute.list_flatten(self._pa_array)
228 index = self._data.index.repeat(counts.fill_null(pa.scalar(0, counts.type)))
229 return Series(
230 flattened,
231 dtype=ArrowDtype(flattened.type),
232 index=index,
233 name=self._data.name,
234 )
237class StructAccessor(ArrowAccessor):
238 """
239 Accessor object for structured data properties of the Series values.
241 Parameters
242 ----------
243 data : Series
244 Series containing Arrow struct data.
245 """
247 def __init__(self, data=None) -> None:
248 super().__init__(
249 data,
250 validation_msg=(
251 "Can only use the '.struct' accessor with 'struct[pyarrow]' "
252 "dtype, not {dtype}."
253 ),
254 )
256 def _is_valid_pyarrow_dtype(self, pyarrow_dtype) -> bool:
257 return pa.types.is_struct(pyarrow_dtype)
259 @property
260 def dtypes(self) -> Series:
261 """
262 Return the dtype object of each child field of the struct.
264 Returns
265 -------
266 pandas.Series
267 The data type of each child field.
269 See Also
270 --------
271 Series.dtype: Return the dtype object of the underlying data.
273 Examples
274 --------
275 >>> import pyarrow as pa
276 >>> s = pd.Series(
277 ... [
278 ... {"version": 1, "project": "pandas"},
279 ... {"version": 2, "project": "pandas"},
280 ... {"version": 1, "project": "numpy"},
281 ... ],
282 ... dtype=pd.ArrowDtype(
283 ... pa.struct([("version", pa.int64()), ("project", pa.string())])
284 ... ),
285 ... )
286 >>> s.struct.dtypes
287 version int64[pyarrow]
288 project string[pyarrow]
289 dtype: object
290 """
291 from pandas import (
292 Index,
293 Series,
294 )
296 pa_type = self._data.dtype.pyarrow_dtype
297 types = [ArrowDtype(struct.type) for struct in pa_type]
298 names = [struct.name for struct in pa_type]
299 return Series(types, index=Index(names))
301 def field(
302 self,
303 name_or_index: list[str]
304 | list[bytes]
305 | list[int]
306 | pc.Expression
307 | bytes
308 | str
309 | int,
310 ) -> Series:
311 """
312 Extract a child field of a struct as a Series.
314 Parameters
315 ----------
316 name_or_index : str | bytes | int | expression | list
317 Name or index of the child field to extract.
319 For list-like inputs, this will index into a nested
320 struct.
322 Returns
323 -------
324 pandas.Series
325 The data corresponding to the selected child field.
327 See Also
328 --------
329 Series.struct.explode : Return all child fields as a DataFrame.
331 Notes
332 -----
333 The name of the resulting Series will be set using the following
334 rules:
336 - For string, bytes, or integer `name_or_index` (or a list of these, for
337 a nested selection), the Series name is set to the selected
338 field's name.
339 - For a :class:`pyarrow.compute.Expression`, this is set to
340 the string form of the expression.
341 - For list-like `name_or_index`, the name will be set to the
342 name of the final field selected.
344 Examples
345 --------
346 >>> import pyarrow as pa
347 >>> s = pd.Series(
348 ... [
349 ... {"version": 1, "project": "pandas"},
350 ... {"version": 2, "project": "pandas"},
351 ... {"version": 1, "project": "numpy"},
352 ... ],
353 ... dtype=pd.ArrowDtype(
354 ... pa.struct([("version", pa.int64()), ("project", pa.string())])
355 ... ),
356 ... )
358 Extract by field name.
360 >>> s.struct.field("project")
361 0 pandas
362 1 pandas
363 2 numpy
364 Name: project, dtype: string[pyarrow]
366 Extract by field index.
368 >>> s.struct.field(0)
369 0 1
370 1 2
371 2 1
372 Name: version, dtype: int64[pyarrow]
374 Or an expression
376 >>> import pyarrow.compute as pc
377 >>> s.struct.field(pc.field("project"))
378 0 pandas
379 1 pandas
380 2 numpy
381 Name: project, dtype: string[pyarrow]
383 For nested struct types, you can pass a list of values to index
384 multiple levels:
386 >>> version_type = pa.struct(
387 ... [
388 ... ("major", pa.int64()),
389 ... ("minor", pa.int64()),
390 ... ]
391 ... )
392 >>> s = pd.Series(
393 ... [
394 ... {"version": {"major": 1, "minor": 5}, "project": "pandas"},
395 ... {"version": {"major": 2, "minor": 1}, "project": "pandas"},
396 ... {"version": {"major": 1, "minor": 26}, "project": "numpy"},
397 ... ],
398 ... dtype=pd.ArrowDtype(
399 ... pa.struct([("version", version_type), ("project", pa.string())])
400 ... ),
401 ... )
402 >>> s.struct.field(["version", "minor"])
403 0 5
404 1 1
405 2 26
406 Name: minor, dtype: int64[pyarrow]
407 >>> s.struct.field([0, 0])
408 0 1
409 1 2
410 2 1
411 Name: major, dtype: int64[pyarrow]
412 """
413 from pandas import Series
415 def get_name(
416 level_name_or_index: list[str]
417 | list[bytes]
418 | list[int]
419 | pc.Expression
420 | bytes
421 | str
422 | int,
423 data: pa.ChunkedArray,
424 ):
425 if isinstance(level_name_or_index, int):
426 name = data.type.field(level_name_or_index).name
427 elif isinstance(level_name_or_index, (str, bytes)):
428 name = level_name_or_index
429 elif isinstance(level_name_or_index, pc.Expression):
430 name = str(level_name_or_index)
431 elif is_list_like(level_name_or_index):
432 # For nested input like [2, 1, 2]
433 # iteratively get the struct and field name. The last
434 # one is used for the name of the index.
435 level_name_or_index = list(reversed(level_name_or_index))
436 selected = data
437 while level_name_or_index:
438 # we need the cast, otherwise mypy complains about
439 # getting ints, bytes, or str here, which isn't possible.
440 level_name_or_index = cast(list, level_name_or_index)
441 name_or_index = level_name_or_index.pop()
442 name = get_name(name_or_index, selected)
443 selected = selected.type.field(selected.type.get_field_index(name))
444 name = selected.name
445 else:
446 raise ValueError(
447 "name_or_index must be an int, str, bytes, "
448 "pyarrow.compute.Expression, or list of those"
449 )
450 return name
452 pa_arr = self._data.array._pa_array
453 name = get_name(name_or_index, pa_arr)
454 field_arr = pc.struct_field(pa_arr, name_or_index)
456 return Series(
457 field_arr,
458 dtype=ArrowDtype(field_arr.type),
459 index=self._data.index,
460 name=name,
461 )
463 def explode(self) -> DataFrame:
464 """
465 Extract all child fields of a struct as a DataFrame.
467 Returns
468 -------
469 pandas.DataFrame
470 The data corresponding to all child fields.
472 See Also
473 --------
474 Series.struct.field : Return a single child field as a Series.
476 Examples
477 --------
478 >>> import pyarrow as pa
479 >>> s = pd.Series(
480 ... [
481 ... {"version": 1, "project": "pandas"},
482 ... {"version": 2, "project": "pandas"},
483 ... {"version": 1, "project": "numpy"},
484 ... ],
485 ... dtype=pd.ArrowDtype(
486 ... pa.struct([("version", pa.int64()), ("project", pa.string())])
487 ... ),
488 ... )
490 >>> s.struct.explode()
491 version project
492 0 1 pandas
493 1 2 pandas
494 2 1 numpy
495 """
496 from pandas import concat
498 pa_type = self._pa_array.type
499 return concat(
500 [self.field(i) for i in range(pa_type.num_fields)], axis="columns"
501 )