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
26 from typing import Any
28 from pandas import (
29 DataFrame,
30 Series,
31 )
34class ArrowAccessor(metaclass=ABCMeta):
35 @abstractmethod
36 def __init__(self, data, validation_msg: str) -> None:
37 self._data = data
38 self._validation_msg = validation_msg
39 self._validate(data)
41 @abstractmethod
42 def _is_valid_pyarrow_dtype(self, pyarrow_dtype) -> bool:
43 pass
45 def _validate(self, data) -> None:
46 dtype = data.dtype
47 if not HAS_PYARROW or not isinstance(dtype, ArrowDtype):
48 # Raise AttributeError so that inspect can handle non-struct Series.
49 raise AttributeError(self._validation_msg.format(dtype=dtype))
51 if not self._is_valid_pyarrow_dtype(dtype.pyarrow_dtype):
52 # Raise AttributeError so that inspect can handle invalid Series.
53 raise AttributeError(self._validation_msg.format(dtype=dtype))
55 @property
56 def _pa_array(self):
57 return self._data.array._pa_array
60def _list_element_neg(chunk: pa.Array, key: int) -> pa.Array:
61 if pa.types.is_fixed_size_list(chunk.type):
62 chunk = chunk.cast(pa.list_(chunk.type.value_type))
63 if pc.any(pc.less(pc.list_value_length(chunk), -key)).as_py():
64 raise IndexError(f"list index {key} out of range")
65 indices = pc.add(chunk.offsets[1:], key)
66 if chunk.null_count:
67 indices = pc.if_else(chunk.is_valid(), indices, None)
68 return chunk.values.take(indices)
71class ListAccessor(ArrowAccessor):
72 """
73 Accessor object for list data properties of the Series values.
75 Parameters
76 ----------
77 data : Series
78 Series containing Arrow list data.
79 """
81 def __init__(self, data=None) -> None:
82 super().__init__(
83 data,
84 validation_msg="Can only use the '.list' accessor with "
85 "'list[pyarrow]' dtype, not {dtype}.",
86 )
88 def _is_valid_pyarrow_dtype(self, pyarrow_dtype) -> bool:
89 return (
90 pa.types.is_list(pyarrow_dtype)
91 or pa.types.is_fixed_size_list(pyarrow_dtype)
92 or pa.types.is_large_list(pyarrow_dtype)
93 )
95 def len(self) -> Series:
96 """
97 Return the length of each list in the Series.
99 Computes the number of elements in each list entry. The result is a
100 Series of integers with the same index as the original Series.
102 Returns
103 -------
104 pandas.Series
105 The length of each list.
107 See Also
108 --------
109 str.len : Python built-in function returning the length of an object.
110 Series.size : Returns the length of the Series.
111 StringMethods.len : Compute the length of each element in the Series/Index.
113 Examples
114 --------
115 >>> import pyarrow as pa
116 >>> s = pd.Series(
117 ... [
118 ... [1, 2, 3],
119 ... [3],
120 ... ],
121 ... dtype=pd.ArrowDtype(pa.list_(pa.int64())),
122 ... )
123 >>> s.list.len()
124 0 3
125 1 1
126 dtype: int32[pyarrow]
127 """
128 from pandas import Series
130 value_lengths = pc.list_value_length(self._pa_array)
131 return Series(
132 value_lengths,
133 dtype=ArrowDtype(value_lengths.type),
134 index=self._data.index,
135 name=self._data.name,
136 )
138 def __getitem__(self, key: int | slice) -> Series:
139 """
140 Index or slice lists in the Series.
142 Retrieves elements at the given integer index or slice from each list
143 in the Series, returning a new Series with the selected values.
145 Parameters
146 ----------
147 key : int | slice
148 Index or slice of indices to access from each list.
150 Returns
151 -------
152 pandas.Series
153 The list at requested index.
155 See Also
156 --------
157 ListAccessor.flatten : Flatten list values.
159 Examples
160 --------
161 >>> import pyarrow as pa
162 >>> s = pd.Series(
163 ... [
164 ... [1, 2, 3],
165 ... [3],
166 ... ],
167 ... dtype=pd.ArrowDtype(pa.list_(pa.int64())),
168 ... )
169 >>> s.list[0]
170 0 1
171 1 3
172 dtype: int64[pyarrow]
173 """
174 from pandas import Series
176 if isinstance(key, int):
177 pa_array = self._pa_array
178 chunks = pa_array.chunks
179 if key < 0:
180 element = pa.chunked_array(
181 [_list_element_neg(chunk, key) for chunk in chunks],
182 type=pa_array.type.value_type,
183 )
184 else:
185 element = pc.list_element(pa_array, key)
186 return Series(
187 element,
188 dtype=ArrowDtype(element.type),
189 index=self._data.index,
190 name=self._data.name,
191 )
192 elif isinstance(key, slice):
193 # TODO: Support negative start/stop/step, ideally this would be added
194 # upstream in pyarrow.
195 start, stop, step = key.start, key.stop, key.step
196 if start is None:
197 # TODO: When adding negative step support
198 # this should be set to last element of array
199 # when step is negative.
200 start = 0
201 if step is None:
202 step = 1
203 sliced = pc.list_slice(self._pa_array, start, stop, step)
204 return Series(
205 sliced,
206 dtype=ArrowDtype(sliced.type),
207 index=self._data.index,
208 name=self._data.name,
209 )
210 else:
211 raise ValueError(f"key must be an int or slice, got {type(key).__name__}")
213 def __iter__(self) -> Iterator[Any]:
214 raise TypeError(f"'{type(self).__name__}' object is not iterable")
216 def flatten(self) -> Series:
217 """
218 Flatten list values.
220 Each list element is expanded into separate rows, preserving the
221 original index. The resulting Series may have a longer length than
222 the original if lists contain more than one element.
224 Returns
225 -------
226 pandas.Series
227 The data from all lists in the series flattened.
229 See Also
230 --------
231 ListAccessor.__getitem__ : Index or slice values in the Series.
233 Examples
234 --------
235 >>> import pyarrow as pa
236 >>> s = pd.Series(
237 ... [
238 ... [1, 2, 3],
239 ... [3],
240 ... ],
241 ... dtype=pd.ArrowDtype(pa.list_(pa.int64())),
242 ... )
243 >>> s.list.flatten()
244 0 1
245 0 2
246 0 3
247 1 3
248 dtype: int64[pyarrow]
249 """
250 from pandas import Series
252 counts = pa.compute.list_value_length(self._pa_array)
253 flattened = pa.compute.list_flatten(self._pa_array)
254 index = self._data.index.repeat(counts.fill_null(pa.scalar(0, counts.type)))
255 return Series(
256 flattened,
257 dtype=ArrowDtype(flattened.type),
258 index=index,
259 name=self._data.name,
260 )
263class StructAccessor(ArrowAccessor):
264 """
265 Accessor object for structured data properties of the Series values.
267 Parameters
268 ----------
269 data : Series
270 Series containing Arrow struct data.
271 """
273 def __init__(self, data=None) -> None:
274 super().__init__(
275 data,
276 validation_msg=(
277 "Can only use the '.struct' accessor with 'struct[pyarrow]' "
278 "dtype, not {dtype}."
279 ),
280 )
282 def _is_valid_pyarrow_dtype(self, pyarrow_dtype) -> bool:
283 return pa.types.is_struct(pyarrow_dtype)
285 @property
286 def dtypes(self) -> Series:
287 """
288 Return the dtype object of each child field of the struct.
290 The returned Series is indexed by the field names of the struct and
291 contains the corresponding :class:`ArrowDtype` for each field.
293 Returns
294 -------
295 pandas.Series
296 The data type of each child field.
298 See Also
299 --------
300 Series.dtype: Return the dtype object of the underlying data.
302 Examples
303 --------
304 >>> import pyarrow as pa
305 >>> s = pd.Series(
306 ... [
307 ... {"version": 1, "project": "pandas"},
308 ... {"version": 2, "project": "pandas"},
309 ... {"version": 1, "project": "numpy"},
310 ... ],
311 ... dtype=pd.ArrowDtype(
312 ... pa.struct([("version", pa.int64()), ("project", pa.string())])
313 ... ),
314 ... )
315 >>> s.struct.dtypes
316 version int64[pyarrow]
317 project string[pyarrow]
318 dtype: object
319 """
320 from pandas import (
321 Index,
322 Series,
323 )
325 pa_type = self._data.dtype.pyarrow_dtype
326 types = [ArrowDtype(struct.type) for struct in pa_type]
327 names = [struct.name for struct in pa_type]
328 return Series(types, index=Index(names))
330 def field(
331 self,
332 name_or_index: list[str]
333 | list[bytes]
334 | list[int]
335 | pc.Expression
336 | bytes
337 | str
338 | int,
339 ) -> Series:
340 """
341 Extract a child field of a struct as a Series.
343 This method accesses individual child fields of an Arrow struct
344 column by name, index, expression, or a list of those for nested
345 structs.
347 Parameters
348 ----------
349 name_or_index : str | bytes | int | expression | list
350 Name or index of the child field to extract.
352 For list-like inputs, this will index into a nested
353 struct.
355 Returns
356 -------
357 pandas.Series
358 The data corresponding to the selected child field.
360 See Also
361 --------
362 Series.struct.explode : Return all child fields as a DataFrame.
364 Notes
365 -----
366 The name of the resulting Series will be set using the following
367 rules:
369 - For string, bytes, or integer `name_or_index` (or a list of these, for
370 a nested selection), the Series name is set to the selected
371 field's name.
372 - For a :class:`pyarrow.compute.Expression`, this is set to
373 the string form of the expression.
374 - For list-like `name_or_index`, the name will be set to the
375 name of the final field selected.
377 Examples
378 --------
379 >>> import pyarrow as pa
380 >>> s = pd.Series(
381 ... [
382 ... {"version": 1, "project": "pandas"},
383 ... {"version": 2, "project": "pandas"},
384 ... {"version": 1, "project": "numpy"},
385 ... ],
386 ... dtype=pd.ArrowDtype(
387 ... pa.struct([("version", pa.int64()), ("project", pa.string())])
388 ... ),
389 ... )
391 Extract by field name.
393 >>> s.struct.field("project")
394 0 pandas
395 1 pandas
396 2 numpy
397 Name: project, dtype: string[pyarrow]
399 Extract by field index.
401 >>> s.struct.field(0)
402 0 1
403 1 2
404 2 1
405 Name: version, dtype: int64[pyarrow]
407 Or an expression
409 >>> import pyarrow.compute as pc
410 >>> s.struct.field(pc.field("project"))
411 0 pandas
412 1 pandas
413 2 numpy
414 Name: project, dtype: string[pyarrow]
416 For nested struct types, you can pass a list of values to index
417 multiple levels:
419 >>> version_type = pa.struct(
420 ... [
421 ... ("major", pa.int64()),
422 ... ("minor", pa.int64()),
423 ... ]
424 ... )
425 >>> s = pd.Series(
426 ... [
427 ... {"version": {"major": 1, "minor": 5}, "project": "pandas"},
428 ... {"version": {"major": 2, "minor": 1}, "project": "pandas"},
429 ... {"version": {"major": 1, "minor": 26}, "project": "numpy"},
430 ... ],
431 ... dtype=pd.ArrowDtype(
432 ... pa.struct([("version", version_type), ("project", pa.string())])
433 ... ),
434 ... )
435 >>> s.struct.field(["version", "minor"])
436 0 5
437 1 1
438 2 26
439 Name: minor, dtype: int64[pyarrow]
440 >>> s.struct.field([0, 0])
441 0 1
442 1 2
443 2 1
444 Name: major, dtype: int64[pyarrow]
445 """
446 from pandas import Series
448 def get_name(
449 level_name_or_index: list[str]
450 | list[bytes]
451 | list[int]
452 | pc.Expression
453 | bytes
454 | str
455 | int,
456 data: pa.ChunkedArray,
457 ):
458 if isinstance(level_name_or_index, int):
459 name = data.type.field(level_name_or_index).name
460 elif isinstance(level_name_or_index, (str, bytes)):
461 name = level_name_or_index
462 elif isinstance(level_name_or_index, pc.Expression):
463 name = str(level_name_or_index)
464 elif is_list_like(level_name_or_index):
465 # For nested input like [2, 1, 2]
466 # iteratively get the struct and field name. The last
467 # one is used for the name of the index.
468 remaining = cast(
469 "list[int | str | bytes]", list(reversed(level_name_or_index))
470 )
471 selected = data
472 while remaining:
473 name_or_index = remaining.pop()
474 name = get_name(name_or_index, selected)
475 selected = selected.type.field(selected.type.get_field_index(name))
476 name = selected.name
477 else:
478 raise ValueError(
479 "name_or_index must be an int, str, bytes, "
480 "pyarrow.compute.Expression, or list of those"
481 )
482 return name
484 pa_arr = self._data.array._pa_array
485 name = get_name(name_or_index, pa_arr)
486 field_arr = pc.struct_field(pa_arr, name_or_index)
488 return Series(
489 field_arr,
490 dtype=ArrowDtype(field_arr.type),
491 index=self._data.index,
492 name=name,
493 )
495 def explode(self) -> DataFrame:
496 """
497 Extract all child fields of a struct as a DataFrame.
499 Each child field of the struct becomes a column in the resulting
500 DataFrame, with column names matching the struct field names.
502 Returns
503 -------
504 pandas.DataFrame
505 The data corresponding to all child fields.
507 See Also
508 --------
509 Series.struct.field : Return a single child field as a Series.
511 Examples
512 --------
513 >>> import pyarrow as pa
514 >>> s = pd.Series(
515 ... [
516 ... {"version": 1, "project": "pandas"},
517 ... {"version": 2, "project": "pandas"},
518 ... {"version": 1, "project": "numpy"},
519 ... ],
520 ... dtype=pd.ArrowDtype(
521 ... pa.struct([("version", pa.int64()), ("project", pa.string())])
522 ... ),
523 ... )
525 >>> s.struct.explode()
526 version project
527 0 1 pandas
528 1 2 pandas
529 2 1 numpy
530 """
531 from pandas import concat
533 pa_type = self._pa_array.type
534 return concat(
535 [self.field(i) for i in range(pa_type.num_fields)], axis="columns"
536 )