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
1from __future__ import annotations
3from datetime import (
4 date,
5 datetime,
6)
7import functools
8import operator
9import re
10import textwrap
11from typing import (
12 TYPE_CHECKING,
13 Any,
14 Literal,
15 Self,
16 cast,
17 overload,
18)
19import unicodedata
20import warnings
22import numpy as np
24from pandas._config import is_nan_na
26from pandas._libs import lib
27from pandas._libs.missing import is_pdna_or_none
28from pandas._libs.tslibs import (
29 Timedelta,
30 Timestamp,
31 timezones,
32)
33from pandas.compat import (
34 HAS_PYARROW,
35 PYARROW_MIN_VERSION,
36 pa_version_under21p0,
37)
38from pandas.errors import Pandas4Warning
39from pandas.util._decorators import (
40 doc,
41 set_module,
42)
43from pandas.util._exceptions import find_stack_level
45from pandas.core.dtypes.cast import (
46 can_hold_element,
47 construct_1d_object_array_from_listlike,
48 infer_dtype_from_scalar,
49)
50from pandas.core.dtypes.common import (
51 is_array_like,
52 is_bool_dtype,
53 is_float_dtype,
54 is_integer,
55 is_list_like,
56 is_numeric_dtype,
57 is_scalar,
58 is_string_dtype,
59 pandas_dtype,
60)
61from pandas.core.dtypes.dtypes import DatetimeTZDtype
62from pandas.core.dtypes.missing import isna
64from pandas.core import (
65 algorithms as algos,
66 missing,
67 ops,
68 roperator,
69)
70from pandas.core.algorithms import map_array
71from pandas.core.arraylike import OpsMixin
72from pandas.core.arrays._arrow_string_mixins import ArrowStringArrayMixin
73from pandas.core.arrays._utils import to_numpy_dtype_inference
74from pandas.core.arrays.base import (
75 ExtensionArray,
76 ExtensionArraySupportsAnyAll,
77)
78from pandas.core.arrays.masked import BaseMaskedArray
79from pandas.core.arrays.string_ import StringDtype
80import pandas.core.common as com
81from pandas.core.construction import extract_array
82from pandas.core.indexers import (
83 check_array_indexer,
84 getitem_returns_view,
85 unpack_tuple_and_ellipses,
86 validate_indices,
87)
88from pandas.core.nanops import check_below_min_count
90from pandas.io._util import _arrow_dtype_mapping
91from pandas.tseries.frequencies import to_offset
93if HAS_PYARROW:
94 import pyarrow as pa
95 import pyarrow.compute as pc
97 from pandas.compat.pyarrow import _safe_fill_null
99 from pandas.core.dtypes.dtypes import ArrowDtype
101 ARROW_CMP_FUNCS = {
102 "eq": pc.equal,
103 "ne": pc.not_equal,
104 "lt": pc.less,
105 "gt": pc.greater,
106 "le": pc.less_equal,
107 "ge": pc.greater_equal,
108 }
110 ARROW_LOGICAL_FUNCS = {
111 "and_": pc.and_kleene,
112 "rand_": lambda x, y: pc.and_kleene(y, x),
113 "or_": pc.or_kleene,
114 "ror_": lambda x, y: pc.or_kleene(y, x),
115 "xor": pc.xor,
116 "rxor": lambda x, y: pc.xor(y, x),
117 }
119 ARROW_BIT_WISE_FUNCS = {
120 "and_": pc.bit_wise_and,
121 "rand_": lambda x, y: pc.bit_wise_and(y, x),
122 "or_": pc.bit_wise_or,
123 "ror_": lambda x, y: pc.bit_wise_or(y, x),
124 "xor": pc.bit_wise_xor,
125 "rxor": lambda x, y: pc.bit_wise_xor(y, x),
126 }
128 def cast_for_truediv(
129 arrow_array: pa.ChunkedArray, pa_object: pa.Array | pa.Scalar
130 ) -> tuple[pa.ChunkedArray, pa.Array | pa.Scalar]:
131 # Ensure int / int -> float mirroring Python/Numpy behavior
132 # as pc.divide_checked(int, int) -> int
133 if pa.types.is_integer(arrow_array.type) and pa.types.is_integer(
134 pa_object.type
135 ):
136 # GH: 56645.
137 # https://github.com/apache/arrow/issues/35563
138 return pc.cast(arrow_array, pa.float64(), safe=False), pc.cast(
139 pa_object, pa.float64(), safe=False
140 )
142 return arrow_array, pa_object
144 def floordiv_compat(
145 left: pa.ChunkedArray | pa.Array | pa.Scalar,
146 right: pa.ChunkedArray | pa.Array | pa.Scalar,
147 ) -> pa.ChunkedArray:
148 # TODO: Replace with pyarrow floordiv kernel.
149 # https://github.com/apache/arrow/issues/39386
150 if pa.types.is_integer(left.type) and pa.types.is_integer(right.type):
151 divided = pc.divide_checked(left, right)
152 if pa.types.is_signed_integer(divided.type):
153 # GH 56676
154 has_remainder = pc.not_equal(pc.multiply(divided, right), left)
155 has_one_negative_operand = pc.less(
156 pc.bit_wise_xor(left, right),
157 pa.scalar(0, type=divided.type),
158 )
159 result = pc.if_else(
160 pc.and_(
161 has_remainder,
162 has_one_negative_operand,
163 ),
164 # GH: 55561
165 pc.subtract(divided, pa.scalar(1, type=divided.type)),
166 divided,
167 )
168 else:
169 result = divided
170 result = result.cast(left.type)
171 else:
172 divided = pc.divide(left, right)
173 result = pc.floor(divided)
174 return result
176 ARROW_ARITHMETIC_FUNCS = {
177 "add": pc.add_checked,
178 "radd": lambda x, y: pc.add_checked(y, x),
179 "sub": pc.subtract_checked,
180 "rsub": lambda x, y: pc.subtract_checked(y, x),
181 "mul": pc.multiply_checked,
182 "rmul": lambda x, y: pc.multiply_checked(y, x),
183 "truediv": lambda x, y: pc.divide(*cast_for_truediv(x, y)),
184 "rtruediv": lambda x, y: pc.divide(*cast_for_truediv(y, x)),
185 "floordiv": lambda x, y: floordiv_compat(x, y),
186 "rfloordiv": lambda x, y: floordiv_compat(y, x),
187 "mod": NotImplemented,
188 "rmod": NotImplemented,
189 "divmod": NotImplemented,
190 "rdivmod": NotImplemented,
191 "pow": pc.power_checked,
192 "rpow": lambda x, y: pc.power_checked(y, x),
193 }
195if TYPE_CHECKING:
196 from collections.abc import (
197 Callable,
198 Sequence,
199 )
201 from pandas._libs.missing import NAType
202 from pandas._typing import (
203 ArrayLike,
204 AxisInt,
205 Dtype,
206 FillnaOptions,
207 InterpolateOptions,
208 Iterator,
209 NpDtype,
210 NumpySorter,
211 NumpyValueArrayLike,
212 PositionalIndexer,
213 Scalar,
214 SortKind,
215 TakeIndexer,
216 TimeAmbiguous,
217 TimeNonexistent,
218 npt,
219 )
221 from pandas.core.dtypes.dtypes import ExtensionDtype
223 from pandas import Series
224 from pandas.core.arrays.datetimes import DatetimeArray
225 from pandas.core.arrays.timedeltas import TimedeltaArray
228def to_pyarrow_type(
229 dtype: ArrowDtype | pa.DataType | Dtype | None,
230) -> pa.DataType | None:
231 """
232 Convert dtype to a pyarrow type instance.
233 """
234 if isinstance(dtype, ArrowDtype):
235 return dtype.pyarrow_dtype
236 elif isinstance(dtype, pa.DataType):
237 return dtype
238 elif isinstance(dtype, DatetimeTZDtype):
239 return pa.timestamp(dtype.unit, dtype.tz)
240 elif dtype:
241 try:
242 # Accepts python types too
243 # Doesn't handle all numpy types
244 return pa.from_numpy_dtype(dtype)
245 except pa.ArrowNotImplementedError:
246 pass
247 return None
250@set_module("pandas.arrays")
251class ArrowExtensionArray(
252 OpsMixin,
253 ExtensionArraySupportsAnyAll,
254 ArrowStringArrayMixin,
255):
256 """
257 Pandas ExtensionArray backed by a PyArrow ChunkedArray.
259 .. warning::
261 ArrowExtensionArray is considered experimental. The implementation and
262 parts of the API may change without warning.
264 Parameters
265 ----------
266 values : pyarrow.Array or pyarrow.ChunkedArray
267 The input data to initialize the ArrowExtensionArray.
269 Attributes
270 ----------
271 None
273 Methods
274 -------
275 None
277 Returns
278 -------
279 ArrowExtensionArray
281 See Also
282 --------
283 array : Create a Pandas array with a specified dtype.
284 DataFrame.to_feather : Write a DataFrame to the binary Feather format.
285 read_feather : Load a feather-format object from the file path.
287 Notes
288 -----
289 Most methods are implemented using `pyarrow compute functions. <https://arrow.apache.org/docs/python/api/compute.html>`__
290 Some methods may either raise an exception or raise a ``PerformanceWarning`` if an
291 associated compute function is not available based on the installed version of PyArrow.
293 Please install the latest version of PyArrow to enable the best functionality and avoid
294 potential bugs in prior versions of PyArrow.
296 Examples
297 --------
298 Create an ArrowExtensionArray with :func:`pandas.array`:
300 >>> pd.array([1, 1, None], dtype="int64[pyarrow]")
301 <ArrowExtensionArray>
302 [1, 1, <NA>]
303 Length: 3, dtype: int64[pyarrow]
304 """ # noqa: E501 (http link too long)
306 _pa_array: pa.ChunkedArray
307 _dtype: ArrowDtype
309 def __init__(self, values: pa.Array | pa.ChunkedArray) -> None:
310 if not HAS_PYARROW:
311 msg = (
312 f"pyarrow>={PYARROW_MIN_VERSION} is required for PyArrow "
313 "backed ArrowExtensionArray."
314 )
315 raise ImportError(msg)
316 if isinstance(values, pa.Array):
317 self._pa_array = pa.chunked_array([values])
318 elif isinstance(values, pa.ChunkedArray):
319 self._pa_array = values
320 else:
321 raise ValueError(
322 f"Unsupported type '{type(values)}' for ArrowExtensionArray"
323 )
324 self._dtype = ArrowDtype(self._pa_array.type)
326 @classmethod
327 def _from_sequence(
328 cls, scalars, *, dtype: Dtype | None = None, copy: bool = False
329 ) -> Self:
330 """
331 Construct a new ExtensionArray from a sequence of scalars.
332 """
333 pa_type = to_pyarrow_type(dtype)
334 pa_array = cls._box_pa_array(scalars, pa_type=pa_type, copy=copy)
335 arr = cls(pa_array)
336 return arr
338 @classmethod
339 def _from_sequence_of_strings(
340 cls, strings, *, dtype: ExtensionDtype, copy: bool = False
341 ) -> Self:
342 """
343 Construct a new ExtensionArray from a sequence of strings.
344 """
345 mask = isna(strings)
347 if isinstance(strings, cls):
348 strings = strings._pa_array
350 pa_type = to_pyarrow_type(dtype)
351 if (
352 pa_type is None
353 or pa.types.is_binary(pa_type)
354 or pa.types.is_string(pa_type)
355 or pa.types.is_large_string(pa_type)
356 ):
357 # pa_type is None: Let pa.array infer
358 # pa_type is string/binary: scalars already correct type
359 scalars = strings
360 elif pa.types.is_timestamp(pa_type):
361 from pandas.core.tools.datetimes import to_datetime
363 scalars = to_datetime(strings, errors="raise")
364 elif pa.types.is_date(pa_type):
365 from pandas.core.tools.datetimes import to_datetime
367 scalars = to_datetime(strings, errors="raise").date
368 scalars = pa.array(scalars, type=pa_type, mask=mask)
369 elif pa.types.is_duration(pa_type):
370 from pandas.core.tools.timedeltas import to_timedelta
372 scalars = to_timedelta(strings, errors="raise")
374 if pa_type.unit != "ns":
375 # GH51175: test_from_sequence_of_strings_pa_array
376 # attempt to parse as int64 reflecting pyarrow's
377 # duration to string casting behavior
378 mask = isna(scalars)
379 if not isinstance(strings, (pa.Array, pa.ChunkedArray)):
380 strings = pa.array(strings, type=pa.string(), mask=mask)
381 strings = pc.if_else(mask, None, strings)
382 try:
383 scalars = strings.cast(pa.int64())
384 except pa.ArrowInvalid:
385 pass
386 elif pa.types.is_time(pa_type):
387 from pandas.core.tools.times import to_time
389 # "coerce" to allow "null times" (None) to not raise
390 scalars = to_time(strings, errors="coerce")
391 elif pa.types.is_boolean(pa_type):
392 # pyarrow string->bool casting is case-insensitive:
393 # "true" or "1" -> True
394 # "false" or "0" -> False
395 # Note: BooleanArray was previously used to parse these strings
396 # and allows "1.0" and "0.0". Pyarrow casting does not support
397 # this, but we allow it here.
398 if isinstance(strings, (pa.Array, pa.ChunkedArray)):
399 scalars = strings
400 else:
401 scalars = pa.array(strings, type=pa.string(), mask=mask)
402 scalars = pc.if_else(pc.equal(scalars, "1.0"), "1", scalars)
403 scalars = pc.if_else(pc.equal(scalars, "0.0"), "0", scalars)
404 scalars = scalars.cast(pa.bool_())
405 elif (
406 pa.types.is_integer(pa_type)
407 or pa.types.is_floating(pa_type)
408 or pa.types.is_decimal(pa_type)
409 ):
410 from pandas.core.tools.numeric import to_numeric
412 scalars = to_numeric(strings, errors="raise")
413 if isinstance(strings, (pa.Array, pa.ChunkedArray)):
414 scalars = strings.cast(pa_type)
415 elif mask is not None:
416 scalars = pa.array(scalars, mask=mask, type=pa_type)
418 else:
419 raise NotImplementedError(
420 f"Converting strings to {pa_type} is not implemented."
421 )
422 return cls._from_sequence(scalars, dtype=pa_type, copy=copy)
424 def _from_pyarrow_array(self, pa_array):
425 """
426 Construct from the pyarrow array result of an operation, for
427 compatibility with ArrowStringArray.
428 """
429 return type(self)(pa_array)
431 def _cast_pointwise_result(self, values) -> ArrayLike:
432 if len(values) == 0:
433 # Retain our dtype
434 return self[:0].copy()
436 try:
437 if self.dtype.kind in "iufc" and not is_nan_na():
438 values = np.asarray(values, dtype=object)
439 mask = is_pdna_or_none(values)
440 arr = pa.array(values, mask=mask)
441 else:
442 arr = pa.array(values, from_pandas=True)
443 except (ValueError, TypeError):
444 # e.g. test_by_column_values_with_same_starting_value with nested
445 # values, one entry of which is an ArrowStringArray
446 # or test_agg_lambda_complex128_dtype_conversion for complex values
447 values = np.asarray(values, dtype=object)
448 return lib.maybe_convert_objects(values, convert_non_numeric=True)
450 if pa.types.is_null(arr.type):
451 if lib.infer_dtype(values) == "decimal":
452 # GH#62522; the specific decimal precision here is arbitrary
453 arr = arr.cast(pa.decimal128(1))
454 if pa.types.is_duration(arr.type):
455 # workaround for https://github.com/apache/arrow/issues/40620
456 result = ArrowExtensionArray._from_sequence(values)
457 if pa.types.is_duration(self._pa_array.type):
458 result = result.astype(self.dtype) # type: ignore[assignment]
459 elif pa.types.is_timestamp(self._pa_array.type):
460 # Try to retain original unit
461 new_dtype = ArrowDtype(pa.duration(self._pa_array.type.unit))
462 try:
463 result = result.astype(new_dtype) # type: ignore[assignment]
464 except ValueError:
465 pass
466 elif pa.types.is_date64(self._pa_array.type):
467 # Try to match unit we get on non-pointwise op
468 dtype = ArrowDtype(pa.duration("ms"))
469 result = result.astype(dtype) # type: ignore[assignment]
470 elif pa.types.is_date(self._pa_array.type):
471 # Try to match unit we get on non-pointwise op
472 dtype = ArrowDtype(pa.duration("s"))
473 result = result.astype(dtype) # type: ignore[assignment]
474 return result
476 elif pa.types.is_date(arr.type) and pa.types.is_date(self._pa_array.type):
477 arr = arr.cast(self._pa_array.type)
478 elif pa.types.is_time(arr.type) and pa.types.is_time(self._pa_array.type):
479 arr = arr.cast(self._pa_array.type)
480 elif pa.types.is_decimal(arr.type) and pa.types.is_decimal(self._pa_array.type):
481 arr = arr.cast(self._pa_array.type)
482 elif pa.types.is_integer(arr.type) and pa.types.is_integer(self._pa_array.type):
483 try:
484 arr = arr.cast(self._pa_array.type)
485 except pa.lib.ArrowInvalid:
486 # e.g. test_combine_add if we can't cast
487 pass
488 elif pa.types.is_floating(arr.type) and pa.types.is_floating(
489 self._pa_array.type
490 ):
491 try:
492 arr = arr.cast(self._pa_array.type)
493 except pa.lib.ArrowInvalid:
494 # e.g. test_combine_add if we can't cast
495 pass
497 if isinstance(self.dtype, StringDtype):
498 if pa.types.is_string(arr.type) or pa.types.is_large_string(arr.type):
499 # ArrowStringArray preserves dtype.na_value
500 return self._from_pyarrow_array(arr)
501 if self.dtype.na_value is np.nan:
502 # ArrowEA has different semantics, so we return numpy-based
503 # result instead
504 values = np.asarray(values, dtype=object)
505 return lib.maybe_convert_objects(values, convert_non_numeric=True)
506 return ArrowExtensionArray(arr)
507 return self._from_pyarrow_array(arr)
509 @classmethod
510 def _box_pa(
511 cls, value, pa_type: pa.DataType | None = None
512 ) -> pa.Array | pa.ChunkedArray | pa.Scalar:
513 """
514 Box value into a pyarrow Array, ChunkedArray or Scalar.
516 Parameters
517 ----------
518 value : any
519 pa_type : pa.DataType | None
521 Returns
522 -------
523 pa.Array or pa.ChunkedArray or pa.Scalar
524 """
525 if isinstance(value, pa.Scalar) or not is_list_like(value):
526 return cls._box_pa_scalar(value, pa_type)
527 return cls._box_pa_array(value, pa_type)
529 @classmethod
530 def _box_pa_scalar(cls, value, pa_type: pa.DataType | None = None) -> pa.Scalar:
531 """
532 Box value into a pyarrow Scalar.
534 Parameters
535 ----------
536 value : any
537 pa_type : pa.DataType | None
539 Returns
540 -------
541 pa.Scalar
542 """
543 if isinstance(value, pa.Scalar):
544 pa_scalar = value
545 elif isna(value) and not (lib.is_float(value) and not is_nan_na()):
546 pa_scalar = pa.scalar(None, type=pa_type)
547 else:
548 # Workaround https://github.com/apache/arrow/issues/37291
549 if isinstance(value, Timedelta):
550 if pa_type is None:
551 pa_type = pa.duration(value.unit)
552 elif value.unit != pa_type.unit:
553 value = value.as_unit(pa_type.unit)
554 value = value._value
555 elif isinstance(value, Timestamp):
556 if pa_type is None:
557 pa_type = pa.timestamp(value.unit, tz=value.tz)
558 elif value.unit != pa_type.unit:
559 value = value.as_unit(pa_type.unit)
560 value = value._value
562 pa_scalar = pa.scalar(value, type=pa_type)
564 if pa_type is not None and pa_scalar.type != pa_type:
565 pa_scalar = pa_scalar.cast(pa_type)
567 return pa_scalar
569 @classmethod
570 def _box_pa_array(
571 cls, value, pa_type: pa.DataType | None = None, copy: bool = False
572 ) -> pa.Array | pa.ChunkedArray:
573 """
574 Box value into a pyarrow Array or ChunkedArray.
576 Parameters
577 ----------
578 value : Sequence
579 pa_type : pa.DataType | None
581 Returns
582 -------
583 pa.Array or pa.ChunkedArray
584 """
585 value = extract_array(value, extract_numpy=True)
586 if isinstance(value, cls):
587 pa_array = value._pa_array
588 elif isinstance(value, (pa.Array, pa.ChunkedArray)):
589 pa_array = value
590 elif isinstance(value, BaseMaskedArray):
591 # GH 52625
592 if copy:
593 value = value.copy()
594 pa_array = value.__arrow_array__()
596 elif hasattr(value, "__arrow_array__"):
597 # e.g. StringArray
598 if copy:
599 value = value.copy()
600 pa_array = value.__arrow_array__()
602 else:
603 if (
604 isinstance(value, np.ndarray)
605 and pa_type is not None
606 and (
607 pa.types.is_large_binary(pa_type)
608 or pa.types.is_large_string(pa_type)
609 )
610 ):
611 # See https://github.com/apache/arrow/issues/35289
612 value = np.asarray(value, dtype=object)
613 elif copy and is_array_like(value):
614 # pa array should not get updated when numpy array is updated
615 value = value.copy()
617 if (
618 pa_type is not None
619 and pa.types.is_duration(pa_type)
620 and (not isinstance(value, np.ndarray) or value.dtype.kind not in "mi")
621 ):
622 # Workaround https://github.com/apache/arrow/issues/37291
623 from pandas.core.tools.timedeltas import to_timedelta
625 value = to_timedelta(value, unit=pa_type.unit).as_unit(pa_type.unit)
626 value = value.to_numpy()
628 if pa_type is not None and pa.types.is_timestamp(pa_type):
629 # Use DatetimeArray to exclude Decimal(NaN) (GH#61774) and
630 # ensure constructor treats tznaive the same as non-pyarrow
631 # dtypes (GH#61775)
632 from pandas.core.arrays.datetimes import (
633 DatetimeArray,
634 tz_to_dtype,
635 )
637 pass_dtype = tz_to_dtype(tz=pa_type.tz, unit=pa_type.unit)
638 value = extract_array(value, extract_numpy=True)
639 if isinstance(value, DatetimeArray):
640 dta = value
641 else:
642 dta = DatetimeArray._from_sequence(
643 value, copy=copy, dtype=pass_dtype
644 )
645 dta_mask = dta.isna()
646 value_i8 = cast("npt.NDArray", dta.view("i8"))
647 if not value_i8.flags["WRITEABLE"]:
648 # e.g. test_setitem_frame_2d_values
649 value_i8 = value_i8.copy()
650 dta = DatetimeArray._from_sequence(value_i8, dtype=dta.dtype)
651 value_i8[dta_mask] = 0 # GH#61776 avoid __sub__ overflow
652 pa_array = pa.array(dta._ndarray, type=pa_type, mask=dta_mask)
653 return pa_array
655 mask = None
656 if is_nan_na():
657 try:
658 arr_value = np.asarray(value)
659 if arr_value.ndim > 1:
660 # e.g. test_fixed_size_list we have list data. ndim > 1
661 # means there were no scalar (NA) entries.
662 mask = np.zeros(len(value), dtype=np.bool_)
663 else:
664 mask = isna(arr_value)
665 except ValueError:
666 # Ragged data that numpy raises on
667 arr_value = construct_1d_object_array_from_listlike(value)
668 mask = isna(arr_value)
669 elif (
670 getattr(value, "dtype", None) is None or value.dtype.kind not in "iumMf"
671 ):
672 arr_value = np.asarray(value, dtype=object)
673 # similar to isna(value) but exclude NaN, NaT, nat-like, nan-like
674 mask = is_pdna_or_none(arr_value)
676 try:
677 pa_array = pa.array(value, type=pa_type, mask=mask)
678 except (pa.ArrowInvalid, pa.ArrowTypeError):
679 # GH50430: let pyarrow infer type, then cast
680 pa_array = pa.array(value, mask=mask)
682 if pa_type is None and pa.types.is_duration(pa_array.type):
683 # Workaround https://github.com/apache/arrow/issues/37291
684 from pandas.core.tools.timedeltas import to_timedelta
686 value = to_timedelta(value)
687 value = value.to_numpy()
688 pa_array = pa.array(value, type=pa_type)
690 if pa.types.is_duration(pa_array.type) and pa_array.null_count > 0:
691 # GH52843: upstream bug for duration types when originally
692 # constructed with data containing numpy NaT.
693 # https://github.com/apache/arrow/issues/35088
694 arr = cls(pa_array)
695 arr = arr.fillna(arr.dtype.na_value)
696 pa_array = arr._pa_array
698 if pa_type is not None and pa_array.type != pa_type:
699 if pa.types.is_dictionary(pa_type):
700 pa_array = pa_array.dictionary_encode()
701 if pa_array.type != pa_type:
702 pa_array = pa_array.cast(pa_type)
703 else:
704 try:
705 pa_array = pa_array.cast(pa_type)
706 except (pa.ArrowNotImplementedError, pa.ArrowTypeError):
707 if pa.types.is_string(pa_array.type) or pa.types.is_large_string(
708 pa_array.type
709 ):
710 # TODO: Move logic in _from_sequence_of_strings into
711 # _box_pa_array
712 dtype = ArrowDtype(pa_type)
713 return cls._from_sequence_of_strings(
714 value, dtype=dtype
715 )._pa_array
716 else:
717 raise
719 return pa_array
721 def __getitem__(self, item: PositionalIndexer):
722 """Select a subset of self.
724 Parameters
725 ----------
726 item : int, slice, or ndarray
727 * int: The position in 'self' to get.
728 * slice: A slice object, where 'start', 'stop', and 'step' are
729 integers or None
730 * ndarray: A 1-d boolean NumPy ndarray the same length as 'self'
732 Returns
733 -------
734 item : scalar or ExtensionArray
736 Notes
737 -----
738 For scalar ``item``, return a scalar value suitable for the array's
739 type. This should be an instance of ``self.dtype.type``.
740 For slice ``key``, return an instance of ``ExtensionArray``, even
741 if the slice is length 0 or 1.
742 For a boolean mask, return an instance of ``ExtensionArray``, filtered
743 to the values where ``item`` is True.
744 """
745 item = check_array_indexer(self, item)
747 if isinstance(item, np.ndarray):
748 if not len(item):
749 # Removable once we migrate StringDtype[pyarrow] to ArrowDtype[string]
750 if (
751 isinstance(self._dtype, StringDtype)
752 and self._dtype.storage == "pyarrow"
753 ):
754 # TODO(infer_string) should this be large_string?
755 pa_dtype = pa.string()
756 else:
757 pa_dtype = self._dtype.pyarrow_dtype
758 result = pa.chunked_array([], type=pa_dtype)
759 return self._from_pyarrow_array(result)
761 elif item.dtype.kind in "iu":
762 return self.take(item)
763 elif item.dtype.kind == "b":
764 return self._from_pyarrow_array(self._pa_array.filter(item))
765 else:
766 raise IndexError(
767 "Only integers, slices and integer or "
768 "boolean arrays are valid indices."
769 )
770 elif isinstance(item, tuple):
771 item = unpack_tuple_and_ellipses(item)
773 if item is Ellipsis:
774 # TODO: should be handled by pyarrow?
775 item = slice(None)
777 if is_scalar(item) and not is_integer(item):
778 # e.g. "foo" or 2.5
779 # exception message copied from numpy
780 raise IndexError(
781 r"only integers, slices (`:`), ellipsis (`...`), numpy.newaxis "
782 r"(`None`) and integer or boolean arrays are valid indices"
783 )
784 # We are not an array indexer, so maybe e.g. a slice or integer
785 # indexer. We dispatch to pyarrow.
786 if isinstance(item, slice):
787 # Arrow bug https://github.com/apache/arrow/issues/38768
788 if item.start == item.stop:
789 pass
790 elif (
791 item.stop is not None
792 and item.stop < -len(self)
793 and item.step is not None
794 and item.step < 0
795 ):
796 item = slice(item.start, None, item.step)
798 value = self._pa_array[item]
799 if isinstance(value, pa.ChunkedArray):
800 result = self._from_pyarrow_array(value)
801 if getitem_returns_view(self, item):
802 result._readonly = self._readonly
803 return result
804 else:
805 pa_type = self._pa_array.type
806 scalar = value.as_py()
807 if scalar is None:
808 return self._dtype.na_value
809 elif pa.types.is_timestamp(pa_type) and pa_type.unit != "ns":
810 # GH 53326
811 return Timestamp(scalar).as_unit(pa_type.unit)
812 elif pa.types.is_duration(pa_type) and pa_type.unit != "ns":
813 # GH 53326
814 return Timedelta(scalar).as_unit(pa_type.unit)
815 else:
816 return scalar
818 def __iter__(self) -> Iterator[Any]:
819 """
820 Iterate over elements of the array.
821 """
822 na_value = self._dtype.na_value
823 # GH 53326
824 pa_type = self._pa_array.type
825 box_timestamp = pa.types.is_timestamp(pa_type) and pa_type.unit != "ns"
826 box_timedelta = pa.types.is_duration(pa_type) and pa_type.unit != "ns"
827 for value in self._pa_array:
828 val = value.as_py()
829 if val is None:
830 yield na_value
831 elif box_timestamp:
832 yield Timestamp(val).as_unit(pa_type.unit)
833 elif box_timedelta:
834 yield Timedelta(val).as_unit(pa_type.unit)
835 else:
836 yield val
838 def __arrow_array__(self, type=None):
839 """Convert myself to a pyarrow ChunkedArray."""
840 return self._pa_array
842 def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
843 # Need to wrap np.array results GH#62800
844 result = super().__array_ufunc__(ufunc, method, *inputs, **kwargs)
845 if type(self) is ArrowExtensionArray:
846 # Exclude ArrowStringArray
847 return type(self)._from_sequence(result)
848 return result
850 def __array__(
851 self, dtype: NpDtype | None = None, copy: bool | None = None
852 ) -> np.ndarray:
853 """Correctly construct numpy arrays when passed to `np.asarray()`."""
854 if copy is False:
855 # TODO: By using `zero_copy_only` it may be possible to implement this
856 raise ValueError(
857 "Unable to avoid copy while creating an array as requested."
858 )
859 elif copy is None:
860 # `to_numpy(copy=False)` has the meaning of NumPy `copy=None`.
861 copy = False
863 return self.to_numpy(dtype=dtype, copy=copy)
865 def __invert__(self) -> Self:
866 # This is a bit wise op for integer types
867 if pa.types.is_integer(self._pa_array.type):
868 return self._from_pyarrow_array(pc.bit_wise_not(self._pa_array))
869 elif pa.types.is_string(self._pa_array.type) or pa.types.is_large_string(
870 self._pa_array.type
871 ):
872 # Raise TypeError instead of pa.ArrowNotImplementedError
873 raise TypeError("__invert__ is not supported for string dtypes")
874 else:
875 return self._from_pyarrow_array(pc.invert(self._pa_array))
877 def __neg__(self) -> Self:
878 try:
879 return self._from_pyarrow_array(pc.negate_checked(self._pa_array))
880 except pa.ArrowNotImplementedError as err:
881 raise TypeError(
882 f"unary '-' not supported for dtype '{self.dtype}'"
883 ) from err
885 def __pos__(self) -> Self:
886 return self._from_pyarrow_array(self._pa_array)
888 def __abs__(self) -> Self:
889 return self._from_pyarrow_array(pc.abs_checked(self._pa_array))
891 # GH 42600: __getstate__/__setstate__ not necessary once
892 # https://issues.apache.org/jira/browse/ARROW-10739 is addressed
893 def __getstate__(self):
894 state = self.__dict__.copy()
895 state["_pa_array"] = self._pa_array.combine_chunks()
896 return state
898 def __setstate__(self, state) -> None:
899 if "_data" in state:
900 data = state.pop("_data")
901 else:
902 data = state["_pa_array"]
903 state["_pa_array"] = pa.chunked_array(data)
904 self.__dict__.update(state)
906 def _cmp_method(self, other, op) -> ArrowExtensionArray:
907 pc_func = ARROW_CMP_FUNCS[op.__name__]
908 ltype = self._pa_array.type
910 if isinstance(other, (ExtensionArray, np.ndarray, list, range)):
911 try:
912 boxed = self._box_pa(other)
913 except pa.lib.ArrowInvalid:
914 # e.g. GH#60228 [1, "b"] we have to operate pointwise
915 res_values = [op(x, y) for x, y in zip(self, other, strict=True)]
916 result = pa.array(res_values, type=pa.bool_(), from_pandas=True)
917 else:
918 rtype = boxed.type
919 if (
920 (pa.types.is_timestamp(ltype) and pa.types.is_date(rtype))
921 or (pa.types.is_timestamp(rtype) and pa.types.is_date(ltype))
922 or isinstance(other, range)
923 ):
924 # GH#62157 match non-pyarrow behavior
925 result = ops.invalid_comparison(self, other, op)
926 result = pa.array(result, type=pa.bool_())
927 else:
928 try:
929 result = pc_func(self._pa_array, boxed)
930 except pa.ArrowNotImplementedError:
931 result = ops.invalid_comparison(self, other, op)
932 result = pa.array(result, type=pa.bool_())
934 elif is_scalar(other):
935 if (isinstance(other, datetime) and pa.types.is_date(ltype)) or (
936 type(other) is date and pa.types.is_timestamp(ltype)
937 ):
938 # GH#62157 match non-pyarrow behavior
939 result = ops.invalid_comparison(self, other, op)
940 result = pa.array(result, type=pa.bool_())
941 else:
942 try:
943 result = pc_func(self._pa_array, self._box_pa(other))
944 except (pa.lib.ArrowNotImplementedError, pa.lib.ArrowInvalid):
945 mask = isna(self) | isna(other)
946 valid = ~mask
947 result = np.zeros(len(self), dtype="bool")
948 np_array = np.array(self)
949 try:
950 result[valid] = op(np_array[valid], other)
951 except TypeError:
952 result = ops.invalid_comparison(self, other, op)
953 result = pa.array(result, type=pa.bool_())
954 result = pc.if_else(valid, result, None)
955 else:
956 raise NotImplementedError(
957 f"{op.__name__} not implemented for {type(other)}"
958 )
959 return ArrowExtensionArray(result)
961 def _op_method_error_message(self, other, op) -> str:
962 if hasattr(other, "dtype"):
963 other_type = f"dtype '{other.dtype}'"
964 else:
965 other_type = f"object of type {type(other)}"
966 return (
967 f"operation '{op.__name__}' not supported for "
968 f"dtype '{self.dtype}' with {other_type}"
969 )
971 def _evaluate_op_method(self, other, op, arrow_funcs) -> Self:
972 pa_type = self._pa_array.type
973 other_original = other
974 other = self._box_pa(other)
976 if (
977 pa.types.is_string(pa_type)
978 or pa.types.is_large_string(pa_type)
979 or pa.types.is_binary(pa_type)
980 ):
981 if op in [operator.add, roperator.radd]:
982 # binary_join_element_wise does not support mixed types, but we
983 # want to allow addition between string and large_string types
984 self_array = self._pa_array
985 if pa.types.is_string(pa_type) and pa.types.is_large_string(other.type):
986 self_array = self._pa_array.cast(pa.large_string())
987 elif pa.types.is_large_string(pa_type) and pa.types.is_string(
988 other.type
989 ):
990 other = other.cast(pa.large_string())
992 sep = pa.scalar("", type=self_array.type)
993 if isinstance(other, pa.Scalar) and pc.is_null(other).as_py():
994 other = other.cast(self_array.type)
995 try:
996 if op is operator.add:
997 result = pc.binary_join_element_wise(self_array, other, sep)
998 elif op is roperator.radd:
999 result = pc.binary_join_element_wise(other, self_array, sep)
1000 except pa.ArrowNotImplementedError as err:
1001 raise TypeError(
1002 self._op_method_error_message(other_original, op)
1003 ) from err
1004 return self._from_pyarrow_array(result)
1005 elif op in [operator.mul, roperator.rmul]:
1006 binary = self._pa_array
1007 integral = other
1008 if not pa.types.is_integer(integral.type):
1009 raise TypeError("Can only string multiply by an integer.")
1010 pa_integral = pc.if_else(pc.less(integral, 0), 0, integral)
1011 result = pc.binary_repeat(binary, pa_integral)
1012 return self._from_pyarrow_array(result)
1013 elif (
1014 pa.types.is_string(other.type)
1015 or pa.types.is_binary(other.type)
1016 or pa.types.is_large_string(other.type)
1017 ) and op in [operator.mul, roperator.rmul]:
1018 binary = other
1019 integral = self._pa_array
1020 if not pa.types.is_integer(integral.type):
1021 raise TypeError("Can only string multiply by an integer.")
1022 pa_integral = pc.if_else(pc.less(integral, 0), 0, integral)
1023 result = pc.binary_repeat(binary, pa_integral)
1024 return self._from_pyarrow_array(result)
1025 if (
1026 isinstance(other, pa.Scalar)
1027 and pc.is_null(other).as_py()
1028 and op.__name__ in ARROW_LOGICAL_FUNCS
1029 ):
1030 # pyarrow kleene ops require null to be typed
1031 other = other.cast(pa_type)
1033 pc_func = arrow_funcs[op.__name__]
1034 if pc_func is NotImplemented:
1035 if pa.types.is_string(pa_type) or pa.types.is_large_string(pa_type):
1036 raise TypeError(self._op_method_error_message(other_original, op))
1037 raise NotImplementedError(f"{op.__name__} not implemented.")
1039 try:
1040 result = pc_func(self._pa_array, other)
1041 except pa.ArrowNotImplementedError as err:
1042 raise TypeError(self._op_method_error_message(other_original, op)) from err
1043 return self._from_pyarrow_array(result)
1045 def _logical_method(self, other, op) -> Self:
1046 # For integer types `^`, `|`, `&` are bitwise operators and return
1047 # integer types. Otherwise these are boolean ops.
1048 if pa.types.is_integer(self._pa_array.type):
1049 return self._evaluate_op_method(other, op, ARROW_BIT_WISE_FUNCS)
1050 elif (
1051 (
1052 pa.types.is_string(self._pa_array.type)
1053 or pa.types.is_large_string(self._pa_array.type)
1054 )
1055 and op in (roperator.ror_, roperator.rand_, roperator.rxor)
1056 and isinstance(other, np.ndarray)
1057 and other.dtype == bool
1058 ):
1059 # GH#60234 backward compatibility for the move to StringDtype in 3.0
1060 op_name = op.__name__[1:].strip("_")
1061 warnings.warn(
1062 f"'{op_name}' operations between boolean dtype and {self.dtype} are "
1063 "deprecated and will raise in a future version. Explicitly "
1064 "cast the strings to a boolean dtype before operating instead.",
1065 Pandas4Warning,
1066 stacklevel=find_stack_level(),
1067 )
1068 return op(other, self.astype(bool))
1069 else:
1070 return self._evaluate_op_method(other, op, ARROW_LOGICAL_FUNCS)
1072 def _str_arith_method_object_fallback(
1073 self, other, op
1074 ) -> Self | npt.NDArray[np.object_]:
1075 mask = isna(self) | isna(other)
1076 valid = ~mask
1078 if is_list_like(other):
1079 if len(other) != len(self):
1080 raise ValueError(
1081 f"Lengths of operands do not match: {len(self)} != {len(other)}"
1082 )
1083 if not is_array_like(other):
1084 other = np.asarray(other)
1085 other = other[valid]
1087 result = np.empty(len(self), dtype=object)
1088 result[mask] = self.dtype.na_value
1089 result[valid] = op(np.asarray(self, dtype=object)[valid], other)
1091 if not lib.is_string_array(result, skipna=True):
1092 return result
1093 return type(self)._from_sequence(result, dtype=self.dtype)
1095 def _arith_method(self, other, op) -> Self | npt.NDArray[np.object_]:
1096 result: Self | npt.NDArray[np.object_]
1097 if pa.types.is_string(self._pa_array.type) or pa.types.is_large_string(
1098 self._pa_array.type
1099 ):
1100 try:
1101 result = self._evaluate_op_method(other, op, ARROW_ARITHMETIC_FUNCS)
1102 except (pa.ArrowInvalid, pa.ArrowTypeError):
1103 result = self._str_arith_method_object_fallback(other, op)
1104 else:
1105 result = self._evaluate_op_method(other, op, ARROW_ARITHMETIC_FUNCS)
1106 if isinstance(result, np.ndarray):
1107 return result
1108 if is_nan_na() and result.dtype.kind == "f":
1109 parr = result._pa_array
1110 mask = pc.is_nan(parr).fill_null(False).to_numpy()
1111 arr = pc.replace_with_mask(parr, mask, pa.scalar(None, type=parr.type))
1112 result = type(self)(arr)
1113 return result
1115 def equals(self, other) -> bool:
1116 if not isinstance(other, ArrowExtensionArray):
1117 return False
1118 # I'm told that pyarrow makes __eq__ behave like pandas' equals;
1119 # TODO: is this documented somewhere?
1120 return self._pa_array == other._pa_array
1122 @property
1123 def dtype(self) -> ArrowDtype:
1124 """
1125 An instance of 'ExtensionDtype'.
1126 """
1127 return self._dtype
1129 @property
1130 def nbytes(self) -> int:
1131 """
1132 The number of bytes needed to store this object in memory.
1133 """
1134 return self._pa_array.nbytes
1136 def __len__(self) -> int:
1137 """
1138 Length of this array.
1140 Returns
1141 -------
1142 length : int
1143 """
1144 return len(self._pa_array)
1146 def __contains__(self, key) -> bool:
1147 # https://github.com/pandas-dev/pandas/pull/51307#issuecomment-1426372604
1148 if isna(key) and key is not self.dtype.na_value:
1149 if lib.is_float(key) and is_nan_na():
1150 return self.dtype.na_value in self
1151 elif self.dtype.kind == "f" and lib.is_float(key):
1152 # Check specifically for NaN
1153 return pc.any(pc.is_nan(self._pa_array)).as_py()
1155 # e.g. date or timestamp types we do not allow None here to match pd.NA
1156 return False
1157 # TODO: maybe complex? object?
1159 return bool(super().__contains__(key))
1161 @property
1162 def _hasna(self) -> bool:
1163 return self._pa_array.null_count > 0
1165 def isna(self) -> npt.NDArray[np.bool_]:
1166 """
1167 Boolean NumPy array indicating if each value is missing.
1169 This should return a 1-D array the same length as 'self'.
1170 """
1171 # GH51630: fast paths
1172 null_count = self._pa_array.null_count
1173 if null_count == 0:
1174 return np.zeros(len(self), dtype=np.bool_)
1175 elif null_count == len(self):
1176 return np.ones(len(self), dtype=np.bool_)
1178 return self._pa_array.is_null().to_numpy()
1180 @overload
1181 def any(self, *, skipna: Literal[True] = ..., **kwargs) -> bool: ...
1183 @overload
1184 def any(self, *, skipna: bool, **kwargs) -> bool | NAType: ...
1186 def any(self, *, skipna: bool = True, **kwargs) -> bool | NAType:
1187 """
1188 Return whether any element is truthy.
1190 Returns False unless there is at least one element that is truthy.
1191 By default, NAs are skipped. If ``skipna=False`` is specified and
1192 missing values are present, similar :ref:`Kleene logic <boolean.kleene>`
1193 is used as for logical operations.
1195 Parameters
1196 ----------
1197 skipna : bool, default True
1198 Exclude NA values. If the entire array is NA and `skipna` is
1199 True, then the result will be False, as for an empty array.
1200 If `skipna` is False, the result will still be True if there is
1201 at least one element that is truthy, otherwise NA will be returned
1202 if there are NA's present.
1204 Returns
1205 -------
1206 bool or :attr:`pandas.NA`
1208 See Also
1209 --------
1210 ArrowExtensionArray.all : Return whether all elements are truthy.
1212 Examples
1213 --------
1214 The result indicates whether any element is truthy (and by default
1215 skips NAs):
1217 >>> pd.array([True, False, True], dtype="boolean[pyarrow]").any()
1218 True
1219 >>> pd.array([True, False, pd.NA], dtype="boolean[pyarrow]").any()
1220 True
1221 >>> pd.array([False, False, pd.NA], dtype="boolean[pyarrow]").any()
1222 False
1223 >>> pd.array([], dtype="boolean[pyarrow]").any()
1224 False
1225 >>> pd.array([pd.NA], dtype="boolean[pyarrow]").any()
1226 False
1227 >>> pd.array([pd.NA], dtype="float64[pyarrow]").any()
1228 False
1230 With ``skipna=False``, the result can be NA if this is logically
1231 required (whether ``pd.NA`` is True or False influences the result):
1233 >>> pd.array([True, False, pd.NA], dtype="boolean[pyarrow]").any(skipna=False)
1234 True
1235 >>> pd.array([1, 0, pd.NA], dtype="boolean[pyarrow]").any(skipna=False)
1236 True
1237 >>> pd.array([False, False, pd.NA], dtype="boolean[pyarrow]").any(skipna=False)
1238 <NA>
1239 >>> pd.array([0, 0, pd.NA], dtype="boolean[pyarrow]").any(skipna=False)
1240 <NA>
1241 """
1242 return self._reduce("any", skipna=skipna, **kwargs)
1244 @overload
1245 def all(self, *, skipna: Literal[True] = ..., **kwargs) -> bool: ...
1247 @overload
1248 def all(self, *, skipna: bool, **kwargs) -> bool | NAType: ...
1250 def all(self, *, skipna: bool = True, **kwargs) -> bool | NAType:
1251 """
1252 Return whether all elements are truthy.
1254 Returns True unless there is at least one element that is falsey.
1255 By default, NAs are skipped. If ``skipna=False`` is specified and
1256 missing values are present, similar :ref:`Kleene logic <boolean.kleene>`
1257 is used as for logical operations.
1259 Parameters
1260 ----------
1261 skipna : bool, default True
1262 Exclude NA values. If the entire array is NA and `skipna` is
1263 True, then the result will be True, as for an empty array.
1264 If `skipna` is False, the result will still be False if there is
1265 at least one element that is falsey, otherwise NA will be returned
1266 if there are NA's present.
1268 Returns
1269 -------
1270 bool or :attr:`pandas.NA`
1272 See Also
1273 --------
1274 ArrowExtensionArray.any : Return whether any element is truthy.
1276 Examples
1277 --------
1278 The result indicates whether all elements are truthy (and by default
1279 skips NAs):
1281 >>> pd.array([True, True, pd.NA], dtype="boolean[pyarrow]").all()
1282 True
1283 >>> pd.array([1, 1, pd.NA], dtype="boolean[pyarrow]").all()
1284 True
1285 >>> pd.array([True, False, pd.NA], dtype="boolean[pyarrow]").all()
1286 False
1287 >>> pd.array([], dtype="boolean[pyarrow]").all()
1288 True
1289 >>> pd.array([pd.NA], dtype="boolean[pyarrow]").all()
1290 True
1291 >>> pd.array([pd.NA], dtype="float64[pyarrow]").all()
1292 True
1294 With ``skipna=False``, the result can be NA if this is logically
1295 required (whether ``pd.NA`` is True or False influences the result):
1297 >>> pd.array([True, True, pd.NA], dtype="boolean[pyarrow]").all(skipna=False)
1298 <NA>
1299 >>> pd.array([1, 1, pd.NA], dtype="boolean[pyarrow]").all(skipna=False)
1300 <NA>
1301 >>> pd.array([True, False, pd.NA], dtype="boolean[pyarrow]").all(skipna=False)
1302 False
1303 >>> pd.array([1, 0, pd.NA], dtype="boolean[pyarrow]").all(skipna=False)
1304 False
1305 """
1306 return self._reduce("all", skipna=skipna, **kwargs)
1308 def argsort(
1309 self,
1310 *,
1311 ascending: bool = True,
1312 kind: SortKind = "quicksort",
1313 na_position: str = "last",
1314 **kwargs,
1315 ) -> np.ndarray:
1316 order = "ascending" if ascending else "descending"
1317 null_placement = {"last": "at_end", "first": "at_start"}.get(na_position, None)
1318 if null_placement is None:
1319 raise ValueError(f"invalid na_position: {na_position}")
1321 result = pc.array_sort_indices(
1322 self._pa_array, order=order, null_placement=null_placement
1323 )
1324 np_result = result.to_numpy()
1325 return np_result.astype(np.intp, copy=False)
1327 def _argmin_max(self, skipna: bool, method: str) -> int:
1328 if self._pa_array.length() in (0, self._pa_array.null_count) or (
1329 self._hasna and not skipna
1330 ):
1331 # For empty or all null, pyarrow returns -1 but pandas expects TypeError
1332 # For skipna=False and data w/ null, pandas expects NotImplementedError
1333 # let ExtensionArray.arg{max|min} raise
1334 return getattr(super(), f"arg{method}")(skipna=skipna)
1336 data = self._pa_array
1337 if pa.types.is_duration(data.type):
1338 data = data.cast(pa.int64())
1340 value = getattr(pc, method)(data, skip_nulls=skipna)
1341 return pc.index(data, value).as_py()
1343 def argmin(self, skipna: bool = True) -> int:
1344 return self._argmin_max(skipna, "min")
1346 def argmax(self, skipna: bool = True) -> int:
1347 return self._argmin_max(skipna, "max")
1349 def copy(self) -> Self:
1350 """
1351 Return a shallow copy of the array.
1353 Underlying ChunkedArray is immutable, so a deep copy is unnecessary.
1355 Returns
1356 -------
1357 type(self)
1358 """
1359 return self._from_pyarrow_array(self._pa_array)
1361 def dropna(self) -> Self:
1362 """
1363 Return ArrowExtensionArray without NA values.
1365 Returns
1366 -------
1367 ArrowExtensionArray
1368 """
1369 return self._from_pyarrow_array(pc.drop_null(self._pa_array))
1371 def _pad_or_backfill(
1372 self,
1373 *,
1374 method: FillnaOptions,
1375 limit: int | None = None,
1376 limit_area: Literal["inside", "outside"] | None = None,
1377 copy: bool = True,
1378 ) -> Self:
1379 if not self._hasna:
1380 return self
1382 if limit is None and limit_area is None:
1383 method = missing.clean_fill_method(method)
1384 try:
1385 if method == "pad":
1386 return self._from_pyarrow_array(
1387 pc.fill_null_forward(self._pa_array)
1388 )
1389 elif method == "backfill":
1390 return self._from_pyarrow_array(
1391 pc.fill_null_backward(self._pa_array)
1392 )
1393 except pa.ArrowNotImplementedError:
1394 # ArrowNotImplementedError: Function 'coalesce' has no kernel
1395 # matching input types (duration[ns], duration[ns])
1396 # TODO: remove try/except wrapper if/when pyarrow implements
1397 # a kernel for duration types.
1398 pass
1400 # TODO: Why do we no longer need the above cases?
1401 # TODO(3.0): after EA.fillna 'method' deprecation is enforced, we can remove
1402 # this method entirely.
1403 return super()._pad_or_backfill(
1404 method=method, limit=limit, limit_area=limit_area, copy=copy
1405 )
1407 @doc(ExtensionArray.fillna)
1408 def fillna(
1409 self,
1410 value: object | ArrayLike,
1411 limit: int | None = None,
1412 copy: bool = True,
1413 ) -> Self:
1414 if not self._hasna:
1415 return self.copy()
1417 if limit is not None:
1418 return super().fillna(value=value, limit=limit, copy=copy)
1420 if isinstance(value, (np.ndarray, ExtensionArray)):
1421 # Similar to check_value_size, but we do not mask here since we may
1422 # end up passing it to the super() method.
1423 if len(value) != len(self):
1424 raise ValueError(
1425 f"Length of 'value' does not match. Got ({len(value)}) "
1426 f" expected {len(self)}"
1427 )
1429 try:
1430 fill_value = self._box_pa(value, pa_type=self._pa_array.type)
1431 except pa.ArrowTypeError as err:
1432 msg = f"Invalid value '{value!s}' for dtype '{self.dtype}'"
1433 raise TypeError(msg) from err
1435 try:
1436 return self._from_pyarrow_array(
1437 _safe_fill_null(self._pa_array, fill_value=fill_value)
1438 )
1439 except pa.ArrowNotImplementedError:
1440 # ArrowNotImplementedError: Function 'coalesce' has no kernel
1441 # matching input types (duration[ns], duration[ns])
1442 # TODO: remove try/except wrapper if/when pyarrow implements
1443 # a kernel for duration types.
1444 pass
1446 return super().fillna(value=value, limit=limit, copy=copy)
1448 def isin(self, values: ArrayLike) -> npt.NDArray[np.bool_]:
1449 # short-circuit to return all False array.
1450 if not len(values):
1451 return np.zeros(len(self), dtype=bool)
1453 value_set = self._box_pa(values)
1454 result = pc.is_in(self._pa_array, value_set=value_set)
1455 # pyarrow 2.0.0 returned nulls, so we explicitly specify dtype to convert nulls
1456 # to False
1457 return np.array(result, dtype=np.bool_)
1459 def _values_for_factorize(self) -> tuple[np.ndarray, Any]:
1460 """
1461 Return an array and missing value suitable for factorization.
1463 Returns
1464 -------
1465 values : ndarray
1466 na_value : pd.NA
1468 Notes
1469 -----
1470 The values returned by this method are also used in
1471 :func:`pandas.util.hash_pandas_object`.
1472 """
1473 values = self._pa_array.to_numpy()
1474 return values, self.dtype.na_value
1476 @doc(ExtensionArray.factorize)
1477 def factorize(
1478 self,
1479 use_na_sentinel: bool = True,
1480 ) -> tuple[np.ndarray, ExtensionArray]:
1481 null_encoding = "mask" if use_na_sentinel else "encode"
1483 data = self._pa_array
1485 if pa.types.is_dictionary(data.type):
1486 if null_encoding == "encode":
1487 # dictionary encode does nothing if an already encoded array is given
1488 data = data.cast(data.type.value_type)
1489 encoded = data.dictionary_encode(null_encoding=null_encoding)
1490 else:
1491 encoded = data
1492 else:
1493 encoded = data.dictionary_encode(null_encoding=null_encoding)
1494 if encoded.length() == 0:
1495 indices = np.array([], dtype=np.intp)
1496 uniques = self._from_pyarrow_array(
1497 pa.chunked_array([], type=encoded.type.value_type)
1498 )
1499 else:
1500 # GH 54844
1501 combined = encoded.combine_chunks()
1502 pa_indices = combined.indices
1503 if pa_indices.null_count > 0:
1504 pa_indices = _safe_fill_null(pa_indices, -1)
1505 indices = pa_indices.to_numpy(zero_copy_only=False, writable=True).astype(
1506 np.intp, copy=False
1507 )
1508 uniques = self._from_pyarrow_array(combined.dictionary)
1510 return indices, uniques
1512 def reshape(self, *args, **kwargs):
1513 raise NotImplementedError(
1514 f"{type(self)} does not support reshape "
1515 f"as backed by a 1D pyarrow.ChunkedArray."
1516 )
1518 def round(self, decimals: int = 0, *args, **kwargs) -> Self:
1519 """
1520 Round each value in the array a to the given number of decimals.
1522 Parameters
1523 ----------
1524 decimals : int, default 0
1525 Number of decimal places to round to. If decimals is negative,
1526 it specifies the number of positions to the left of the decimal point.
1527 *args, **kwargs
1528 Additional arguments and keywords have no effect.
1530 Returns
1531 -------
1532 ArrowExtensionArray
1533 Rounded values of the ArrowExtensionArray.
1535 See Also
1536 --------
1537 DataFrame.round : Round values of a DataFrame.
1538 Series.round : Round values of a Series.
1539 """
1540 return self._from_pyarrow_array(pc.round(self._pa_array, ndigits=decimals))
1542 @doc(ExtensionArray.searchsorted)
1543 def searchsorted(
1544 self,
1545 value: NumpyValueArrayLike | ExtensionArray,
1546 side: Literal["left", "right"] = "left",
1547 sorter: NumpySorter | None = None,
1548 ) -> npt.NDArray[np.intp] | np.intp:
1549 if self._hasna:
1550 raise ValueError(
1551 "searchsorted requires array to be sorted, which is impossible "
1552 "with NAs present."
1553 )
1554 if isinstance(value, ExtensionArray):
1555 value = value.astype(object)
1556 # Base class searchsorted would cast to object, which is *much* slower.
1557 dtype = None
1558 if isinstance(self.dtype, ArrowDtype):
1559 pa_dtype = self.dtype.pyarrow_dtype
1560 if (
1561 pa.types.is_timestamp(pa_dtype) or pa.types.is_duration(pa_dtype)
1562 ) and pa_dtype.unit == "ns":
1563 # np.array[datetime/timedelta].searchsorted(datetime/timedelta)
1564 # erroneously fails when numpy type resolution is nanoseconds
1565 dtype = object
1566 return self.to_numpy(dtype=dtype).searchsorted(value, side=side, sorter=sorter)
1568 def take(
1569 self,
1570 indices: TakeIndexer,
1571 allow_fill: bool = False,
1572 fill_value: Any = None,
1573 ) -> ArrowExtensionArray:
1574 """
1575 Take elements from an array.
1577 Parameters
1578 ----------
1579 indices : sequence of int or one-dimensional np.ndarray of int
1580 Indices to be taken.
1581 allow_fill : bool, default False
1582 How to handle negative values in `indices`.
1584 * False: negative values in `indices` indicate positional indices
1585 from the right (the default). This is similar to
1586 :func:`numpy.take`.
1588 * True: negative values in `indices` indicate
1589 missing values. These values are set to `fill_value`. Any other
1590 other negative values raise a ``ValueError``.
1592 fill_value : any, optional
1593 Fill value to use for NA-indices when `allow_fill` is True.
1594 This may be ``None``, in which case the default NA value for
1595 the type, ``self.dtype.na_value``, is used.
1597 For many ExtensionArrays, there will be two representations of
1598 `fill_value`: a user-facing "boxed" scalar, and a low-level
1599 physical NA value. `fill_value` should be the user-facing version,
1600 and the implementation should handle translating that to the
1601 physical version for processing the take if necessary.
1603 Returns
1604 -------
1605 ExtensionArray
1607 Raises
1608 ------
1609 IndexError
1610 When the indices are out of bounds for the array.
1611 ValueError
1612 When `indices` contains negative values other than ``-1``
1613 and `allow_fill` is True.
1615 See Also
1616 --------
1617 numpy.take
1618 api.extensions.take
1620 Notes
1621 -----
1622 ExtensionArray.take is called by ``Series.__getitem__``, ``.loc``,
1623 ``iloc``, when `indices` is a sequence of values. Additionally,
1624 it's called by :meth:`Series.reindex`, or any other method
1625 that causes realignment, with a `fill_value`.
1626 """
1627 indices_array = np.asanyarray(indices)
1629 if len(self._pa_array) == 0 and (indices_array >= 0).any():
1630 raise IndexError("cannot do a non-empty take")
1631 if indices_array.size > 0 and indices_array.max() >= len(self._pa_array):
1632 raise IndexError("out of bounds value in 'indices'.")
1634 if allow_fill:
1635 fill_mask = indices_array < 0
1636 if fill_mask.any():
1637 validate_indices(indices_array, len(self._pa_array))
1638 # TODO(ARROW-9433): Treat negative indices as NULL
1639 indices_array = pa.array(indices_array, mask=fill_mask)
1640 result = self._pa_array.take(indices_array)
1641 if isna(fill_value):
1642 return self._from_pyarrow_array(result)
1643 # TODO: ArrowNotImplementedError: Function fill_null has no
1644 # kernel matching input types (array[string], scalar[string])
1645 result = self._from_pyarrow_array(result)
1646 result[fill_mask] = fill_value
1647 return result
1648 # return type(self)(pc.fill_null(result, pa.scalar(fill_value)))
1649 else:
1650 # Nothing to fill
1651 return self._from_pyarrow_array(self._pa_array.take(indices))
1652 else: # allow_fill=False
1653 # TODO(ARROW-9432): Treat negative indices as indices from the right.
1654 if (indices_array < 0).any():
1655 # Don't modify in-place
1656 indices_array = np.copy(indices_array)
1657 indices_array[indices_array < 0] += len(self._pa_array)
1658 return self._from_pyarrow_array(self._pa_array.take(indices_array))
1660 def _maybe_convert_datelike_array(self):
1661 """Maybe convert to a datelike array."""
1662 pa_type = self._pa_array.type
1663 if pa.types.is_timestamp(pa_type):
1664 return self._to_datetimearray()
1665 elif pa.types.is_duration(pa_type):
1666 return self._to_timedeltaarray()
1667 return self
1669 def _to_datetimearray(self) -> DatetimeArray:
1670 """Convert a pyarrow timestamp typed array to a DatetimeArray."""
1671 from pandas.core.arrays.datetimes import (
1672 DatetimeArray,
1673 tz_to_dtype,
1674 )
1676 pa_type = self._pa_array.type
1677 assert pa.types.is_timestamp(pa_type)
1678 np_dtype = np.dtype(f"M8[{pa_type.unit}]")
1679 dtype = tz_to_dtype(pa_type.tz, pa_type.unit)
1680 np_array = self._pa_array.to_numpy()
1681 np_array = np_array.astype(np_dtype, copy=False)
1682 return DatetimeArray._simple_new(np_array, dtype=dtype)
1684 def _to_timedeltaarray(self) -> TimedeltaArray:
1685 """Convert a pyarrow duration typed array to a TimedeltaArray."""
1686 from pandas.core.arrays.timedeltas import TimedeltaArray
1688 pa_type = self._pa_array.type
1689 assert pa.types.is_duration(pa_type)
1690 np_dtype = np.dtype(f"m8[{pa_type.unit}]")
1691 np_array = self._pa_array.to_numpy()
1692 np_array = np_array.astype(np_dtype, copy=False)
1693 return TimedeltaArray._simple_new(np_array, dtype=np_dtype)
1695 def _values_for_json(self) -> np.ndarray:
1696 if is_numeric_dtype(self.dtype):
1697 return np.asarray(self, dtype=object)
1698 return super()._values_for_json()
1700 @doc(ExtensionArray.to_numpy)
1701 def to_numpy(
1702 self,
1703 dtype: npt.DTypeLike | None = None,
1704 copy: bool = False,
1705 na_value: object = lib.no_default,
1706 ) -> np.ndarray:
1707 original_na_value = na_value
1708 dtype, na_value = to_numpy_dtype_inference(self, dtype, na_value, self._hasna)
1709 pa_type = self._pa_array.type
1710 if not self._hasna or isna(na_value) or pa.types.is_null(pa_type):
1711 data = self
1712 else:
1713 data = self.fillna(na_value)
1714 copy = False
1716 if pa.types.is_timestamp(pa_type) or pa.types.is_duration(pa_type):
1717 # GH 55997
1718 if dtype != object and na_value is self.dtype.na_value:
1719 na_value = lib.no_default
1720 result = data._maybe_convert_datelike_array().to_numpy(
1721 dtype=dtype, na_value=na_value
1722 )
1723 elif pa.types.is_time(pa_type) or pa.types.is_date(pa_type):
1724 # convert to list of python datetime.time objects before
1725 # wrapping in ndarray
1726 result = np.array(list(data), dtype=dtype)
1727 if data._hasna:
1728 result[data.isna()] = na_value
1729 elif pa.types.is_null(pa_type):
1730 if dtype is not None and isna(na_value):
1731 na_value = None
1732 result = np.full(len(data), fill_value=na_value, dtype=dtype)
1733 elif not data._hasna or (
1734 pa.types.is_floating(pa_type)
1735 and (
1736 na_value is np.nan
1737 or (
1738 original_na_value is lib.no_default
1739 and is_float_dtype(dtype)
1740 and is_nan_na()
1741 )
1742 )
1743 ):
1744 result = data._pa_array.to_numpy()
1745 if dtype is not None:
1746 result = result.astype(dtype, copy=False)
1747 if copy:
1748 result = result.copy()
1749 else:
1750 if dtype is None:
1751 empty = pa.array([], type=pa_type).to_numpy(zero_copy_only=False)
1752 if can_hold_element(empty, na_value):
1753 dtype = empty.dtype
1754 else:
1755 dtype = np.object_
1756 result = np.empty(len(data), dtype=dtype)
1757 mask = data.isna()
1758 result[mask] = na_value
1759 result[~mask] = data[~mask]._pa_array.to_numpy()
1760 return result
1762 def map(self, mapper, na_action: Literal["ignore"] | None = None):
1763 if is_numeric_dtype(self.dtype):
1764 return map_array(self.to_numpy(), mapper, na_action=na_action)
1765 else:
1766 # For "mM" cases, the super() method passes `self` without the
1767 # to_numpy call, which inside map_array casts to ndarray[object].
1768 # Without the to_numpy() call, NA is preserved instead of changed
1769 # to None.
1770 return super().map(mapper, na_action)
1772 @doc(ExtensionArray.duplicated)
1773 def duplicated(
1774 self, keep: Literal["first", "last", False] = "first"
1775 ) -> npt.NDArray[np.bool_]:
1776 pa_type = self._pa_array.type
1777 if pa.types.is_floating(pa_type) or pa.types.is_integer(pa_type):
1778 values = self.to_numpy(na_value=0)
1779 elif pa.types.is_boolean(pa_type):
1780 values = self.to_numpy(na_value=False)
1781 elif pa.types.is_temporal(pa_type):
1782 if pa_type.bit_width == 32:
1783 pa_type = pa.int32()
1784 else:
1785 pa_type = pa.int64()
1786 arr = self.astype(ArrowDtype(pa_type))
1787 values = arr.to_numpy(na_value=0)
1788 else:
1789 # factorize the values to avoid the performance penalty of
1790 # converting to object dtype
1791 values = self.factorize()[0]
1793 mask = self.isna() if self._hasna else None
1794 return algos.duplicated(values, keep=keep, mask=mask)
1796 def unique(self) -> Self:
1797 """
1798 Compute the ArrowExtensionArray of unique values.
1800 Returns
1801 -------
1802 ArrowExtensionArray
1803 """
1804 pa_result = pc.unique(self._pa_array)
1805 return self._from_pyarrow_array(pa_result)
1807 def value_counts(self, dropna: bool = True) -> Series:
1808 """
1809 Return a Series containing counts of each unique value.
1811 Parameters
1812 ----------
1813 dropna : bool, default True
1814 Don't include counts of missing values.
1816 Returns
1817 -------
1818 counts : Series
1820 See Also
1821 --------
1822 Series.value_counts
1823 """
1824 from pandas import (
1825 Index,
1826 Series,
1827 )
1829 data = self._pa_array
1830 vc = data.value_counts()
1832 values = vc.field(0)
1833 counts = vc.field(1)
1834 if dropna and data.null_count > 0:
1835 mask = values.is_valid()
1836 values = values.filter(mask)
1837 counts = counts.filter(mask)
1839 counts = ArrowExtensionArray(counts)
1841 index = Index(self._from_pyarrow_array(values), copy=False)
1843 return Series(counts, index=index, name="count", copy=False)
1845 @classmethod
1846 def _concat_same_type(cls, to_concat) -> Self:
1847 """
1848 Concatenate multiple ArrowExtensionArrays.
1850 Parameters
1851 ----------
1852 to_concat : sequence of ArrowExtensionArrays
1854 Returns
1855 -------
1856 ArrowExtensionArray
1857 """
1858 chunks = [array for ea in to_concat for array in ea._pa_array.iterchunks()]
1859 if to_concat[0].dtype == "string":
1860 # StringDtype has no attribute pyarrow_dtype
1861 pa_dtype = pa.large_string()
1862 else:
1863 pa_dtype = to_concat[0].dtype.pyarrow_dtype
1864 arr = pa.chunked_array(chunks, type=pa_dtype)
1865 return to_concat[0]._from_pyarrow_array(arr)
1867 def _accumulate(
1868 self, name: str, *, skipna: bool = True, **kwargs
1869 ) -> ArrowExtensionArray | ExtensionArray:
1870 """
1871 Return an ExtensionArray performing an accumulation operation.
1873 The underlying data type might change.
1875 Parameters
1876 ----------
1877 name : str
1878 Name of the function, supported values are:
1879 - cummin
1880 - cummax
1881 - cumsum
1882 - cumprod
1883 skipna : bool, default True
1884 If True, skip NA values.
1885 **kwargs
1886 Additional keyword arguments passed to the accumulation function.
1887 Currently, there is no supported kwarg.
1889 Returns
1890 -------
1891 array
1893 Raises
1894 ------
1895 NotImplementedError : subclass does not define accumulations
1896 """
1897 if is_string_dtype(self):
1898 return self._str_accumulate(name=name, skipna=skipna, **kwargs)
1900 pyarrow_name = {
1901 "cummax": "cumulative_max",
1902 "cummin": "cumulative_min",
1903 "cumprod": "cumulative_prod_checked",
1904 "cumsum": "cumulative_sum_checked",
1905 }.get(name, name)
1906 pyarrow_meth = getattr(pc, pyarrow_name, None)
1907 if pyarrow_meth is None:
1908 return super()._accumulate(name, skipna=skipna, **kwargs)
1910 data_to_accum = self._pa_array
1912 pa_dtype = data_to_accum.type
1914 convert_to_int = (
1915 pa.types.is_temporal(pa_dtype) and name in ["cummax", "cummin"]
1916 ) or (pa.types.is_duration(pa_dtype) and name == "cumsum")
1918 if convert_to_int:
1919 if pa_dtype.bit_width == 32:
1920 data_to_accum = data_to_accum.cast(pa.int32())
1921 else:
1922 data_to_accum = data_to_accum.cast(pa.int64())
1924 try:
1925 result = pyarrow_meth(data_to_accum, skip_nulls=skipna, **kwargs)
1926 except pa.ArrowNotImplementedError as err:
1927 msg = f"operation '{name}' not supported for dtype '{self.dtype}'"
1928 raise TypeError(msg) from err
1930 if convert_to_int:
1931 result = result.cast(pa_dtype)
1933 return self._from_pyarrow_array(result)
1935 def _str_accumulate(
1936 self, name: str, *, skipna: bool = True, **kwargs
1937 ) -> ArrowExtensionArray | ExtensionArray:
1938 """
1939 Accumulate implementation for strings, see `_accumulate` docstring for details.
1941 pyarrow.compute does not implement these methods for strings.
1942 """
1943 if name == "cumprod":
1944 msg = f"operation '{name}' not supported for dtype '{self.dtype}'"
1945 raise TypeError(msg)
1947 # We may need to strip out trailing NA values
1948 tail: pa.array | None = None
1949 na_mask: pa.array | None = None
1950 pa_array = self._pa_array
1951 np_func = {
1952 "cumsum": np.cumsum,
1953 "cummin": np.minimum.accumulate,
1954 "cummax": np.maximum.accumulate,
1955 }[name]
1957 if self._hasna:
1958 na_mask = pc.is_null(pa_array)
1959 if pc.all(na_mask) == pa.scalar(True):
1960 return self._from_pyarrow_array(pa_array)
1961 if skipna:
1962 if name == "cumsum":
1963 pa_array = _safe_fill_null(pa_array, "")
1964 else:
1965 # We can retain the running min/max by forward/backward filling.
1966 pa_array = pc.fill_null_forward(pa_array)
1967 pa_array = pc.fill_null_backward(pa_array)
1968 else:
1969 # When not skipping NA values, the result should be null from
1970 # the first NA value onward.
1971 idx = pc.index(na_mask, True).as_py()
1972 tail = pa.nulls(len(pa_array) - idx, type=pa_array.type)
1973 pa_array = pa_array[:idx]
1975 # error: Cannot call function of unknown type
1976 pa_result = pa.array(np_func(pa_array), type=pa_array.type) # type: ignore[operator]
1978 if tail is not None:
1979 pa_result = pa.concat_arrays([pa_result, tail])
1980 elif na_mask is not None:
1981 pa_result = pc.if_else(na_mask, None, pa_result)
1983 result = self._from_pyarrow_array(pa_result)
1984 return result
1986 def _reduce_pyarrow(self, name: str, *, skipna: bool = True, **kwargs) -> pa.Scalar:
1987 """
1988 Return a pyarrow scalar result of performing the reduction operation.
1990 Parameters
1991 ----------
1992 name : str
1993 Name of the function, supported values are:
1994 { any, all, min, max, sum, mean, median, prod,
1995 std, var, sem, kurt, skew }.
1996 skipna : bool, default True
1997 If True, skip NaN values.
1998 **kwargs
1999 Additional keyword arguments passed to the reduction function.
2000 Currently, `ddof` is the only supported kwarg.
2002 Returns
2003 -------
2004 pyarrow scalar
2006 Raises
2007 ------
2008 TypeError : subclass does not define reductions
2009 """
2010 pa_type = self._pa_array.type
2012 data_to_reduce = self._pa_array
2014 if name in ["any", "all"] and (
2015 pa.types.is_integer(pa_type)
2016 or pa.types.is_floating(pa_type)
2017 or pa.types.is_duration(pa_type)
2018 or pa.types.is_decimal(pa_type)
2019 ):
2020 # pyarrow only supports any/all for boolean dtype, we allow
2021 # for other dtypes, matching our non-pyarrow behavior
2023 if pa.types.is_duration(pa_type):
2024 data_to_cmp = self._pa_array.cast(pa.int64())
2025 else:
2026 data_to_cmp = self._pa_array
2028 not_eq = pc.not_equal(data_to_cmp, 0)
2029 data_to_reduce = not_eq
2031 elif name in ["min", "max", "sum"] and pa.types.is_duration(pa_type):
2032 data_to_reduce = self._pa_array.cast(pa.int64())
2034 elif name in ["median", "mean", "std", "sem"] and pa.types.is_temporal(pa_type):
2035 nbits = pa_type.bit_width
2036 if nbits == 32:
2037 data_to_reduce = self._pa_array.cast(pa.int32())
2038 else:
2039 data_to_reduce = self._pa_array.cast(pa.int64())
2041 if name == "sem":
2043 def pyarrow_meth(data, skip_nulls, **kwargs):
2044 numerator = pc.stddev(data, skip_nulls=skip_nulls, **kwargs)
2045 denominator = pc.sqrt_checked(pc.count(self._pa_array))
2046 return pc.divide_checked(numerator, denominator)
2048 elif name == "sum" and (
2049 pa.types.is_string(pa_type) or pa.types.is_large_string(pa_type)
2050 ):
2052 def pyarrow_meth(data, skip_nulls, min_count=0): # type: ignore[misc]
2053 mask = pc.is_null(data) if data.null_count > 0 else None
2054 if skip_nulls:
2055 if min_count > 0 and check_below_min_count(
2056 (len(data),),
2057 None if mask is None else mask.to_numpy(),
2058 min_count,
2059 ):
2060 return pa.scalar(None, type=data.type)
2061 if data.null_count > 0:
2062 # binary_join returns null if there is any null ->
2063 # have to filter out any nulls
2064 data = data.filter(pc.invert(mask))
2065 elif mask is not None or check_below_min_count(
2066 (len(data),), None, min_count
2067 ):
2068 return pa.scalar(None, type=data.type)
2070 if pa.types.is_large_string(data.type):
2071 # binary_join only supports string, not large_string
2072 data = data.cast(pa.string())
2073 data_list = pa.ListArray.from_arrays(
2074 [0, len(data)], data.combine_chunks()
2075 )[0]
2076 return pc.binary_join(data_list, "")
2078 else:
2079 pyarrow_name = {
2080 "median": "quantile",
2081 "prod": "product",
2082 "std": "stddev",
2083 "var": "variance",
2084 }.get(name, name)
2085 # error: Incompatible types in assignment
2086 # (expression has type "Optional[Any]", variable has type
2087 # "Callable[[Any, Any, KwArg(Any)], Any]")
2088 pyarrow_meth = getattr(pc, pyarrow_name, None) # type: ignore[assignment]
2089 if pyarrow_meth is None:
2090 # Let ExtensionArray._reduce raise the TypeError
2091 return super()._reduce(name, skipna=skipna, **kwargs)
2093 # GH51624: pyarrow defaults to min_count=1, pandas behavior is min_count=0
2094 if name in ["any", "all"] and "min_count" not in kwargs:
2095 kwargs["min_count"] = 0
2096 elif name == "median":
2097 # GH 52679: Use quantile instead of approximate_median
2098 kwargs["q"] = 0.5
2100 try:
2101 result = pyarrow_meth(data_to_reduce, skip_nulls=skipna, **kwargs)
2102 except (AttributeError, NotImplementedError, TypeError) as err:
2103 msg = (
2104 f"'{type(self).__name__}' with dtype {self.dtype} "
2105 f"does not support operation '{name}' with pyarrow "
2106 f"version {pa.__version__}. '{name}' may be supported by "
2107 f"upgrading pyarrow."
2108 )
2109 raise TypeError(msg) from err
2110 if name == "median":
2111 # GH 52679: Use quantile instead of approximate_median; returns array
2112 result = result[0]
2114 if name in ["min", "max", "sum"] and pa.types.is_duration(pa_type):
2115 result = result.cast(pa_type)
2116 if name in ["median", "mean"] and pa.types.is_temporal(pa_type):
2117 nbits = pa_type.bit_width
2118 if nbits == 32:
2119 result = result.cast(pa.int32(), safe=False)
2120 else:
2121 result = result.cast(pa.int64(), safe=False)
2122 result = result.cast(pa_type)
2123 if name in ["std", "sem"] and pa.types.is_temporal(pa_type):
2124 result = result.cast(pa.int64(), safe=False)
2125 if pa.types.is_duration(pa_type):
2126 result = result.cast(pa_type)
2127 elif pa.types.is_time(pa_type):
2128 result = result.cast(pa.duration(pa_type.unit))
2129 elif pa.types.is_date(pa_type):
2130 # go with closest available unit, i.e. "s"
2131 result = result.cast(pa.duration("s"))
2132 else:
2133 # i.e. timestamp
2134 result = result.cast(pa.duration(pa_type.unit))
2136 return result
2138 def _reduce(
2139 self, name: str, *, skipna: bool = True, keepdims: bool = False, **kwargs
2140 ):
2141 """
2142 Return a scalar result of performing the reduction operation.
2144 Parameters
2145 ----------
2146 name : str
2147 Name of the function, supported values are:
2148 { any, all, min, max, sum, mean, median, prod,
2149 std, var, sem, kurt, skew }.
2150 skipna : bool, default True
2151 If True, skip NaN values.
2152 **kwargs
2153 Additional keyword arguments passed to the reduction function.
2154 Currently, `ddof` is the only supported kwarg.
2156 Returns
2157 -------
2158 scalar
2160 Raises
2161 ------
2162 TypeError : subclass does not define reductions
2163 """
2164 result = self._reduce_calc(name, skipna=skipna, keepdims=keepdims, **kwargs)
2165 if isinstance(result, pa.Array):
2166 return self._from_pyarrow_array(result)
2167 else:
2168 return result
2170 def _reduce_calc(
2171 self, name: str, *, skipna: bool = True, keepdims: bool = False, **kwargs
2172 ):
2173 pa_result = self._reduce_pyarrow(name, skipna=skipna, **kwargs)
2175 if keepdims:
2176 if isinstance(pa_result, pa.Scalar):
2177 result = pa.array([pa_result.as_py()], type=pa_result.type)
2178 else:
2179 result = pa.array(
2180 [pa_result],
2181 type=to_pyarrow_type(infer_dtype_from_scalar(pa_result)[0]),
2182 )
2183 return result
2185 if pc.is_null(pa_result).as_py():
2186 return self.dtype.na_value
2187 elif isinstance(pa_result, pa.Scalar):
2188 result = pa_result.as_py()
2189 pa_type = pa_result.type
2190 if pa.types.is_duration(pa_type) and pa_type.unit != "ns":
2191 return Timedelta(result).as_unit(pa_type.unit)
2192 elif pa.types.is_timestamp(pa_type) and pa_type.unit != "ns":
2193 return Timestamp(result).as_unit(pa_type.unit)
2194 return result
2195 else:
2196 return pa_result
2198 def _explode(self):
2199 """
2200 See Series.explode.__doc__.
2201 """
2202 # child class explode method supports only list types; return
2203 # default implementation for non list types.
2204 if not hasattr(self.dtype, "pyarrow_dtype") or (
2205 not pa.types.is_list(self.dtype.pyarrow_dtype)
2206 and not pa.types.is_large_list(self.dtype.pyarrow_dtype)
2207 ):
2208 return super()._explode()
2209 values = self
2210 counts = pa.compute.list_value_length(values._pa_array)
2211 counts = counts.fill_null(1).to_numpy()
2212 fill_value = pa.scalar([None], type=self._pa_array.type)
2213 mask = counts == 0
2214 if mask.any():
2215 # pc.if_else here is similar to `values[mask] = fill_value`
2216 # but this avoids an object-dtype round-trip.
2217 pa_values = pc.if_else(~mask, values._pa_array, fill_value)
2218 values = self._from_pyarrow_array(pa_values)
2219 counts = counts.copy()
2220 counts[mask] = 1
2221 values = values.fillna(fill_value)
2222 values = self._from_pyarrow_array(pa.compute.list_flatten(values._pa_array))
2223 return values, counts
2225 def __setitem__(self, key, value) -> None:
2226 """Set one or more values inplace.
2228 Parameters
2229 ----------
2230 key : int, ndarray, or slice
2231 When called from, e.g. ``Series.__setitem__``, ``key`` will be
2232 one of
2234 * scalar int
2235 * ndarray of integers.
2236 * boolean ndarray
2237 * slice object
2239 value : ExtensionDtype.type, Sequence[ExtensionDtype.type], or object
2240 value or values to be set of ``key``.
2242 Returns
2243 -------
2244 None
2245 """
2246 if self._readonly:
2247 raise ValueError("Cannot modify read-only array")
2249 # GH50085: unwrap 1D indexers
2250 if isinstance(key, tuple) and len(key) == 1:
2251 key = key[0]
2253 key = check_array_indexer(self, key)
2254 value = self._maybe_convert_setitem_value(value)
2256 if com.is_null_slice(key):
2257 # fast path (GH50248)
2258 if (
2259 isinstance(value, (pa.Array, pa.ChunkedArray))
2260 and value.type == self._pa_array.type
2261 and len(value) == len(self)
2262 ):
2263 data = value
2264 else:
2265 data = self._if_else(True, value, self._pa_array)
2267 elif is_integer(key):
2268 # fast path
2269 key = cast(int, key)
2270 n = len(self)
2271 if key < 0:
2272 key += n
2273 if not 0 <= key < n:
2274 raise IndexError(
2275 f"index {key} is out of bounds for axis 0 with size {n}"
2276 )
2277 if isinstance(value, pa.Scalar):
2278 value = value.as_py()
2279 elif is_list_like(value):
2280 raise ValueError("Length of indexer and values mismatch")
2281 chunks = [
2282 *self._pa_array[:key].chunks,
2283 pa.array([value], type=self._pa_array.type, from_pandas=is_nan_na()),
2284 *self._pa_array[key + 1 :].chunks,
2285 ]
2286 data = pa.chunked_array(chunks).combine_chunks()
2288 elif is_bool_dtype(key):
2289 key = np.asarray(key, dtype=np.bool_)
2290 data = self._replace_with_mask(self._pa_array, key, value)
2292 elif is_scalar(value) or isinstance(value, pa.Scalar):
2293 mask = np.zeros(len(self), dtype=np.bool_)
2294 mask[key] = True
2295 data = self._if_else(mask, value, self._pa_array)
2297 else:
2298 indices = np.arange(len(self))[key]
2299 if len(indices) != len(value):
2300 raise ValueError("Length of indexer and values mismatch")
2301 if len(indices) == 0:
2302 return
2303 # GH#58530 wrong item assignment by repeated key
2304 _, argsort = np.unique(indices, return_index=True)
2305 indices = indices[argsort]
2306 value = value.take(argsort)
2307 mask = np.zeros(len(self), dtype=np.bool_)
2308 mask[indices] = True
2309 data = self._replace_with_mask(self._pa_array, mask, value)
2311 if isinstance(data, pa.Array):
2312 data = pa.chunked_array([data])
2313 self._pa_array = data
2315 def _rank_calc(
2316 self,
2317 *,
2318 axis: AxisInt = 0,
2319 method: str = "average",
2320 na_option: str = "keep",
2321 ascending: bool = True,
2322 pct: bool = False,
2323 ):
2324 if axis != 0:
2325 ranked = super()._rank(
2326 axis=axis,
2327 method=method,
2328 na_option=na_option,
2329 ascending=ascending,
2330 pct=pct,
2331 )
2332 # keep dtypes consistent with the implementation below
2333 if method == "average" or pct:
2334 pa_type = pa.float64()
2335 else:
2336 pa_type = pa.uint64()
2337 result = pa.array(ranked, type=pa_type, from_pandas=is_nan_na())
2338 return result
2340 data = self._pa_array.combine_chunks()
2341 sort_keys = "ascending" if ascending else "descending"
2342 null_placement = "at_start" if na_option == "top" else "at_end"
2343 tiebreaker = "min" if method == "average" else method
2345 result = pc.rank(
2346 data,
2347 sort_keys=sort_keys,
2348 null_placement=null_placement,
2349 tiebreaker=tiebreaker,
2350 )
2352 if na_option == "keep":
2353 mask = pc.is_null(self._pa_array)
2354 null = pa.scalar(None, type=result.type)
2355 result = pc.if_else(mask, null, result)
2357 if method == "average":
2358 result_max = pc.rank(
2359 data,
2360 sort_keys=sort_keys,
2361 null_placement=null_placement,
2362 tiebreaker="max",
2363 )
2364 result_max = result_max.cast(pa.float64())
2365 result_min = result.cast(pa.float64())
2366 result = pc.divide(pc.add(result_min, result_max), 2)
2368 if pct:
2369 if not pa.types.is_floating(result.type):
2370 result = result.cast(pa.float64())
2371 if method == "dense":
2372 divisor = pc.max(result)
2373 else:
2374 divisor = pc.count(result)
2375 result = pc.divide(result, divisor)
2377 return result
2379 def _rank(
2380 self,
2381 *,
2382 axis: AxisInt = 0,
2383 method: str = "average",
2384 na_option: str = "keep",
2385 ascending: bool = True,
2386 pct: bool = False,
2387 ) -> Self:
2388 """
2389 See Series.rank.__doc__.
2390 """
2391 return self._convert_rank_result(
2392 self._rank_calc(
2393 axis=axis,
2394 method=method,
2395 na_option=na_option,
2396 ascending=ascending,
2397 pct=pct,
2398 )
2399 )
2401 def _quantile(self, qs: npt.NDArray[np.float64], interpolation: str) -> Self:
2402 """
2403 Compute the quantiles of self for each quantile in `qs`.
2405 Parameters
2406 ----------
2407 qs : np.ndarray[float64]
2408 interpolation: str
2410 Returns
2411 -------
2412 same type as self
2413 """
2414 pa_dtype = self._pa_array.type
2416 data = self._pa_array
2417 if pa.types.is_temporal(pa_dtype):
2418 # https://github.com/apache/arrow/issues/33769 in these cases
2419 # we can cast to ints and back
2420 nbits = pa_dtype.bit_width
2421 if nbits == 32:
2422 data = data.cast(pa.int32())
2423 else:
2424 data = data.cast(pa.int64())
2426 result = pc.quantile(data, q=qs, interpolation=interpolation)
2428 if pa.types.is_temporal(pa_dtype):
2429 if pa.types.is_floating(result.type):
2430 result = pc.floor(result)
2431 nbits = pa_dtype.bit_width
2432 if nbits == 32:
2433 result = result.cast(pa.int32())
2434 else:
2435 result = result.cast(pa.int64())
2436 result = result.cast(pa_dtype)
2438 return self._from_pyarrow_array(result)
2440 def _mode(self, dropna: bool = True) -> Self:
2441 """
2442 Returns the mode(s) of the ExtensionArray.
2444 Always returns `ExtensionArray` even if only one value.
2446 Parameters
2447 ----------
2448 dropna : bool, default True
2449 Don't consider counts of NA values.
2451 Returns
2452 -------
2453 same type as self
2454 Sorted, if possible.
2455 """
2456 pa_type = self._pa_array.type
2457 if pa.types.is_temporal(pa_type):
2458 nbits = pa_type.bit_width
2459 if nbits == 32:
2460 data = self._pa_array.cast(pa.int32())
2461 elif nbits == 64:
2462 data = self._pa_array.cast(pa.int64())
2463 else:
2464 raise NotImplementedError(pa_type)
2465 else:
2466 data = self._pa_array
2468 if dropna:
2469 data = data.drop_null()
2471 res = pc.value_counts(data)
2472 most_common = res.field("values").filter(
2473 pc.equal(res.field("counts"), pc.max(res.field("counts")))
2474 )
2476 if pa.types.is_temporal(pa_type):
2477 most_common = most_common.cast(pa_type)
2479 most_common = most_common.take(pc.array_sort_indices(most_common))
2480 return self._from_pyarrow_array(most_common)
2482 def _maybe_convert_setitem_value(self, value):
2483 """Maybe convert value to be pyarrow compatible."""
2484 try:
2485 value = self._box_pa(value, self._pa_array.type)
2486 except pa.ArrowTypeError as err:
2487 msg = f"Invalid value '{value!s}' for dtype '{self.dtype}'"
2488 raise TypeError(msg) from err
2489 return value
2491 def interpolate(
2492 self,
2493 *,
2494 method: InterpolateOptions,
2495 axis: int,
2496 index,
2497 limit,
2498 limit_direction,
2499 limit_area,
2500 copy: bool,
2501 **kwargs,
2502 ) -> Self:
2503 """
2504 See NDFrame.interpolate.__doc__.
2505 """
2506 # NB: we return type(self) even if copy=False
2507 if not self.dtype._is_numeric:
2508 raise TypeError(f"Cannot interpolate with {self.dtype} dtype")
2510 if (
2511 method == "linear"
2512 and limit_area is None
2513 and limit is None
2514 and limit_direction == "forward"
2515 ):
2516 values = self._pa_array.combine_chunks()
2517 na_value = pa.array([None], type=values.type)
2518 y_diff_2 = pc.fill_null_backward(pc.pairwise_diff_checked(values, period=2))
2519 prev_values = pa.concat_arrays([na_value, values[:-2], na_value])
2520 interps = pc.add_checked(prev_values, pc.divide_checked(y_diff_2, 2))
2521 return self._from_pyarrow_array(pc.coalesce(self._pa_array, interps))
2523 mask = self.isna()
2524 if self.dtype.kind == "f":
2525 data = self._pa_array.to_numpy()
2526 elif self.dtype.kind in "iu":
2527 data = self.to_numpy(dtype="f8", na_value=0.0)
2528 else:
2529 raise NotImplementedError(
2530 f"interpolate is not implemented for dtype={self.dtype}"
2531 )
2533 missing.interpolate_2d_inplace(
2534 data,
2535 method=method,
2536 axis=0,
2537 index=index,
2538 limit=limit,
2539 limit_direction=limit_direction,
2540 limit_area=limit_area,
2541 mask=mask,
2542 **kwargs,
2543 )
2544 return self._from_pyarrow_array(self._box_pa_array(pa.array(data, mask=mask)))
2546 @classmethod
2547 def _if_else(
2548 cls,
2549 cond: npt.NDArray[np.bool_] | bool,
2550 left: ArrayLike | Scalar,
2551 right: ArrayLike | Scalar,
2552 ) -> pa.Array:
2553 """
2554 Choose values based on a condition.
2556 Analogous to pyarrow.compute.if_else, with logic
2557 to fallback to numpy for unsupported types.
2559 Parameters
2560 ----------
2561 cond : npt.NDArray[np.bool_] or bool
2562 left : ArrayLike | Scalar
2563 right : ArrayLike | Scalar
2565 Returns
2566 -------
2567 pa.Array
2568 """
2570 # TODO: Remove this part when pa.if_else is fixed (GH#64320)
2571 def _maybe_combine(arr):
2572 if not isinstance(arr, pa.ChunkedArray) or not (
2573 pa.types.is_string(arr.type) or pa.types.is_large_string(arr.type)
2574 ):
2575 return arr
2576 if not any(c.offset != 0 for c in arr.chunks):
2577 return arr
2578 try:
2579 return arr.combine_chunks()
2580 except (pa.ArrowInvalid, pa.ArrowCapacityError, MemoryError):
2581 return None
2583 left_c, right_c = _maybe_combine(left), _maybe_combine(right)
2584 if left_c is not None and right_c is not None:
2585 try:
2586 return pc.if_else(cond, left_c, right_c)
2587 except pa.ArrowNotImplementedError:
2588 pass
2589 if left_c is not None:
2590 left = left_c
2591 if right_c is not None:
2592 right = right_c
2594 def _to_numpy_and_type(value) -> tuple[np.ndarray, pa.DataType | None]:
2595 if isinstance(value, (pa.Array, pa.ChunkedArray)):
2596 pa_type = value.type
2597 elif isinstance(value, pa.Scalar):
2598 pa_type = value.type
2599 value = value.as_py()
2600 else:
2601 pa_type = None
2602 return np.array(value, dtype=object), pa_type
2604 left, left_type = _to_numpy_and_type(left)
2605 right, right_type = _to_numpy_and_type(right)
2606 pa_type = left_type or right_type
2607 result = np.where(cond, left, right)
2608 return pa.array(result, type=pa_type, from_pandas=is_nan_na())
2610 @classmethod
2611 def _replace_with_mask(
2612 cls,
2613 values: pa.Array | pa.ChunkedArray,
2614 mask: npt.NDArray[np.bool_] | bool,
2615 replacements: ArrayLike | Scalar,
2616 ) -> pa.Array | pa.ChunkedArray:
2617 """
2618 Replace items selected with a mask.
2620 Analogous to pyarrow.compute.replace_with_mask, with logic
2621 to fallback to numpy for unsupported types.
2623 Parameters
2624 ----------
2625 values : pa.Array or pa.ChunkedArray
2626 mask : npt.NDArray[np.bool_] or bool
2627 replacements : ArrayLike or Scalar
2628 Replacement value(s)
2630 Returns
2631 -------
2632 pa.Array or pa.ChunkedArray
2633 """
2634 if isinstance(replacements, pa.ChunkedArray):
2635 # replacements must be array or scalar, not ChunkedArray
2636 replacements = replacements.combine_chunks()
2637 if isinstance(values, pa.ChunkedArray) and pa.types.is_boolean(values.type):
2638 # GH#52059 replace_with_mask segfaults for chunked array
2639 # https://github.com/apache/arrow/issues/34634
2640 values = values.combine_chunks()
2641 try:
2642 return pc.replace_with_mask(values, mask, replacements)
2643 except pa.ArrowNotImplementedError:
2644 pass
2645 if isinstance(replacements, pa.Array):
2646 replacements = np.array(replacements, dtype=object)
2647 elif isinstance(replacements, pa.Scalar):
2648 replacements = replacements.as_py()
2650 result = np.array(values, dtype=object)
2651 result[mask] = replacements
2652 return pa.array(result, type=values.type, from_pandas=is_nan_na())
2654 # ------------------------------------------------------------------
2655 # GroupBy Methods
2657 def _to_masked(self):
2658 pa_dtype = self._pa_array.type
2660 if pa.types.is_floating(pa_dtype) or pa.types.is_integer(pa_dtype):
2661 na_value = 1
2662 elif pa.types.is_boolean(pa_dtype):
2663 na_value = True
2664 else:
2665 raise NotImplementedError
2667 dtype = _arrow_dtype_mapping()[pa_dtype]
2668 mask = self.isna()
2669 arr = self.to_numpy(dtype=dtype.numpy_dtype, na_value=na_value)
2670 return dtype.construct_array_type()(arr, mask)
2672 def _groupby_op(
2673 self,
2674 *,
2675 how: str,
2676 has_dropped_na: bool,
2677 min_count: int,
2678 ngroups: int,
2679 ids: npt.NDArray[np.intp],
2680 **kwargs,
2681 ):
2682 if isinstance(self.dtype, StringDtype):
2683 if how in [
2684 "prod",
2685 "mean",
2686 "median",
2687 "cumsum",
2688 "cumprod",
2689 "std",
2690 "sem",
2691 "var",
2692 "skew",
2693 ]:
2694 raise TypeError(
2695 f"dtype '{self.dtype}' does not support operation '{how}'"
2696 )
2697 return super()._groupby_op(
2698 how=how,
2699 has_dropped_na=has_dropped_na,
2700 min_count=min_count,
2701 ngroups=ngroups,
2702 ids=ids,
2703 **kwargs,
2704 )
2706 # maybe convert to a compatible dtype optimized for groupby
2707 values: ExtensionArray
2708 pa_type = self._pa_array.type
2709 if pa.types.is_timestamp(pa_type):
2710 values = self._to_datetimearray()
2711 elif pa.types.is_duration(pa_type):
2712 values = self._to_timedeltaarray()
2713 else:
2714 values = self._to_masked()
2716 result = values._groupby_op(
2717 how=how,
2718 has_dropped_na=has_dropped_na,
2719 min_count=min_count,
2720 ngroups=ngroups,
2721 ids=ids,
2722 **kwargs,
2723 )
2724 if isinstance(result, np.ndarray):
2725 return result
2726 elif isinstance(result, BaseMaskedArray):
2727 pa_result = result.__arrow_array__()
2728 return self._from_pyarrow_array(pa_result)
2729 else:
2730 # DatetimeArray, TimedeltaArray
2731 pa_result = pa.array(result)
2732 return self._from_pyarrow_array(pa_result)
2734 def _apply_elementwise(self, func: Callable) -> list[list[Any]]:
2735 """Apply a callable to each element while maintaining the chunking structure."""
2736 return [
2737 [
2738 None if val is None else func(val)
2739 for val in chunk.to_numpy(zero_copy_only=False)
2740 ]
2741 for chunk in self._pa_array.iterchunks()
2742 ]
2744 def _convert_bool_result(self, result, na=lib.no_default, method_name=None):
2745 if na is not lib.no_default and not isna(na): # pyright: ignore [reportGeneralTypeIssues]
2746 result = result.fill_null(na)
2747 return self._from_pyarrow_array(result)
2749 def _convert_int_result(self, result):
2750 return self._from_pyarrow_array(result)
2752 def _convert_rank_result(self, result):
2753 return self._from_pyarrow_array(result)
2755 def _str_count(self, pat: str, flags: int = 0) -> Self:
2756 if flags:
2757 raise NotImplementedError(f"count not implemented with {flags=}")
2758 return self._from_pyarrow_array(pc.count_substring_regex(self._pa_array, pat))
2760 def _str_repeat(self, repeats: int | Sequence[int]) -> Self:
2761 if not isinstance(repeats, int):
2762 raise NotImplementedError(
2763 f"repeat is not implemented when repeats is {type(repeats).__name__}"
2764 )
2765 return self._from_pyarrow_array(pc.binary_repeat(self._pa_array, repeats))
2767 def _str_join(self, sep: str) -> Self:
2768 if pa.types.is_string(self._pa_array.type) or pa.types.is_large_string(
2769 self._pa_array.type
2770 ):
2771 result = self._apply_elementwise(list)
2772 result = pa.chunked_array(result, type=pa.list_(pa.string()))
2773 else:
2774 result = self._pa_array
2775 return self._from_pyarrow_array(pc.binary_join(result, sep))
2777 def _str_partition(self, sep: str, expand: bool) -> Self:
2778 predicate = lambda val: val.partition(sep)
2779 result = self._apply_elementwise(predicate)
2780 return self._from_pyarrow_array(pa.chunked_array(result))
2782 def _str_rpartition(self, sep: str, expand: bool) -> Self:
2783 predicate = lambda val: val.rpartition(sep)
2784 result = self._apply_elementwise(predicate)
2785 return self._from_pyarrow_array(pa.chunked_array(result))
2787 def _str_casefold(self) -> Self:
2788 predicate = lambda val: val.casefold()
2789 result = self._apply_elementwise(predicate)
2790 return self._from_pyarrow_array(pa.chunked_array(result))
2792 def _str_encode(self, encoding: str, errors: str = "strict") -> Self:
2793 predicate = lambda val: val.encode(encoding, errors)
2794 result = self._apply_elementwise(predicate)
2795 return self._from_pyarrow_array(pa.chunked_array(result))
2797 def _str_extract(self, pat: str, flags: int = 0, expand: bool = True):
2798 if flags:
2799 raise NotImplementedError("Only flags=0 is implemented.")
2800 groups = re.compile(pat).groupindex.keys()
2801 if len(groups) == 0:
2802 raise ValueError(f"{pat=} must contain a symbolic group name.")
2803 result = pc.extract_regex(self._pa_array, pat)
2804 if expand:
2805 return {
2806 col: self._from_pyarrow_array(pc.struct_field(result, [i]))
2807 for col, i in zip(groups, range(result.type.num_fields), strict=True)
2808 }
2809 else:
2810 return type(self)(pc.struct_field(result, [0]))
2812 def _str_findall(self, pat: str, flags: int = 0) -> Self:
2813 regex = re.compile(pat, flags=flags)
2814 predicate = lambda val: regex.findall(val)
2815 result = self._apply_elementwise(predicate)
2816 return self._from_pyarrow_array(pa.chunked_array(result))
2818 def _str_get_dummies(self, sep: str = "|", dtype: NpDtype | None = None):
2819 if dtype is None:
2820 dtype = np.bool_
2821 split = pc.split_pattern(self._pa_array, sep)
2822 flattened_values = pc.list_flatten(split)
2823 uniques = flattened_values.unique()
2824 uniques_sorted = uniques.take(pa.compute.array_sort_indices(uniques))
2825 lengths = pc.list_value_length(split).fill_null(0).to_numpy()
2826 n_rows = len(self)
2827 n_cols = len(uniques)
2828 indices = pc.index_in(flattened_values, uniques_sorted).to_numpy()
2829 indices = indices + np.arange(n_rows).repeat(lengths) * n_cols
2830 _dtype = pandas_dtype(dtype)
2831 dummies_dtype: NpDtype
2832 if isinstance(_dtype, np.dtype):
2833 dummies_dtype = _dtype
2834 else:
2835 dummies_dtype = np.bool_
2836 dummies = np.zeros(n_rows * n_cols, dtype=dummies_dtype)
2837 dummies[indices] = True
2838 dummies = dummies.reshape((n_rows, n_cols))
2839 result = self._from_pyarrow_array(pa.array(list(dummies)))
2840 return result, uniques_sorted.to_pylist()
2842 def _str_index(self, sub: str, start: int = 0, end: int | None = None) -> Self:
2843 predicate = lambda val: val.index(sub, start, end)
2844 result = self._apply_elementwise(predicate)
2845 return self._from_pyarrow_array(pa.chunked_array(result))
2847 def _str_rindex(self, sub: str, start: int = 0, end: int | None = None) -> Self:
2848 predicate = lambda val: val.rindex(sub, start, end)
2849 result = self._apply_elementwise(predicate)
2850 return self._from_pyarrow_array(pa.chunked_array(result))
2852 def _str_normalize(self, form: Literal["NFC", "NFD", "NFKC", "NFKD"]) -> Self:
2853 predicate = lambda val: unicodedata.normalize(form, val)
2854 result = self._apply_elementwise(predicate)
2855 return self._from_pyarrow_array(pa.chunked_array(result))
2857 def _str_rfind(self, sub: str, start: int = 0, end=None) -> Self:
2858 predicate = lambda val: val.rfind(sub, start, end)
2859 result = self._apply_elementwise(predicate)
2860 return self._from_pyarrow_array(pa.chunked_array(result))
2862 def _str_split(
2863 self,
2864 pat: str | None = None,
2865 n: int | None = -1,
2866 expand: bool = False,
2867 regex: bool | None = None,
2868 ) -> Self:
2869 if n in {-1, 0}:
2870 n = None
2871 if pat is None:
2872 split_func = pc.utf8_split_whitespace
2873 elif regex:
2874 split_func = functools.partial(pc.split_pattern_regex, pattern=pat)
2875 else:
2876 split_func = functools.partial(pc.split_pattern, pattern=pat)
2877 return self._from_pyarrow_array(split_func(self._pa_array, max_splits=n))
2879 def _str_rsplit(self, pat: str | None = None, n: int | None = -1) -> Self:
2880 if n in {-1, 0}:
2881 n = None
2882 if pat is None:
2883 return self._from_pyarrow_array(
2884 pc.utf8_split_whitespace(self._pa_array, max_splits=n, reverse=True)
2885 )
2886 return self._from_pyarrow_array(
2887 pc.split_pattern(self._pa_array, pat, max_splits=n, reverse=True)
2888 )
2890 def _str_translate(self, table: dict[int, str]) -> Self:
2891 predicate = lambda val: val.translate(table)
2892 result = self._apply_elementwise(predicate)
2893 return self._from_pyarrow_array(pa.chunked_array(result))
2895 def _str_wrap(self, width: int, **kwargs) -> Self:
2896 kwargs["width"] = width
2897 tw = textwrap.TextWrapper(**kwargs)
2898 predicate = lambda val: "\n".join(tw.wrap(val))
2899 result = self._apply_elementwise(predicate)
2900 return self._from_pyarrow_array(pa.chunked_array(result))
2902 def _str_zfill(self, width: int) -> Self:
2903 if pa_version_under21p0:
2904 predicate = lambda val: val.zfill(width)
2905 result = self._apply_elementwise(predicate)
2906 return type(self)(pa.chunked_array(result))
2907 return type(self)(pc.utf8_zfill(self._pa_array, width))
2909 @property
2910 def _dt_days(self) -> Self:
2911 return self._from_pyarrow_array(
2912 pa.array(
2913 self._to_timedeltaarray().components.days,
2914 from_pandas=True,
2915 type=pa.int32(),
2916 )
2917 )
2919 @property
2920 def _dt_hours(self) -> Self:
2921 return self._from_pyarrow_array(
2922 pa.array(
2923 self._to_timedeltaarray().components.hours,
2924 from_pandas=True,
2925 type=pa.int32(),
2926 )
2927 )
2929 @property
2930 def _dt_minutes(self) -> Self:
2931 return self._from_pyarrow_array(
2932 pa.array(
2933 self._to_timedeltaarray().components.minutes,
2934 from_pandas=True,
2935 type=pa.int32(),
2936 )
2937 )
2939 @property
2940 def _dt_seconds(self) -> Self:
2941 return self._from_pyarrow_array(
2942 pa.array(
2943 self._to_timedeltaarray().components.seconds,
2944 from_pandas=True,
2945 type=pa.int32(),
2946 )
2947 )
2949 @property
2950 def _dt_milliseconds(self) -> Self:
2951 return self._from_pyarrow_array(
2952 pa.array(
2953 self._to_timedeltaarray().components.milliseconds,
2954 from_pandas=True,
2955 type=pa.int32(),
2956 )
2957 )
2959 @property
2960 def _dt_microseconds(self) -> Self:
2961 return self._from_pyarrow_array(
2962 pa.array(
2963 self._to_timedeltaarray().components.microseconds,
2964 from_pandas=True,
2965 type=pa.int32(),
2966 )
2967 )
2969 @property
2970 def _dt_nanoseconds(self) -> Self:
2971 return self._from_pyarrow_array(
2972 pa.array(
2973 self._to_timedeltaarray().components.nanoseconds,
2974 from_pandas=True,
2975 type=pa.int32(),
2976 )
2977 )
2979 def _dt_to_pytimedelta(self) -> np.ndarray:
2980 data = self._pa_array.to_pylist()
2981 if self._dtype.pyarrow_dtype.unit == "ns":
2982 data = [None if ts is None else ts.to_pytimedelta() for ts in data]
2983 return np.array(data, dtype=object)
2985 def _dt_total_seconds(self) -> Self:
2986 unit = self._pa_array.type.unit
2987 unit_per_second = {"s": 1.0, "ms": 1e3, "us": 1e6, "ns": 1e9}
2988 result = pc.divide(pc.cast(self._pa_array, pa.int64()), unit_per_second[unit])
2989 return self._from_pyarrow_array(result)
2991 def _dt_as_unit(self, unit: str) -> Self:
2992 pa_type = self._pa_array.type
2993 if pa.types.is_timestamp(pa_type):
2994 target_type = pa.timestamp(unit, tz=pa_type.tz)
2995 elif pa.types.is_duration(pa_type):
2996 target_type = pa.duration(unit)
2997 else:
2998 raise NotImplementedError(f"as_unit not implemented for {pa_type}")
2999 # Use safe=False to allow truncation, matching pandas as_unit behavior
3000 result = pc.cast(self._pa_array, target_type, safe=False)
3001 return self._from_pyarrow_array(result)
3003 @property
3004 def _dt_year(self) -> Self:
3005 result = pc.year(self._pa_array)
3006 return self._from_pyarrow_array(result)
3008 @property
3009 def _dt_day(self) -> Self:
3010 result = pc.day(self._pa_array)
3011 return self._from_pyarrow_array(result)
3013 @property
3014 def _dt_day_of_week(self) -> Self:
3015 result = pc.day_of_week(self._pa_array)
3016 return self._from_pyarrow_array(result)
3018 _dt_dayofweek = _dt_day_of_week
3019 _dt_weekday = _dt_day_of_week
3021 @property
3022 def _dt_day_of_year(self) -> Self:
3023 result = pc.day_of_year(self._pa_array)
3024 return self._from_pyarrow_array(result)
3026 _dt_dayofyear = _dt_day_of_year
3028 @property
3029 def _dt_hour(self) -> Self:
3030 result = pc.hour(self._pa_array)
3031 return self._from_pyarrow_array(result)
3033 def _dt_isocalendar(self) -> Self:
3034 result = pc.iso_calendar(self._pa_array)
3035 return self._from_pyarrow_array(result)
3037 @property
3038 def _dt_is_leap_year(self) -> Self:
3039 result = pc.is_leap_year(self._pa_array)
3040 return self._from_pyarrow_array(result)
3042 @property
3043 def _dt_is_month_start(self) -> Self:
3044 result = pc.equal(pc.day(self._pa_array), 1)
3045 return self._from_pyarrow_array(result)
3047 @property
3048 def _dt_is_month_end(self) -> Self:
3049 result = pc.equal(
3050 pc.days_between(
3051 pc.floor_temporal(self._pa_array, unit="day"),
3052 pc.ceil_temporal(self._pa_array, unit="month"),
3053 ),
3054 1,
3055 )
3056 return self._from_pyarrow_array(result)
3058 @property
3059 def _dt_is_year_start(self) -> Self:
3060 result = pc.and_(
3061 pc.equal(pc.month(self._pa_array), 1),
3062 pc.equal(pc.day(self._pa_array), 1),
3063 )
3064 return self._from_pyarrow_array(result)
3066 @property
3067 def _dt_is_year_end(self) -> Self:
3068 result = pc.and_(
3069 pc.equal(pc.month(self._pa_array), 12),
3070 pc.equal(pc.day(self._pa_array), 31),
3071 )
3072 return self._from_pyarrow_array(result)
3074 @property
3075 def _dt_is_quarter_start(self) -> Self:
3076 result = pc.equal(
3077 pc.floor_temporal(self._pa_array, unit="quarter"),
3078 pc.floor_temporal(self._pa_array, unit="day"),
3079 )
3080 return self._from_pyarrow_array(result)
3082 @property
3083 def _dt_is_quarter_end(self) -> Self:
3084 result = pc.equal(
3085 pc.days_between(
3086 pc.floor_temporal(self._pa_array, unit="day"),
3087 pc.ceil_temporal(self._pa_array, unit="quarter"),
3088 ),
3089 1,
3090 )
3091 return self._from_pyarrow_array(result)
3093 @property
3094 def _dt_days_in_month(self) -> Self:
3095 result = pc.days_between(
3096 pc.floor_temporal(self._pa_array, unit="month"),
3097 pc.ceil_temporal(self._pa_array, unit="month"),
3098 )
3099 return self._from_pyarrow_array(result)
3101 _dt_daysinmonth = _dt_days_in_month
3103 @property
3104 def _dt_microsecond(self) -> Self:
3105 # GH 59154
3106 us = pc.microsecond(self._pa_array)
3107 ms_to_us = pc.multiply(pc.millisecond(self._pa_array), 1000)
3108 result = pc.add(us, ms_to_us)
3109 return self._from_pyarrow_array(result)
3111 @property
3112 def _dt_minute(self) -> Self:
3113 result = pc.minute(self._pa_array)
3114 return self._from_pyarrow_array(result)
3116 @property
3117 def _dt_month(self) -> Self:
3118 result = pc.month(self._pa_array)
3119 return self._from_pyarrow_array(result)
3121 @property
3122 def _dt_nanosecond(self) -> Self:
3123 result = pc.nanosecond(self._pa_array)
3124 return self._from_pyarrow_array(result)
3126 @property
3127 def _dt_quarter(self) -> Self:
3128 result = pc.quarter(self._pa_array)
3129 return self._from_pyarrow_array(result)
3131 @property
3132 def _dt_second(self) -> Self:
3133 result = pc.second(self._pa_array)
3134 return self._from_pyarrow_array(result)
3136 @property
3137 def _dt_date(self) -> Self:
3138 result = self._pa_array.cast(pa.date32())
3139 return self._from_pyarrow_array(result)
3141 @property
3142 def _dt_time(self) -> Self:
3143 unit = (
3144 self.dtype.pyarrow_dtype.unit
3145 if self.dtype.pyarrow_dtype.unit in {"us", "ns"}
3146 else "ns"
3147 )
3148 result = self._pa_array.cast(pa.time64(unit))
3149 return self._from_pyarrow_array(result)
3151 @property
3152 def _dt_tz(self):
3153 return timezones.maybe_get_tz(self.dtype.pyarrow_dtype.tz)
3155 @property
3156 def _dt_unit(self):
3157 return self.dtype.pyarrow_dtype.unit
3159 def _dt_normalize(self) -> Self:
3160 result = pc.floor_temporal(self._pa_array, 1, "day")
3161 return self._from_pyarrow_array(result)
3163 def _dt_strftime(self, format: str) -> Self:
3164 result = pc.strftime(self._pa_array, format=format)
3165 return self._from_pyarrow_array(result)
3167 def _round_temporally(
3168 self,
3169 method: Literal["ceil", "floor", "round"],
3170 freq,
3171 ambiguous: TimeAmbiguous = "raise",
3172 nonexistent: TimeNonexistent = "raise",
3173 ) -> Self:
3174 if ambiguous != "raise":
3175 raise NotImplementedError("ambiguous is not supported.")
3176 if nonexistent != "raise":
3177 raise NotImplementedError("nonexistent is not supported.")
3178 offset = to_offset(freq)
3179 if offset is None:
3180 raise ValueError(f"Must specify a valid frequency: {freq}")
3181 pa_supported_unit = {
3182 "Y": "year",
3183 "YS": "year",
3184 "Q": "quarter",
3185 "QS": "quarter",
3186 "M": "month",
3187 "MS": "month",
3188 "W": "week",
3189 "D": "day",
3190 "h": "hour",
3191 "min": "minute",
3192 "s": "second",
3193 "ms": "millisecond",
3194 "us": "microsecond",
3195 "ns": "nanosecond",
3196 }
3197 unit = pa_supported_unit.get(offset._prefix, None)
3198 if unit is None:
3199 raise ValueError(f"{freq=} is not supported")
3200 multiple = offset.n
3201 rounding_method = getattr(pc, f"{method}_temporal")
3202 result = rounding_method(self._pa_array, multiple=multiple, unit=unit)
3203 return self._from_pyarrow_array(result)
3205 def _dt_ceil(
3206 self,
3207 freq,
3208 ambiguous: TimeAmbiguous = "raise",
3209 nonexistent: TimeNonexistent = "raise",
3210 ) -> Self:
3211 return self._round_temporally("ceil", freq, ambiguous, nonexistent)
3213 def _dt_floor(
3214 self,
3215 freq,
3216 ambiguous: TimeAmbiguous = "raise",
3217 nonexistent: TimeNonexistent = "raise",
3218 ) -> Self:
3219 return self._round_temporally("floor", freq, ambiguous, nonexistent)
3221 def _dt_round(
3222 self,
3223 freq,
3224 ambiguous: TimeAmbiguous = "raise",
3225 nonexistent: TimeNonexistent = "raise",
3226 ) -> Self:
3227 return self._round_temporally("round", freq, ambiguous, nonexistent)
3229 def _dt_day_name(self, locale: str | None = None) -> Self:
3230 if locale is None:
3231 locale = "C"
3232 result = pc.strftime(self._pa_array, format="%A", locale=locale)
3233 return self._from_pyarrow_array(result)
3235 def _dt_month_name(self, locale: str | None = None) -> Self:
3236 if locale is None:
3237 locale = "C"
3238 result = pc.strftime(self._pa_array, format="%B", locale=locale)
3239 return self._from_pyarrow_array(result)
3241 def _dt_to_pydatetime(self) -> Series:
3242 from pandas import Series
3244 if pa.types.is_date(self.dtype.pyarrow_dtype):
3245 raise ValueError(
3246 f"to_pydatetime cannot be called with {self.dtype.pyarrow_dtype} type. "
3247 "Convert to pyarrow timestamp type."
3248 )
3249 data = self._pa_array.to_pylist()
3250 if self._dtype.pyarrow_dtype.unit == "ns":
3251 data = [None if ts is None else ts.to_pydatetime(warn=False) for ts in data]
3252 return Series(data, dtype=object)
3254 def _dt_tz_localize(
3255 self,
3256 tz,
3257 ambiguous: TimeAmbiguous = "raise",
3258 nonexistent: TimeNonexistent = "raise",
3259 ) -> Self:
3260 if ambiguous != "raise":
3261 raise NotImplementedError(f"{ambiguous=} is not supported")
3262 nonexistent_pa = {
3263 "raise": "raise",
3264 "shift_backward": "earliest",
3265 "shift_forward": "latest",
3266 }.get(
3267 nonexistent, # type: ignore[arg-type]
3268 None,
3269 )
3270 if nonexistent_pa is None:
3271 raise NotImplementedError(f"{nonexistent=} is not supported")
3272 if tz is None:
3273 result = pc.local_timestamp(self._pa_array)
3274 else:
3275 result = pc.assume_timezone(
3276 self._pa_array, str(tz), ambiguous=ambiguous, nonexistent=nonexistent_pa
3277 )
3278 return self._from_pyarrow_array(result)
3280 def _dt_tz_convert(self, tz) -> Self:
3281 if self.dtype.pyarrow_dtype.tz is None:
3282 raise TypeError(
3283 "Cannot convert tz-naive timestamps, use tz_localize to localize"
3284 )
3285 current_unit = self.dtype.pyarrow_dtype.unit
3286 result = self._pa_array.cast(pa.timestamp(current_unit, tz))
3287 return self._from_pyarrow_array(result)
3290def transpose_homogeneous_pyarrow(
3291 arrays: Sequence[ArrowExtensionArray],
3292) -> list[ArrowExtensionArray]:
3293 """Transpose arrow extension arrays in a list, but faster.
3295 Input should be a list of arrays of equal length and all have the same
3296 dtype. The caller is responsible for ensuring validity of input data.
3297 """
3298 arrays = list(arrays)
3299 nrows, ncols = len(arrays[0]), len(arrays)
3300 indices = np.arange(nrows * ncols).reshape(ncols, nrows).T.reshape(-1)
3301 arr = pa.chunked_array([chunk for arr in arrays for chunk in arr._pa_array.chunks])
3302 arr = arr.take(indices)
3303 return [ArrowExtensionArray(arr.slice(i * ncols, ncols)) for i in range(nrows)]