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 time,
7)
8from decimal import Decimal
9import functools
10import operator
11import re
12import textwrap
13from typing import (
14 TYPE_CHECKING,
15 Any,
16 Literal,
17 Self,
18 cast,
19 overload,
20)
21import warnings
23import numpy as np
25from pandas._config import is_nan_na
27from pandas._libs import lib
28from pandas._libs.missing import is_pdna_or_none
29from pandas._libs.tslibs import (
30 BaseOffset,
31 Timedelta,
32 Timestamp,
33 timezones,
34)
35from pandas.compat import (
36 HAS_PYARROW,
37 PYARROW_MIN_VERSION,
38 pa_version_under21p0,
39 pa_version_under25p0,
40)
41from pandas.compat.numpy import function as nv
42from pandas.errors import (
43 OutOfBoundsDatetime,
44 OutOfBoundsTimedelta,
45 Pandas4Warning,
46)
47from pandas.util._decorators import (
48 cache_readonly,
49 set_module,
50)
51from pandas.util._exceptions import find_stack_level
53from pandas.core.dtypes.cast import (
54 can_hold_element,
55 construct_1d_object_array_from_listlike,
56 infer_dtype_from_scalar,
57)
58from pandas.core.dtypes.common import (
59 is_array_like_deprecate_non_pandas,
60 is_arrow_temporal_dtype,
61 is_bool_dtype,
62 is_float_dtype,
63 is_integer,
64 is_list_like,
65 is_numeric_dtype,
66 is_scalar,
67 is_string_dtype,
68 needs_i8_conversion,
69 pandas_dtype,
70)
71from pandas.core.dtypes.dtypes import (
72 ArrowDtype,
73 CategoricalDtype,
74 DatetimeTZDtype,
75)
76from pandas.core.dtypes.generic import (
77 ABCDataFrame,
78 ABCIndex,
79 ABCSeries,
80)
81from pandas.core.dtypes.missing import isna
83from pandas.core import (
84 algorithms as algos,
85 arraylike,
86 missing,
87 ops,
88 roperator,
89)
90from pandas.core.arraylike import OpsMixin
91from pandas.core.arrays._arrow_string_mixins import ArrowStringArrayMixin
92from pandas.core.arrays._utils import to_numpy_dtype_inference
93from pandas.core.arrays.base import (
94 ExtensionArray,
95 ExtensionArrayNaResult,
96)
97from pandas.core.arrays.masked import BaseMaskedArray
98from pandas.core.arrays.string_ import StringDtype
99import pandas.core.common as com
100from pandas.core.construction import extract_array
101from pandas.core.indexers import (
102 check_array_indexer,
103 getitem_returns_view,
104 unpack_tuple_and_ellipses,
105 validate_indices,
106)
107from pandas.core.nanops import check_below_min_count
109from pandas.io._util import _arrow_dtype_mapping
110from pandas.tseries.frequencies import to_offset
112if HAS_PYARROW:
113 import pyarrow as pa
114 import pyarrow.compute as pc
116 from pandas.compat.pyarrow import _safe_fill_null
118 ARROW_CMP_FUNCS = {
119 "eq": pc.equal,
120 "ne": pc.not_equal,
121 "lt": pc.less,
122 "gt": pc.greater,
123 "le": pc.less_equal,
124 "ge": pc.greater_equal,
125 }
127 ARROW_LOGICAL_FUNCS = {
128 "and_": pc.and_kleene,
129 "rand_": lambda x, y: pc.and_kleene(y, x),
130 "or_": pc.or_kleene,
131 "ror_": lambda x, y: pc.or_kleene(y, x),
132 "xor": pc.xor,
133 "rxor": lambda x, y: pc.xor(y, x),
134 }
136 ARROW_BIT_WISE_FUNCS = {
137 "and_": pc.bit_wise_and,
138 "rand_": lambda x, y: pc.bit_wise_and(y, x),
139 "or_": pc.bit_wise_or,
140 "ror_": lambda x, y: pc.bit_wise_or(y, x),
141 "xor": pc.bit_wise_xor,
142 "rxor": lambda x, y: pc.bit_wise_xor(y, x),
143 }
145 def cast_for_truediv(
146 arrow_array: pa.ChunkedArray, pa_object: pa.Array | pa.Scalar
147 ) -> tuple[pa.ChunkedArray, pa.Array | pa.Scalar]:
148 # Ensure int / int -> float mirroring Python/Numpy behavior
149 # as pc.divide_checked(int, int) -> int
150 if pa.types.is_integer(arrow_array.type) and pa.types.is_integer(
151 pa_object.type
152 ):
153 # GH: 56645.
154 # https://github.com/apache/arrow/issues/35563
155 return pc.cast(arrow_array, pa.float64(), safe=False), pc.cast(
156 pa_object, pa.float64(), safe=False
157 )
159 return arrow_array, pa_object
161 def floordiv_compat(
162 left: pa.ChunkedArray | pa.Array | pa.Scalar,
163 right: pa.ChunkedArray | pa.Array | pa.Scalar,
164 ) -> pa.ChunkedArray:
165 # TODO: Replace with pyarrow floordiv kernel.
166 # https://github.com/apache/arrow/issues/39386
167 if pa.types.is_integer(left.type) and pa.types.is_integer(right.type):
168 divided = pc.divide_checked(left, right)
169 if pa.types.is_signed_integer(divided.type):
170 # GH 56676
171 has_remainder = pc.not_equal(pc.multiply(divided, right), left)
172 has_one_negative_operand = pc.less(
173 pc.bit_wise_xor(left, right),
174 pa.scalar(0, type=divided.type),
175 )
176 result = pc.if_else(
177 pc.and_(
178 has_remainder,
179 has_one_negative_operand,
180 ),
181 # GH: 55561
182 pc.subtract(divided, pa.scalar(1, type=divided.type)),
183 divided,
184 )
185 else:
186 result = divided
187 result = result.cast(left.type)
188 else:
189 divided = pc.divide(left, right)
190 result = pc.floor(divided)
191 return result
193 # Conversion factors for duration units to various time components.
194 # Maps component -> {unit: divisor} to extract that component from the unit.
195 # Units coarser than a component are absent (e.g. "s" is missing from
196 # "millisecond"), in which case that component is always 0.
197 _DURATION_DIVISORS: dict[str, dict[str, int]] = {
198 "day": {
199 "s": 86400,
200 "ms": 86_400_000,
201 "us": 86_400_000_000,
202 "ns": 86_400_000_000_000,
203 },
204 "hour": {
205 "s": 3600,
206 "ms": 3_600_000,
207 "us": 3_600_000_000,
208 "ns": 3_600_000_000_000,
209 },
210 "minute": {"s": 60, "ms": 60_000, "us": 60_000_000, "ns": 60_000_000_000},
211 "second": {"s": 1, "ms": 1_000, "us": 1_000_000, "ns": 1_000_000_000},
212 "millisecond": {"ms": 1, "us": 1_000, "ns": 1_000_000},
213 "microsecond": {"us": 1, "ns": 1_000},
214 "nanosecond": {"ns": 1},
215 }
217 def floor_div_int64(
218 arr: pa.Array | pa.ChunkedArray, divisor: int
219 ) -> pa.Array | pa.ChunkedArray:
220 """
221 Floor division for an int64 (chunked) array, handling negatives correctly.
223 PyArrow's ``pc.divide`` truncates toward zero, but we need floor division
224 (toward negative infinity) to match pandas behavior for negative durations.
225 Uses a fast path for all-non-negative arrays (the common case).
226 """
227 min_scalar = pc.min(arr)
228 if not min_scalar.is_valid or min_scalar.as_py() >= 0:
229 # Fast path: truncation equals floor when all values are non-negative
230 return pc.divide(arr, divisor)
231 # Truncated division (toward zero) and remainder with the dividend's sign
232 trunc_div = pc.divide(arr, divisor)
233 remainder = pc.subtract(arr, pc.multiply(trunc_div, divisor))
234 # floor(a/b) = trunc(a/b) - 1 when a < 0 and remainder != 0
235 needs_adjustment = pc.and_(pc.less(arr, 0), pc.not_equal(remainder, 0))
236 return pc.subtract(trunc_div, pc.if_else(needs_adjustment, 1, 0))
238 def mod_int(
239 arr: pa.Array | pa.ChunkedArray, modulo: int
240 ) -> pa.Array | pa.ChunkedArray:
241 """
242 Compute ``arr % modulo`` for a non-negative int (chunked) array.
244 PyArrow lacks a modulo kernel, so emulate it with divide/multiply/subtract.
245 TODO: revisit once Arrow has a modulo kernel.
246 """
247 quotient = pc.divide(arr, modulo)
248 return pc.subtract(arr, pc.multiply(quotient, modulo))
250 ARROW_ARITHMETIC_FUNCS = {
251 "add": pc.add_checked,
252 "radd": lambda x, y: pc.add_checked(y, x),
253 "sub": pc.subtract_checked,
254 "rsub": lambda x, y: pc.subtract_checked(y, x),
255 "mul": pc.multiply_checked,
256 "rmul": lambda x, y: pc.multiply_checked(y, x),
257 "truediv": lambda x, y: pc.divide(*cast_for_truediv(x, y)),
258 "rtruediv": lambda x, y: pc.divide(*cast_for_truediv(y, x)),
259 "floordiv": lambda x, y: floordiv_compat(x, y),
260 "rfloordiv": lambda x, y: floordiv_compat(y, x),
261 "mod": NotImplemented,
262 "rmod": NotImplemented,
263 "divmod": NotImplemented,
264 "rdivmod": NotImplemented,
265 "pow": pc.power_checked,
266 "rpow": lambda x, y: pc.power_checked(y, x),
267 }
269if TYPE_CHECKING:
270 from collections.abc import (
271 Callable,
272 Hashable,
273 Sequence,
274 )
276 from pandas._libs.missing import NAType
277 from pandas._typing import (
278 ArrayLike,
279 AxisInt,
280 Dtype,
281 DtypeObj,
282 FillnaOptions,
283 InterpolateOptions,
284 Iterator,
285 NpDtype,
286 NumpySorter,
287 NumpyValueArrayLike,
288 PositionalIndexer,
289 RankMethod,
290 RankNaOption,
291 Scalar,
292 SortKind,
293 TakeIndexer,
294 TimeAmbiguous,
295 TimeNonexistent,
296 npt,
297 )
299 from pandas.core.dtypes.dtypes import ExtensionDtype
301 from pandas import Series
302 from pandas.core.arrays.datetimes import DatetimeArray
303 from pandas.core.arrays.timedeltas import TimedeltaArray
306def to_pyarrow_type(
307 dtype: ArrowDtype | pa.DataType | Dtype | None,
308) -> pa.DataType | None:
309 """
310 Convert dtype to a pyarrow type instance.
311 """
312 if isinstance(dtype, ArrowDtype):
313 return dtype.pyarrow_dtype
314 elif isinstance(dtype, pa.DataType):
315 return dtype
316 elif isinstance(dtype, DatetimeTZDtype):
317 return pa.timestamp(dtype.unit, dtype.tz)
318 elif dtype:
319 try:
320 # Accepts python types too
321 # Doesn't handle all numpy types
322 return pa.from_numpy_dtype(dtype)
323 except pa.ArrowNotImplementedError:
324 pass
325 return None
328def _is_varbinary_type(pa_type: pa.DataType) -> bool:
329 """
330 Whether this is one of string, large_string, binary and large_binary.
332 pc.if_else misreads a non-zero offset for exactly these four, silently
333 truncating values (GH#64320, https://github.com/apache/arrow/issues/49410).
334 Other offset-carrying layouts such as list and map are unaffected.
335 """
336 return (
337 pa.types.is_string(pa_type)
338 or pa.types.is_large_string(pa_type)
339 or pa.types.is_binary(pa_type)
340 or pa.types.is_large_binary(pa_type)
341 )
344def _is_temporal_pa_type(pa_type: pa.DataType) -> bool | None:
345 """
346 Whether `pa_type` is temporal, or None if it does not settle that.
347 """
348 if pa.types.is_dictionary(pa_type):
349 return _is_temporal_pa_type(pa_type.value_type)
350 if pa.types.is_temporal(pa_type):
351 return True
352 if (
353 pa.types.is_integer(pa_type)
354 or pa.types.is_floating(pa_type)
355 or pa.types.is_boolean(pa_type)
356 or pa.types.is_decimal(pa_type)
357 ):
358 return False
359 # a string can parse to a timestamp, and nothing else here rules either way
360 return None
363def _is_temporal_dtype(dtype: DtypeObj) -> bool | None:
364 """
365 Whether `dtype` is temporal, or None if it does not settle that.
366 """
367 if isinstance(dtype, ArrowDtype):
368 return _is_temporal_pa_type(dtype.pyarrow_dtype)
369 if isinstance(dtype, CategoricalDtype):
370 return _is_temporal_dtype(dtype.categories.dtype)
371 if dtype.kind in "mM" or needs_i8_conversion(dtype):
372 # kind catches wrappers such as SparseDtype("M8[ns]"); needs_i8_conversion
373 # catches PeriodDtype, whose kind is "O"
374 return True
375 if dtype.kind in "iufbc":
376 return False
377 # Anything unrecognized has to land here: a wrong False would reject a
378 # valid assignment, see test_setitem_temporal_still_accepted
379 return None
382def _is_all_na(value) -> bool:
383 """
384 Whether every entry of `value` is NA.
385 """
386 if isinstance(value, (pa.Array, pa.ChunkedArray)):
387 # pyarrow keeps NaN distinct from null, so a float NaN built as an arrow
388 # value is a value here, unlike the numpy spelling below
389 return value.null_count == len(value)
390 return bool(np.asarray(isna(value)).all())
393def _is_temporal_value(value) -> bool | None:
394 """
395 Whether `value` is temporal, or None if its type does not settle that.
397 pyarrow converts a temporal value into an integer array without complaint
398 (`pa.array(dta, type=pa.int64())` succeeds), so a setitem has to reject the
399 reinterpretation before handing the value to `_box_pa` (GH#68419).
400 """
401 if is_scalar(value) and isna(value):
402 # NA of any flavor is settable into any dtype; a NaT scalar carries
403 # an M8/m8 dtype, so this has to come before the dtype lookup
404 return None
405 if type(value) in (int, float, bool):
406 # the common scalar setitem: skips infer_dtype_from_scalar (~10%), and
407 # settles an int too wide for any integer dtype, which it maps to object
408 return False
410 value = extract_array(value, extract_numpy=True)
411 if isinstance(value, pa.Scalar):
412 # a typed null scalar is just "assign NA"; isna() does not recognize
413 # pa.Scalar, so the guard above does not catch it
414 return None if not value.is_valid else _is_temporal_pa_type(value.type)
415 if isinstance(value, (pa.Array, pa.ChunkedArray)):
416 return _is_temporal_pa_type(value.type)
418 dtype = getattr(value, "dtype", None)
419 if dtype is None:
420 if is_list_like(value):
421 return None
422 if isinstance(value, (date, time)):
423 # infer_dtype_from_scalar maps both to object, which would leave the
424 # scalar disagreeing with the date32/time64 array forms
425 return True
426 if isinstance(value, Decimal):
427 # object again, this time disagreeing with the decimal128 array form
428 return False
429 dtype = infer_dtype_from_scalar(value)[0]
430 return _is_temporal_dtype(dtype)
433def _boxing_may_borrow_memory(pa_type: pa.DataType) -> bool:
434 """
435 Whether ``pa.array`` on this type can return a view on caller-owned memory.
437 Zero-copy over numpy/masked arrays for fixed-width layouts; character
438 layouts always repack, so copying those would cost a full copy of the
439 character data for no safety gain. Nested types may have a zero-copy child.
440 """
441 return not (
442 _is_varbinary_type(pa_type)
443 or pa.types.is_string_view(pa_type)
444 or pa.types.is_binary_view(pa_type)
445 )
448def _copy_pyarrow_buffers(
449 pa_array: pa.Array | pa.ChunkedArray,
450) -> pa.Array | pa.ChunkedArray:
451 """
452 Return an equal array that owns its buffers (GH#67990).
454 ``pa.concat_arrays`` reuses the ``dictionary`` child rather than copying it,
455 so dictionary types are rebuilt from copies of both halves.
456 """
457 if isinstance(pa_array, pa.ChunkedArray):
458 return pa.chunked_array(
459 [_copy_pyarrow_buffers(chunk) for chunk in pa_array.chunks],
460 type=pa_array.type,
461 )
462 if pa.types.is_dictionary(pa_array.type):
463 return pa.DictionaryArray.from_arrays(
464 _copy_pyarrow_buffers(pa_array.indices),
465 _copy_pyarrow_buffers(pa_array.dictionary),
466 ordered=pa_array.type.ordered,
467 )
468 return pa.concat_arrays([pa_array])
471@set_module("pandas.arrays")
472class ArrowExtensionArray(
473 OpsMixin,
474 ExtensionArrayNaResult,
475 ArrowStringArrayMixin,
476):
477 """
478 Pandas ExtensionArray backed by a PyArrow ChunkedArray.
480 .. warning::
482 ArrowExtensionArray is considered experimental. The implementation and
483 parts of the API may change without warning.
485 Parameters
486 ----------
487 values : pyarrow.Array or pyarrow.ChunkedArray
488 The input data to initialize the ArrowExtensionArray.
490 Attributes
491 ----------
492 None
494 Methods
495 -------
496 None
498 Returns
499 -------
500 ArrowExtensionArray
502 See Also
503 --------
504 array : Create a Pandas array with a specified dtype.
505 DataFrame.to_feather : Write a DataFrame to the binary Feather format.
506 read_feather : Load a feather-format object from the file path.
508 Notes
509 -----
510 Most methods are implemented using `pyarrow compute functions. <https://arrow.apache.org/docs/python/api/compute.html>`__
511 Some methods may either raise an exception or raise a ``PerformanceWarning`` if an
512 associated compute function is not available based on the installed version of PyArrow.
514 Please install the latest version of PyArrow to enable the best functionality and avoid
515 potential bugs in prior versions of PyArrow.
517 Examples
518 --------
519 Create an ArrowExtensionArray with :func:`pandas.array`:
521 >>> pd.array([1, 1, None], dtype="int64[pyarrow]")
522 <ArrowExtensionArray>
523 [1, 1, <NA>]
524 Length: 3, dtype: int64[pyarrow]
525 """ # noqa: E501 (http link too long)
527 _pa_array: pa.ChunkedArray
528 _dtype: ArrowDtype
529 # results from calls to methods decorated with cache_readonly get added here
530 _cache: dict[str, pa.ChunkedArray]
532 def __init__(self, values: pa.Array | pa.ChunkedArray) -> None:
533 if not HAS_PYARROW:
534 msg = (
535 f"pyarrow>={PYARROW_MIN_VERSION} is required for PyArrow "
536 "backed ArrowExtensionArray."
537 )
538 raise ImportError(msg)
539 if isinstance(values, pa.Array):
540 self._pa_array = pa.chunked_array([values])
541 elif isinstance(values, pa.ChunkedArray):
542 self._pa_array = values
543 else:
544 raise ValueError(
545 f"Unsupported type '{type(values)}' for ArrowExtensionArray"
546 )
547 self._dtype = ArrowDtype(self._pa_array.type)
548 self._cache = {}
550 @classmethod
551 def _from_sequence(
552 cls, scalars, *, dtype: Dtype | None = None, copy: bool = False
553 ) -> Self:
554 """
555 Construct a new ExtensionArray from a sequence of scalars.
556 """
557 pa_type = to_pyarrow_type(dtype)
558 pa_array = cls._box_pa_array(scalars, pa_type=pa_type, copy=copy)
559 arr = cls(pa_array)
560 return arr
562 @classmethod
563 def _from_sequence_of_strings(
564 cls, strings, *, dtype: ExtensionDtype, copy: bool = False
565 ) -> Self:
566 """
567 Construct a new ExtensionArray from a sequence of strings.
568 """
569 if isinstance(strings, cls):
570 strings = strings._pa_array
572 is_pa_array = isinstance(strings, (pa.Array, pa.ChunkedArray))
574 pa_type = to_pyarrow_type(dtype)
575 if (
576 pa_type is None
577 or pa.types.is_binary(pa_type)
578 or pa.types.is_string(pa_type)
579 or pa.types.is_large_string(pa_type)
580 ):
581 # pa_type is None: Let pa.array infer
582 # pa_type is string/binary: scalars already correct type
583 scalars = strings
584 elif pa.types.is_timestamp(pa_type):
585 from pandas.core.tools.datetimes import to_datetime
587 scalars = to_datetime(strings, errors="raise")
588 elif pa.types.is_date(pa_type):
589 from pandas.core.tools.datetimes import to_datetime
591 scalars = to_datetime(strings, errors="raise").date
592 scalars = pa.array(scalars, type=pa_type, mask=isna(strings))
593 elif pa.types.is_duration(pa_type):
594 from pandas.core.tools.timedeltas import to_timedelta
596 scalars = to_timedelta(strings, errors="raise")
598 if pa_type.unit != "ns":
599 # GH51175: test_from_sequence_of_strings_pa_array
600 # attempt to parse as int64 reflecting pyarrow's
601 # duration to string casting behavior
602 mask = isna(scalars)
603 if not is_pa_array:
604 strings = pa.array(strings, type=pa.string(), mask=mask)
605 if _is_varbinary_type(strings.type):
606 # GH#64320: replace_with_mask rather than if_else, which
607 # truncates values read through a non-zero offset. Only for
608 # the affected types; replace_with_mask has no kernel for
609 # several of the others and falls back to a lossy numpy path
610 strings = cls._replace_with_mask(
611 strings, mask, pa.scalar(None, type=strings.type)
612 )
613 else:
614 strings = pc.if_else(mask, None, strings)
615 try:
616 scalars = strings.cast(pa.int64())
617 except pa.ArrowInvalid:
618 pass
619 elif pa.types.is_time(pa_type):
620 from pandas.core.tools.times import to_time
622 # "coerce" to allow "null times" (None) to not raise
623 scalars = to_time(strings, errors="coerce")
624 elif pa.types.is_boolean(pa_type):
625 # pyarrow string->bool casting is case-insensitive:
626 # "true" or "1" -> True
627 # "false" or "0" -> False
628 # Note: BooleanArray was previously used to parse these strings
629 # and allows "1.0" and "0.0". Pyarrow casting does not support
630 # this, but we allow it here.
631 if is_pa_array:
632 scalars = strings
633 else:
634 scalars = pa.array(strings, type=pa.string(), mask=isna(strings))
635 scalars = pc.if_else(pc.equal(scalars, "1.0"), "1", scalars)
636 scalars = pc.if_else(pc.equal(scalars, "0.0"), "0", scalars)
637 scalars = scalars.cast(pa.bool_())
638 elif (
639 pa.types.is_integer(pa_type)
640 or pa.types.is_floating(pa_type)
641 or pa.types.is_decimal(pa_type)
642 ):
643 from pandas.core.tools.numeric import to_numeric
645 if is_pa_array:
646 # to_numeric only to reject spellings pyarrow's cast accepts
647 # but we do not, e.g. "0x1F"
648 to_numeric(strings, errors="raise")
649 scalars = strings.cast(pa_type)
650 else:
651 if pa.types.is_integer(pa_type):
652 # GH#56135: the default backend widens to float64 as soon
653 # as an NA is present, rounding integers above 2**53
654 scalars = extract_array(
655 to_numeric(
656 strings, errors="raise", dtype_backend="numpy_nullable"
657 ),
658 extract_numpy=True,
659 )
660 if not isinstance(scalars, BaseMaskedArray):
661 # the nullable backend returns object outside int64,
662 # and outside uint64 it also drops the NAs
663 scalars = to_numeric(strings, errors="raise")
664 else:
665 scalars = to_numeric(strings, errors="raise")
667 mask = isna(strings)
668 # GH#66834: to_numeric coerces "" to NaN instead of raising
669 if (isna(scalars) & ~mask).any():
670 raise ValueError(f"could not convert string to {pa_type}: ''")
671 if isinstance(scalars, BaseMaskedArray):
672 # the check above leaves scalars._mask a subset of mask
673 scalars = scalars._data
674 scalars = pa.array(scalars, mask=mask, type=pa_type)
676 else:
677 raise NotImplementedError(
678 f"Converting strings to {pa_type} is not implemented."
679 )
680 return cls._from_sequence(scalars, dtype=pa_type, copy=copy)
682 def _from_pyarrow_array(self, pa_array):
683 """
684 Construct from a pyarrow Array/ChunkedArray result of an operation.
686 Avoids full __init__ overhead by reusing the dtype when the pyarrow
687 type is unchanged.
688 """
689 assert isinstance(pa_array, (pa.Array, pa.ChunkedArray))
690 obj = type(self).__new__(type(self))
691 if isinstance(pa_array, pa.Array):
692 pa_array = pa.chunked_array([pa_array])
693 obj._pa_array = pa_array
694 pa_type = pa_array.type
695 obj._dtype = (
696 self._dtype if pa_type == self._dtype.pyarrow_dtype else ArrowDtype(pa_type)
697 )
698 obj._cache = {}
699 return obj
701 def _cast_pointwise_result(self, values) -> ArrayLike:
702 if len(values) == 0:
703 # Retain our dtype
704 return self[:0].copy()
706 try:
707 if self.dtype.kind in "iufc" and not is_nan_na():
708 values = np.asarray(values, dtype=object)
709 mask = is_pdna_or_none(values)
710 arr = pa.array(values, mask=mask)
711 elif hasattr(values, "__arrow_array__"):
712 arr = values.__arrow_array__()
713 elif hasattr(values, "_ndarray"):
714 # DatetimeArray, TimedeltaArray: pass underlying ndarray
715 # to avoid pyarrow calling Series.values (which is deprecated
716 # for tz-aware datetimes)
717 arr = pa.array(values._ndarray, from_pandas=True)
718 if hasattr(values.dtype, "tz") and values.dtype.tz is not None:
719 arr = arr.cast(
720 pa.timestamp(values.dtype.unit, tz=str(values.dtype.tz))
721 )
722 else:
723 arr = pa.array(values, from_pandas=True)
724 except (ValueError, TypeError, OverflowError):
725 # e.g. test_by_column_values_with_same_starting_value with nested
726 # values, one entry of which is an ArrowStringArray
727 # or test_agg_lambda_complex128_dtype_conversion for complex values
728 # or OverflowError for e.g. 2**70, GH#70233
729 values = np.asarray(values, dtype=object)
730 return lib.maybe_convert_objects(values, convert_non_numeric=True)
732 if pa.types.is_null(arr.type):
733 if lib.infer_dtype(values) == "decimal":
734 # GH#62522; the specific decimal precision here is arbitrary
735 arr = arr.cast(pa.decimal128(1))
736 if pa.types.is_duration(arr.type):
737 # workaround for https://github.com/apache/arrow/issues/40620
738 result = ArrowExtensionArray._from_sequence(values)
739 if pa.types.is_duration(self._pa_array.type):
740 result = result.astype(self.dtype) # type: ignore[assignment]
741 elif pa.types.is_timestamp(self._pa_array.type):
742 # Try to retain original unit
743 new_dtype = ArrowDtype(pa.duration(self._pa_array.type.unit))
744 try:
745 result = result.astype(new_dtype) # type: ignore[assignment]
746 except ValueError:
747 pass
748 elif pa.types.is_date64(self._pa_array.type):
749 # Try to match unit we get on non-pointwise op
750 dtype = ArrowDtype(pa.duration("ms"))
751 result = result.astype(dtype) # type: ignore[assignment]
752 elif pa.types.is_date(self._pa_array.type):
753 # Try to match unit we get on non-pointwise op
754 dtype = ArrowDtype(pa.duration("s"))
755 result = result.astype(dtype) # type: ignore[assignment]
756 return result
758 elif pa.types.is_timestamp(arr.type) and pa.types.is_timestamp(
759 self._pa_array.type
760 ):
761 if arr.type.tz == self._pa_array.type.tz:
762 try:
763 arr = arr.cast(self._pa_array.type)
764 except pa.lib.ArrowInvalid:
765 # GH#62523 a finer-resolution result can't be cast down to
766 # the original coarser unit without losing data; keep the
767 # inferred unit rather than raising
768 pass
770 elif pa.types.is_date(arr.type) and pa.types.is_date(self._pa_array.type):
771 arr = arr.cast(self._pa_array.type)
772 elif pa.types.is_time(arr.type) and pa.types.is_time(self._pa_array.type):
773 try:
774 arr = arr.cast(self._pa_array.type)
775 except pa.lib.ArrowInvalid:
776 # GH#62523 as above for the timestamp branch
777 pass
778 elif pa.types.is_decimal(arr.type) and pa.types.is_decimal(self._pa_array.type):
779 try:
780 arr = arr.cast(self._pa_array.type)
781 except pa.lib.ArrowInvalid:
782 # GH#62523 rescaling to the original decimal type would lose
783 # data; keep the inferred precision/scale rather than raising
784 pass
785 elif pa.types.is_integer(arr.type) and pa.types.is_integer(self._pa_array.type):
786 try:
787 arr = arr.cast(self._pa_array.type)
788 except pa.lib.ArrowInvalid:
789 # e.g. test_combine_add if we can't cast
790 pass
791 elif pa.types.is_floating(arr.type) and pa.types.is_floating(
792 self._pa_array.type
793 ):
794 try:
795 arr = arr.cast(self._pa_array.type)
796 except pa.lib.ArrowInvalid:
797 # e.g. test_combine_add if we can't cast
798 pass
800 return self._from_pyarrow_array(arr)
802 @classmethod
803 def _box_pa(
804 cls, value, pa_type: pa.DataType | None = None
805 ) -> pa.Array | pa.ChunkedArray | pa.Scalar:
806 """
807 Box value into a pyarrow Array, ChunkedArray or Scalar.
809 Parameters
810 ----------
811 value : any
812 pa_type : pa.DataType | None
814 Returns
815 -------
816 pa.Array or pa.ChunkedArray or pa.Scalar
817 """
818 if isinstance(value, pa.Scalar) or not is_list_like(value):
819 return cls._box_pa_scalar(value, pa_type)
820 return cls._box_pa_array(value, pa_type)
822 @classmethod
823 def _box_pa_scalar(cls, value, pa_type: pa.DataType | None = None) -> pa.Scalar:
824 """
825 Box value into a pyarrow Scalar.
827 Parameters
828 ----------
829 value : any
830 pa_type : pa.DataType | None
832 Returns
833 -------
834 pa.Scalar
835 """
836 if isinstance(value, pa.Scalar):
837 pa_scalar = value
838 elif isna(value) and not (lib.is_float(value) and not is_nan_na()):
839 pa_scalar = pa.scalar(None, type=pa_type)
840 else:
841 if (
842 pa_type is not None
843 and pa.types.is_timestamp(pa_type)
844 and isinstance(value, datetime)
845 and (value.tzinfo is None) != (pa_type.tz is None)
846 ):
847 # a Timestamp and a plain datetime both store their UTC epoch
848 # under pa_type, so a mismatch names a different instant
849 # (GH#69029). ArrowTypeError so the callers that normalize it
850 # report this like any other rejected value
851 raise pa.ArrowTypeError(
852 "Cannot mix tz-aware and tz-naive datetime-like values"
853 )
855 # Workaround https://github.com/apache/arrow/issues/37291. Only a
856 # duration or timestamp target reconciles units here; anything else
857 # is pa.scalar's to accept or reject (GH#68419)
858 if isinstance(value, Timedelta) and (
859 pa_type is None or pa.types.is_duration(pa_type)
860 ):
861 if pa_type is None:
862 pa_type = pa.duration(value.unit)
863 elif value.unit != pa_type.unit:
864 value = value.as_unit(pa_type.unit)
865 value = value._value
866 elif isinstance(value, Timestamp) and (
867 pa_type is None or pa.types.is_timestamp(pa_type)
868 ):
869 if pa_type is None:
870 pa_type = pa.timestamp(value.unit, tz=value.tz)
871 elif value.unit != pa_type.unit:
872 value = value.as_unit(pa_type.unit)
873 value = value._value
875 pa_scalar = pa.scalar(value, type=pa_type)
877 if pa_type is not None and pa_scalar.type != pa_type:
878 if (
879 pa_scalar.is_valid
880 and pa.types.is_timestamp(pa_scalar.type)
881 and pa.types.is_timestamp(pa_type)
882 and (pa_scalar.type.tz is None) != (pa_type.tz is None)
883 ):
884 # the cast keeps the UTC epoch, so it would name a different
885 # instant; a pa.Scalar value reaches the boundary only here
886 # (GH#69029)
887 raise pa.ArrowTypeError(
888 "Cannot mix tz-aware and tz-naive datetime-like values"
889 )
890 pa_scalar = pa_scalar.cast(pa_type)
892 return pa_scalar
894 @classmethod
895 def _box_pa_array(
896 cls, value, pa_type: pa.DataType | None = None, copy: bool = False
897 ) -> pa.Array | pa.ChunkedArray:
898 """
899 Box value into a pyarrow Array or ChunkedArray.
901 Parameters
902 ----------
903 value : Sequence
904 pa_type : pa.DataType | None
906 Returns
907 -------
908 pa.Array or pa.ChunkedArray
909 """
910 value = extract_array(value, extract_numpy=True)
911 if isinstance(value, cls):
912 pa_array = value._pa_array
913 elif isinstance(value, (pa.Array, pa.ChunkedArray)):
914 pa_array = value
915 elif isinstance(value, BaseMaskedArray):
916 # GH 52625
917 if copy:
918 value = value.copy()
919 pa_array = value.__arrow_array__()
921 elif hasattr(value, "__arrow_array__"):
922 # e.g. StringArray
923 if copy:
924 value = value.copy()
925 pa_array = value.__arrow_array__()
927 else:
928 if (
929 isinstance(value, np.ndarray)
930 and pa_type is not None
931 and (
932 pa.types.is_large_binary(pa_type)
933 or pa.types.is_large_string(pa_type)
934 )
935 ):
936 # See https://github.com/apache/arrow/issues/35289
937 value = np.asarray(value, dtype=object)
938 elif copy and is_array_like_deprecate_non_pandas(value):
939 # pa array should not get updated when numpy array is updated
940 value = value.copy()
942 if (
943 pa_type is not None
944 and pa.types.is_duration(pa_type)
945 and (not isinstance(value, np.ndarray) or value.dtype.kind not in "mi")
946 ):
947 # Workaround https://github.com/apache/arrow/issues/37291
948 from pandas.core.tools.timedeltas import to_timedelta
950 value = to_timedelta(value, unit=pa_type.unit).as_unit(pa_type.unit)
951 value = value.to_numpy()
953 if pa_type is not None and pa.types.is_timestamp(pa_type):
954 # Use DatetimeArray to exclude Decimal(NaN) (GH#61774) and
955 # ensure constructor treats tznaive the same as non-pyarrow
956 # dtypes (GH#61775)
957 from pandas.core.arrays.datetimes import (
958 DatetimeArray,
959 tz_to_dtype,
960 )
962 pass_dtype = tz_to_dtype(tz=pa_type.tz, unit=pa_type.unit)
963 value = extract_array(value, extract_numpy=True)
964 if isinstance(value, DatetimeArray):
965 dta = value
966 else:
967 dta = DatetimeArray._from_sequence(
968 value, copy=copy, dtype=pass_dtype
969 )
970 dta_mask = dta.isna()
971 if dta_mask.any():
972 # GH#61776 avoid __sub__ overflow. Work on a copy because
973 # dta can share memory with the caller's data
974 value_i8 = cast("npt.NDArray[np.int64]", dta.view("i8")).copy()
975 value_i8[dta_mask] = 0
976 dta = DatetimeArray._from_sequence(value_i8, dtype=dta.dtype)
977 pa_array = pa.array(dta._ndarray, type=pa_type, mask=dta_mask)
978 return pa_array
980 mask = None
981 if is_nan_na():
982 try:
983 arr_value = np.asarray(value)
984 if arr_value.ndim > 1:
985 # e.g. test_fixed_size_list we have list data. ndim > 1
986 # means there were no scalar (NA) entries.
987 mask = np.zeros(len(value), dtype=np.bool_)
988 else:
989 mask = isna(arr_value)
990 except ValueError:
991 # Ragged data that numpy raises on
992 arr_value = construct_1d_object_array_from_listlike(value)
993 mask = isna(arr_value)
994 elif (
995 getattr(value, "dtype", None) is None or value.dtype.kind not in "iumMf"
996 ):
997 arr_value = np.asarray(value, dtype=object)
998 # similar to isna(value) but exclude NaN, NaT, nat-like, nan-like
999 mask = is_pdna_or_none(arr_value)
1001 try:
1002 pa_array = pa.array(value, type=pa_type, mask=mask)
1003 except (pa.ArrowInvalid, pa.ArrowTypeError):
1004 # GH50430: let pyarrow infer type, then cast
1005 pa_array = pa.array(value, mask=mask)
1007 if pa_type is None and pa.types.is_duration(pa_array.type):
1008 # Workaround https://github.com/apache/arrow/issues/37291
1009 from pandas.core.tools.timedeltas import to_timedelta
1011 value = to_timedelta(value)
1012 value = value.to_numpy()
1013 pa_array = pa.array(value, type=pa_type)
1015 if pa.types.is_duration(pa_array.type) and pa_array.null_count > 0:
1016 # GH52843: upstream bug for duration types when originally
1017 # constructed with data containing numpy NaT.
1018 # https://github.com/apache/arrow/issues/35088
1019 arr = cls(pa_array)
1020 arr = arr.fillna(arr.dtype.na_value)
1021 pa_array = arr._pa_array
1023 if pa_type is not None and pa_array.type != pa_type:
1024 if pa.types.is_dictionary(pa_type):
1025 pa_array = pa_array.dictionary_encode()
1026 if pa_array.type != pa_type:
1027 pa_array = pa_array.cast(pa_type)
1028 else:
1029 try:
1030 pa_array = pa_array.cast(pa_type)
1031 except (pa.ArrowNotImplementedError, pa.ArrowTypeError):
1032 if pa.types.is_string(pa_array.type) or pa.types.is_large_string(
1033 pa_array.type
1034 ):
1035 # TODO: Move logic in _from_sequence_of_strings into
1036 # _box_pa_array
1037 dtype = ArrowDtype(pa_type)
1038 return cls._from_sequence_of_strings(
1039 value, dtype=dtype
1040 )._pa_array
1041 else:
1042 raise
1044 return pa_array
1046 def __getitem__(self, item: PositionalIndexer):
1047 """Select a subset of self.
1049 Parameters
1050 ----------
1051 item : int, slice, or ndarray
1052 * int: The position in 'self' to get.
1053 * slice: A slice object, where 'start', 'stop', and 'step' are
1054 integers or None
1055 * ndarray: A 1-d boolean NumPy ndarray the same length as 'self'
1057 Returns
1058 -------
1059 item : scalar or ExtensionArray
1061 Notes
1062 -----
1063 For scalar ``item``, return a scalar value suitable for the array's
1064 type. This should be an instance of ``self.dtype.type``.
1065 For slice ``key``, return an instance of ``ExtensionArray``, even
1066 if the slice is length 0 or 1.
1067 For a boolean mask, return an instance of ``ExtensionArray``, filtered
1068 to the values where ``item`` is True.
1069 """
1070 if isinstance(item, slice):
1071 # Fast path for slices: avoid check_array_indexer (no-op for
1072 # slices) and getitem_returns_view (always True for slices).
1073 # Arrow bug https://github.com/apache/arrow/issues/38768
1074 if item.start == item.stop:
1075 pass
1076 elif (
1077 item.stop is not None
1078 and item.stop < -len(self)
1079 and item.step is not None
1080 and item.step < 0
1081 ):
1082 item = slice(item.start, None, item.step)
1083 result = self._from_pyarrow_array(self._pa_array[item])
1084 result._readonly = self._readonly
1085 return result
1087 item = check_array_indexer(self, item)
1089 if isinstance(item, np.ndarray):
1090 if not len(item):
1091 # Removable once we migrate StringDtype[pyarrow] to ArrowDtype[string]
1092 if (
1093 isinstance(self._dtype, StringDtype)
1094 and self._dtype.storage == "pyarrow"
1095 ):
1096 # TODO(infer_string) should this be large_string?
1097 pa_dtype = pa.string()
1098 else:
1099 pa_dtype = self._dtype.pyarrow_dtype
1100 result = pa.chunked_array([], type=pa_dtype)
1101 return self._from_pyarrow_array(result)
1103 elif item.dtype.kind in "iu":
1104 return self.take(item)
1105 elif item.dtype.kind == "b":
1106 return self._from_pyarrow_array(self._pa_array.filter(item))
1107 else:
1108 raise IndexError(
1109 "Only integers, slices and integer or "
1110 "boolean arrays are valid indices."
1111 )
1112 elif isinstance(item, tuple):
1113 item = unpack_tuple_and_ellipses(item)
1115 if item is Ellipsis:
1116 # TODO: should be handled by pyarrow?
1117 item = slice(None)
1119 if is_scalar(item) and not is_integer(item):
1120 # e.g. "foo" or 2.5
1121 # exception message copied from numpy
1122 raise IndexError(
1123 r"only integers, slices (`:`), ellipsis (`...`), numpy.newaxis "
1124 r"(`None`) and integer or boolean arrays are valid indices"
1125 )
1126 # We are not an array indexer, so maybe e.g. a slice or integer
1127 # indexer. We dispatch to pyarrow.
1128 value = self._pa_array[item]
1129 if isinstance(value, pa.ChunkedArray):
1130 result = self._from_pyarrow_array(value)
1131 if getitem_returns_view(self, item):
1132 result._readonly = self._readonly
1133 return result
1134 else:
1135 pa_type = self._pa_array.type
1136 scalar = value.as_py()
1137 if scalar is None:
1138 return self._dtype.na_value
1139 elif pa.types.is_timestamp(pa_type) and pa_type.unit != "ns":
1140 # GH 53326
1141 return Timestamp(scalar).as_unit(pa_type.unit)
1142 elif pa.types.is_duration(pa_type) and pa_type.unit != "ns":
1143 # GH 53326
1144 return Timedelta(scalar).as_unit(pa_type.unit)
1145 else:
1146 return scalar
1148 def __iter__(self) -> Iterator[Any]:
1149 """
1150 Iterate over elements of the array.
1151 """
1152 na_value = self._dtype.na_value
1153 # GH 53326
1154 pa_type = self._pa_array.type
1155 box_timestamp = pa.types.is_timestamp(pa_type) and pa_type.unit != "ns"
1156 box_timedelta = pa.types.is_duration(pa_type) and pa_type.unit != "ns"
1157 for value in self._pa_array:
1158 val = value.as_py()
1159 if val is None:
1160 yield na_value
1161 elif box_timestamp:
1162 yield Timestamp(val).as_unit(pa_type.unit)
1163 elif box_timedelta:
1164 yield Timedelta(val).as_unit(pa_type.unit)
1165 else:
1166 yield val
1168 def __arrow_array__(self, type=None):
1169 """Convert myself to a pyarrow ChunkedArray."""
1170 return self._pa_array
1172 def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
1173 if (
1174 not is_nan_na()
1175 and type(self) is ArrowExtensionArray
1176 and (
1177 pa.types.is_floating(self._pa_array.type)
1178 or pa.types.is_integer(self._pa_array.type)
1179 )
1180 ):
1181 # this branch returns without reaching ExtensionArray.__array_ufunc__
1182 ops.disallow_datetimelike_logical_ufunc(ufunc, inputs)
1184 # GH#62506 - when distinguish_nan_and_na is True,
1185 # default_array_ufunc converts to numpy via np.asarray which
1186 # produces object dtype that most ufuncs can't handle.
1187 # Replicate ExtensionArray.__array_ufunc__ but convert with
1188 # na_value=np.nan (float, not object) in the default fallback,
1189 # then re-mask NA positions in the result.
1190 if any(
1191 isinstance(other, (ABCSeries, ABCIndex, ABCDataFrame))
1192 for other in inputs
1193 ):
1194 return NotImplemented
1196 result = arraylike.maybe_dispatch_ufunc_to_dunder_op(
1197 self, ufunc, method, *inputs, **kwargs
1198 )
1199 if result is not NotImplemented:
1200 return result
1202 if "out" in kwargs:
1203 return arraylike.dispatch_ufunc_with_out(
1204 self, ufunc, method, *inputs, **kwargs
1205 )
1207 if method == "reduce":
1208 result = arraylike.dispatch_reduction_ufunc(
1209 self, ufunc, method, *inputs, **kwargs
1210 )
1211 if result is not NotImplemented:
1212 return result
1214 # Default: convert to numpy with NaN for NA instead of
1215 # object dtype, run the ufunc, then re-mask.
1216 mask = self.isna()
1217 new_inputs = [
1218 self.to_numpy(na_value=np.nan) if x is self else x for x in inputs
1219 ]
1220 result = getattr(ufunc, method)(*new_inputs, **kwargs)
1222 remask = functools.partial(pa.array, mask=mask, from_pandas=False)
1223 if isinstance(result, tuple):
1224 return tuple(type(self)(remask(res)) for res in result)
1225 return type(self)(remask(result))
1227 # Need to wrap np.array results GH#62800
1228 result = super().__array_ufunc__(ufunc, method, *inputs, **kwargs)
1229 if type(self) is ArrowExtensionArray:
1230 # Exclude ArrowStringArray
1231 return type(self)._from_sequence(result)
1232 return result
1234 def __array__(
1235 self, dtype: NpDtype | None = None, copy: bool | None = None
1236 ) -> np.ndarray:
1237 """Correctly construct numpy arrays when passed to `np.asarray()`."""
1238 if copy is False:
1239 # TODO: By using `zero_copy_only` it may be possible to implement this
1240 raise ValueError(
1241 "Unable to avoid copy while creating an array as requested."
1242 )
1243 elif copy is None:
1244 # `to_numpy(copy=False)` has the meaning of NumPy `copy=None`.
1245 copy = False
1247 return self.to_numpy(dtype=dtype, copy=copy)
1249 def __invert__(self) -> Self:
1250 # This is a bit wise op for integer types
1251 if pa.types.is_integer(self._pa_array.type):
1252 return self._from_pyarrow_array(pc.bit_wise_not(self._pa_array))
1253 elif pa.types.is_string(self._pa_array.type) or pa.types.is_large_string(
1254 self._pa_array.type
1255 ):
1256 # Raise TypeError instead of pa.ArrowNotImplementedError
1257 raise TypeError("__invert__ is not supported for string dtypes")
1258 else:
1259 return self._from_pyarrow_array(pc.invert(self._pa_array))
1261 def __neg__(self) -> Self:
1262 try:
1263 return self._from_pyarrow_array(pc.negate_checked(self._pa_array))
1264 except pa.ArrowNotImplementedError as err:
1265 raise TypeError(
1266 f"unary '-' not supported for dtype '{self.dtype}'"
1267 ) from err
1269 def __pos__(self) -> Self:
1270 return self._from_pyarrow_array(self._pa_array)
1272 def __abs__(self) -> Self:
1273 return self._from_pyarrow_array(pc.abs_checked(self._pa_array))
1275 # GH 42600: __getstate__/__setstate__ not necessary once
1276 # https://issues.apache.org/jira/browse/ARROW-10739 is addressed
1277 def __getstate__(self):
1278 state = self.__dict__.copy()
1279 state["_pa_array"] = self._pa_array.combine_chunks()
1280 # cached properties can be recomputed; don't bloat the pickle
1281 state["_cache"] = {}
1282 return state
1284 def __setstate__(self, state) -> None:
1285 if "_data" in state:
1286 data = state.pop("_data")
1287 else:
1288 data = state["_pa_array"]
1289 state["_pa_array"] = pa.chunked_array(data)
1290 # pickles created before _cache existed
1291 state.setdefault("_cache", {})
1292 self.__dict__.update(state)
1294 def _cmp_method(self, other, op) -> ArrowExtensionArray:
1295 pc_func = ARROW_CMP_FUNCS[op.__name__]
1296 ltype = self._pa_array.type
1298 if isinstance(other, range):
1299 # GH#63429 our callers defer this to the EA for EA-backed values
1300 ops.maybe_warn_listlike(other)
1302 if isinstance(other, (ExtensionArray, np.ndarray, list, range)):
1303 ops.raise_if_2d(other)
1304 if is_list_like(other) and len(self) != len(other):
1305 # is_list_like excludes 0-dim ndarrays, which are scalars here
1306 raise ValueError("Lengths must match to compare")
1307 try:
1308 boxed = self._box_pa(other)
1309 except pa.lib.ArrowInvalid:
1310 # e.g. GH#60228 [1, "b"] we have to operate pointwise
1311 res_values = [op(x, y) for x, y in zip(self, other, strict=True)]
1312 result = pa.array(res_values, type=pa.bool_(), from_pandas=True)
1313 else:
1314 rtype = boxed.type
1315 if (pa.types.is_timestamp(ltype) and pa.types.is_date(rtype)) or (
1316 pa.types.is_timestamp(rtype) and pa.types.is_date(ltype)
1317 ):
1318 # GH#62157 match non-pyarrow behavior
1319 result = ops.invalid_comparison(self, other, op)
1320 result = pa.array(result, type=pa.bool_())
1321 else:
1322 try:
1323 result = pc_func(self._pa_array, boxed)
1324 except pa.ArrowNotImplementedError:
1325 result = ops.invalid_comparison(self, other, op)
1326 result = pa.array(result, type=pa.bool_())
1328 elif is_scalar(other):
1329 if (isinstance(other, datetime) and pa.types.is_date(ltype)) or (
1330 type(other) is date and pa.types.is_timestamp(ltype)
1331 ):
1332 # GH#62157 match non-pyarrow behavior
1333 result = ops.invalid_comparison(self, other, op)
1334 result = pa.array(result, type=pa.bool_())
1335 else:
1336 try:
1337 result = pc_func(self._pa_array, self._box_pa(other))
1338 except (pa.lib.ArrowNotImplementedError, pa.lib.ArrowInvalid):
1339 mask = isna(self) | isna(other)
1340 valid = ~mask
1341 result = np.zeros(len(self), dtype="bool")
1342 np_array = np.array(self)
1343 try:
1344 result[valid] = op(np_array[valid], other)
1345 except TypeError:
1346 result = ops.invalid_comparison(self, other, op)
1347 result = pa.array(result, type=pa.bool_())
1348 result = pc.if_else(valid, result, None)
1349 else:
1350 raise NotImplementedError(
1351 f"{op.__name__} not implemented for {type(other)}"
1352 )
1353 return ArrowExtensionArray(result)
1355 def _op_method_error_message(self, other, op) -> str:
1356 if hasattr(other, "dtype"):
1357 other_type = f"dtype '{other.dtype}'"
1358 else:
1359 other_type = f"object of type {type(other)}"
1360 return (
1361 f"operation '{op.__name__}' not supported for "
1362 f"dtype '{self.dtype}' with {other_type}"
1363 )
1365 def _evaluate_op_method(self, other, op, arrow_funcs) -> Self:
1366 if (
1367 is_list_like(other)
1368 and not isinstance(other, (np.ndarray, ExtensionArray, list))
1369 and not ops.has_castable_attr(other)
1370 ):
1371 warnings.warn(
1372 f"Operation with {type(other).__name__} is deprecated. "
1373 "In a future version these will be treated as scalar-like. "
1374 "To retain the old behavior, explicitly wrap in a Series "
1375 "instead.",
1376 Pandas4Warning,
1377 stacklevel=find_stack_level(),
1378 )
1380 pa_type = self._pa_array.type
1381 other_original = other
1382 ops.raise_if_2d(other)
1383 other = self._box_pa(other)
1385 if (
1386 pa.types.is_string(pa_type)
1387 or pa.types.is_large_string(pa_type)
1388 or pa.types.is_binary(pa_type)
1389 ):
1390 if op in [operator.add, roperator.radd]:
1391 # binary_join_element_wise does not support mixed types, but we
1392 # want to allow addition between string and large_string types
1393 self_array = self._pa_array
1394 if pa.types.is_string(pa_type) and pa.types.is_large_string(other.type):
1395 self_array = self._pa_array.cast(pa.large_string())
1396 elif pa.types.is_large_string(pa_type) and pa.types.is_string(
1397 other.type
1398 ):
1399 other = other.cast(pa.large_string())
1401 sep = pa.scalar("", type=self_array.type)
1402 if isinstance(other, pa.Scalar):
1403 if pc.is_null(other).as_py():
1404 other = other.cast(self_array.type)
1405 elif len(other) == 0 and pa.types.is_null(other.type):
1406 # GH#40624 a length-zero object-dtype other boxes to the
1407 # pyarrow null type, which has no join kernel
1408 other = other.cast(self_array.type)
1409 try:
1410 if op is operator.add:
1411 result = pc.binary_join_element_wise(self_array, other, sep)
1412 elif op is roperator.radd:
1413 result = pc.binary_join_element_wise(other, self_array, sep)
1414 except pa.ArrowNotImplementedError as err:
1415 raise TypeError(
1416 self._op_method_error_message(other_original, op)
1417 ) from err
1418 return self._from_pyarrow_array(result)
1419 elif op in [operator.mul, roperator.rmul]:
1420 binary = self._pa_array
1421 integral = other
1422 if not pa.types.is_integer(integral.type):
1423 raise TypeError("Can only string multiply by an integer.")
1424 pa_integral = pc.if_else(pc.less(integral, 0), 0, integral)
1425 result = pc.binary_repeat(binary, pa_integral)
1426 return self._from_pyarrow_array(result)
1427 elif (
1428 pa.types.is_string(other.type)
1429 or pa.types.is_binary(other.type)
1430 or pa.types.is_large_string(other.type)
1431 ) and op in [operator.mul, roperator.rmul]:
1432 binary = other
1433 integral = self._pa_array
1434 if not pa.types.is_integer(integral.type):
1435 raise TypeError("Can only string multiply by an integer.")
1436 pa_integral = pc.if_else(pc.less(integral, 0), 0, integral)
1437 result = pc.binary_repeat(binary, pa_integral)
1438 return self._from_pyarrow_array(result)
1439 if (
1440 isinstance(other, pa.Scalar)
1441 and pc.is_null(other).as_py()
1442 and op.__name__ in ARROW_LOGICAL_FUNCS
1443 ):
1444 # pyarrow kleene ops require null to be typed
1445 other = other.cast(pa_type)
1447 pc_func = arrow_funcs[op.__name__]
1448 if pc_func is NotImplemented:
1449 if pa.types.is_string(pa_type) or pa.types.is_large_string(pa_type):
1450 raise TypeError(self._op_method_error_message(other_original, op))
1451 raise NotImplementedError(f"{op.__name__} not implemented.")
1453 try:
1454 result = pc_func(self._pa_array, other)
1455 except pa.ArrowNotImplementedError as err:
1456 raise TypeError(self._op_method_error_message(other_original, op)) from err
1457 return self._from_pyarrow_array(result)
1459 def _logical_method(self, other, op) -> Self:
1460 # checked here too because the GH#60234 arm below returns without
1461 # reaching _evaluate_op_method
1462 ops.raise_if_2d(other)
1464 # For integer types `^`, `|`, `&` are bitwise operators and return
1465 # integer types. Otherwise these are boolean ops.
1466 if pa.types.is_integer(self._pa_array.type):
1467 return self._evaluate_op_method(other, op, ARROW_BIT_WISE_FUNCS)
1468 elif (
1469 (
1470 pa.types.is_string(self._pa_array.type)
1471 or pa.types.is_large_string(self._pa_array.type)
1472 )
1473 and op in (roperator.ror_, roperator.rand_, roperator.rxor)
1474 and isinstance(other, np.ndarray)
1475 and other.dtype == bool
1476 ):
1477 # GH#60234 backward compatibility for the move to StringDtype in 3.0
1478 op_name = op.__name__[1:].strip("_")
1479 warnings.warn(
1480 f"'{op_name}' operations between boolean dtype and {self.dtype} are "
1481 "deprecated and will raise in a future version. Explicitly "
1482 "cast the strings to a boolean dtype before operating instead.",
1483 Pandas4Warning,
1484 stacklevel=find_stack_level(),
1485 )
1486 return op(other, self.astype(bool))
1487 else:
1488 return self._evaluate_op_method(other, op, ARROW_LOGICAL_FUNCS)
1490 def _str_arith_method_object_fallback(
1491 self, other, op
1492 ) -> Self | npt.NDArray[np.object_]:
1493 mask = isna(self) | isna(other)
1494 valid = ~mask
1496 if is_list_like(other):
1497 if len(other) != len(self):
1498 raise ValueError(
1499 f"Lengths of operands do not match: {len(self)} != {len(other)}"
1500 )
1501 if not is_array_like_deprecate_non_pandas(other):
1502 other = np.asarray(other)
1503 other = other[valid]
1505 result = np.empty(len(self), dtype=object)
1506 result[mask] = self.dtype.na_value
1507 result[valid] = op(np.asarray(self, dtype=object)[valid], other)
1509 if not lib.is_string_array(result, skipna=True):
1510 return result
1511 return type(self)._from_sequence(result, dtype=self.dtype)
1513 def _arith_method(self, other, op) -> Self | npt.NDArray[np.object_]:
1514 if isinstance(other, BaseOffset) and pa.types.is_date(self._pa_array.type):
1515 # Cast date32/date64 → timestamp, apply offset via DatetimeArray, cast back
1516 ts_array = type(self)(self._pa_array.cast(pa.timestamp("us")))
1517 dt_array = ts_array._to_datetimearray()
1519 shifted = op(dt_array, other)
1520 check = shifted[~shifted.isna()] if shifted._hasna else shifted
1521 if not check.is_normalized:
1522 raise TypeError(
1523 f"DateOffset {other} is intra-day and cannot be "
1524 f"applied to date32/date64 arrays"
1525 )
1526 result_pa = pa.array(shifted._ndarray, from_pandas=True).cast(
1527 self._pa_array.type
1528 )
1529 return self._from_pyarrow_array(result_pa)
1531 result: Self | npt.NDArray[np.object_]
1532 if pa.types.is_string(self._pa_array.type) or pa.types.is_large_string(
1533 self._pa_array.type
1534 ):
1535 try:
1536 result = self._evaluate_op_method(other, op, ARROW_ARITHMETIC_FUNCS)
1537 except (pa.ArrowInvalid, pa.ArrowTypeError):
1538 result = self._str_arith_method_object_fallback(other, op)
1539 else:
1540 result = self._evaluate_op_method(other, op, ARROW_ARITHMETIC_FUNCS)
1541 if isinstance(result, np.ndarray):
1542 return result
1543 if is_nan_na() and result.dtype.kind == "f":
1544 parr = result._pa_array
1545 mask = pc.is_nan(parr).fill_null(False).to_numpy()
1546 arr = pc.replace_with_mask(parr, mask, pa.scalar(None, type=parr.type))
1547 result = type(self)(arr)
1548 return result
1550 def equals(self, other) -> bool:
1551 if not isinstance(other, ArrowExtensionArray):
1552 return False
1553 # I'm told that pyarrow makes __eq__ behave like pandas' equals;
1554 # TODO: is this documented somewhere?
1555 return self._pa_array == other._pa_array
1557 @property
1558 def dtype(self) -> ArrowDtype:
1559 """
1560 An instance of 'ExtensionDtype'.
1561 """
1562 return self._dtype
1564 @property
1565 def nbytes(self) -> int:
1566 """
1567 The number of bytes needed to store this object in memory.
1568 """
1569 return self._pa_array.nbytes
1571 def __len__(self) -> int:
1572 """
1573 Length of this array.
1575 Returns
1576 -------
1577 length : int
1578 """
1579 return len(self._pa_array)
1581 def __contains__(self, key) -> bool:
1582 # https://github.com/pandas-dev/pandas/pull/51307#issuecomment-1426372604
1583 if isna(key) and key is not self.dtype.na_value:
1584 if lib.is_float(key) and is_nan_na():
1585 return self.dtype.na_value in self
1586 elif self.dtype.kind == "f" and lib.is_float(key):
1587 # Check specifically for NaN
1588 return pc.any(pc.is_nan(self._pa_array)).as_py()
1590 # e.g. date or timestamp types we do not allow None here to match pd.NA
1591 return False
1592 # TODO: maybe complex? object?
1594 return bool(super().__contains__(key))
1596 @property
1597 def _hasna(self) -> bool:
1598 return self._pa_array.null_count > 0
1600 def isna(self) -> npt.NDArray[np.bool_]:
1601 """
1602 Boolean NumPy array indicating if each value is missing.
1604 This should return a 1-D array the same length as 'self'.
1605 """
1606 # GH51630: fast paths
1607 null_count = self._pa_array.null_count
1608 if null_count == 0:
1609 return np.zeros(len(self), dtype=np.bool_)
1610 elif null_count == len(self):
1611 return np.ones(len(self), dtype=np.bool_)
1613 return self._pa_array.is_null().to_numpy()
1615 @overload
1616 def any(self, *, skipna: Literal[True] = ..., **kwargs) -> bool: ...
1618 @overload
1619 def any(self, *, skipna: bool, **kwargs) -> bool | NAType: ...
1621 def any(self, *, skipna: bool = True, **kwargs) -> bool | NAType:
1622 """
1623 Return whether any element is truthy.
1625 Returns False unless there is at least one element that is truthy.
1626 By default, NAs are skipped. If ``skipna=False`` is specified and
1627 missing values are present, similar :ref:`Kleene logic <boolean.kleene>`
1628 is used as for logical operations.
1630 Parameters
1631 ----------
1632 skipna : bool, default True
1633 Exclude NA values. If the entire array is NA and `skipna` is
1634 True, then the result will be False, as for an empty array.
1635 If `skipna` is False, the result will still be True if there is
1636 at least one element that is truthy, otherwise NA will be returned
1637 if there are NA's present.
1639 Returns
1640 -------
1641 bool or :attr:`pandas.NA`
1643 See Also
1644 --------
1645 ArrowExtensionArray.all : Return whether all elements are truthy.
1647 Examples
1648 --------
1649 The result indicates whether any element is truthy (and by default
1650 skips NAs):
1652 >>> pd.array([True, False, True], dtype="boolean[pyarrow]").any()
1653 True
1654 >>> pd.array([True, False, pd.NA], dtype="boolean[pyarrow]").any()
1655 True
1656 >>> pd.array([False, False, pd.NA], dtype="boolean[pyarrow]").any()
1657 False
1658 >>> pd.array([], dtype="boolean[pyarrow]").any()
1659 False
1660 >>> pd.array([pd.NA], dtype="boolean[pyarrow]").any()
1661 False
1662 >>> pd.array([pd.NA], dtype="float64[pyarrow]").any()
1663 False
1665 With ``skipna=False``, the result can be NA if this is logically
1666 required (whether ``pd.NA`` is True or False influences the result):
1668 >>> pd.array([True, False, pd.NA], dtype="boolean[pyarrow]").any(skipna=False)
1669 True
1670 >>> pd.array([1, 0, pd.NA], dtype="boolean[pyarrow]").any(skipna=False)
1671 True
1672 >>> pd.array([False, False, pd.NA], dtype="boolean[pyarrow]").any(skipna=False)
1673 <NA>
1674 >>> pd.array([0, 0, pd.NA], dtype="boolean[pyarrow]").any(skipna=False)
1675 <NA>
1676 """
1677 return self._reduce("any", skipna=skipna, **kwargs)
1679 @overload
1680 def all(self, *, skipna: Literal[True] = ..., **kwargs) -> bool: ...
1682 @overload
1683 def all(self, *, skipna: bool, **kwargs) -> bool | NAType: ...
1685 def all(self, *, skipna: bool = True, **kwargs) -> bool | NAType:
1686 """
1687 Return whether all elements are truthy.
1689 Returns True unless there is at least one element that is falsey.
1690 By default, NAs are skipped. If ``skipna=False`` is specified and
1691 missing values are present, similar :ref:`Kleene logic <boolean.kleene>`
1692 is used as for logical operations.
1694 Parameters
1695 ----------
1696 skipna : bool, default True
1697 Exclude NA values. If the entire array is NA and `skipna` is
1698 True, then the result will be True, as for an empty array.
1699 If `skipna` is False, the result will still be False if there is
1700 at least one element that is falsey, otherwise NA will be returned
1701 if there are NA's present.
1703 Returns
1704 -------
1705 bool or :attr:`pandas.NA`
1707 See Also
1708 --------
1709 ArrowExtensionArray.any : Return whether any element is truthy.
1711 Examples
1712 --------
1713 The result indicates whether all elements are truthy (and by default
1714 skips NAs):
1716 >>> pd.array([True, True, pd.NA], dtype="boolean[pyarrow]").all()
1717 True
1718 >>> pd.array([1, 1, pd.NA], dtype="boolean[pyarrow]").all()
1719 True
1720 >>> pd.array([True, False, pd.NA], dtype="boolean[pyarrow]").all()
1721 False
1722 >>> pd.array([], dtype="boolean[pyarrow]").all()
1723 True
1724 >>> pd.array([pd.NA], dtype="boolean[pyarrow]").all()
1725 True
1726 >>> pd.array([pd.NA], dtype="float64[pyarrow]").all()
1727 True
1729 With ``skipna=False``, the result can be NA if this is logically
1730 required (whether ``pd.NA`` is True or False influences the result):
1732 >>> pd.array([True, True, pd.NA], dtype="boolean[pyarrow]").all(skipna=False)
1733 <NA>
1734 >>> pd.array([1, 1, pd.NA], dtype="boolean[pyarrow]").all(skipna=False)
1735 <NA>
1736 >>> pd.array([True, False, pd.NA], dtype="boolean[pyarrow]").all(skipna=False)
1737 False
1738 >>> pd.array([1, 0, pd.NA], dtype="boolean[pyarrow]").all(skipna=False)
1739 False
1740 """
1741 return self._reduce("all", skipna=skipna, **kwargs)
1743 def argsort(
1744 self,
1745 *,
1746 ascending: bool = True,
1747 kind: SortKind = "quicksort",
1748 na_position: str = "last",
1749 **kwargs,
1750 ) -> np.ndarray:
1751 order = "ascending" if ascending else "descending"
1752 null_placement = {"last": "at_end", "first": "at_start"}.get(na_position, None)
1753 if null_placement is None:
1754 raise ValueError(f"invalid na_position: {na_position}")
1756 result = pc.array_sort_indices(
1757 self._pa_array, order=order, null_placement=null_placement
1758 )
1759 np_result = result.to_numpy()
1760 return np_result.astype(np.intp, copy=False)
1762 def sort(
1763 self,
1764 *,
1765 ascending: bool = True,
1766 kind: SortKind = "quicksort",
1767 na_position: str = "last",
1768 ) -> None:
1769 # This override replaces self._pa_array directly, bypassing __setitem__,
1770 # so enforce the read-only guard here to stay consistent with it and
1771 # with the base ExtensionArray.sort.
1772 if self._readonly:
1773 raise ValueError("Cannot modify read-only array")
1774 sort_indices = self.argsort(
1775 ascending=ascending, kind=kind, na_position=na_position
1776 )
1777 sorted_array = self.take(sort_indices)
1778 self._pa_array = sorted_array._pa_array
1779 # Invalidate any cache_readonly properties that depend on _pa_array
1780 self._cache.clear()
1782 def _argmin_max(self, skipna: bool, method: str) -> int:
1783 if self._pa_array.length() in (0, self._pa_array.null_count) or (
1784 self._hasna and not skipna
1785 ):
1786 # For empty or all null, pyarrow returns -1 but pandas expects TypeError
1787 # For skipna=False and data w/ null, pandas expects NotImplementedError
1788 # let ExtensionArray.arg{max|min} raise
1789 return getattr(super(), f"arg{method}")(skipna=skipna)
1791 data = self._pa_array
1792 if pa.types.is_duration(data.type):
1793 data = data.cast(pa.int64())
1795 value = getattr(pc, method)(data, skip_nulls=skipna)
1796 return pc.index(data, value).as_py()
1798 def argmin(self, skipna: bool = True) -> int:
1799 return self._argmin_max(skipna, "min")
1801 def argmax(self, skipna: bool = True) -> int:
1802 return self._argmin_max(skipna, "max")
1804 def copy(self) -> Self:
1805 """
1806 Return a shallow copy of the array.
1808 Underlying ChunkedArray is immutable, so a deep copy is unnecessary.
1810 Returns
1811 -------
1812 type(self)
1813 """
1814 return self._from_pyarrow_array(self._pa_array)
1816 @overload
1817 def view(self, dtype: None = ...) -> Self: ...
1819 @overload
1820 def view(self, dtype: Dtype | None = ...) -> ArrayLike: ...
1822 def view(self, dtype: Dtype | None = None) -> ArrayLike:
1823 # Override since the ExtensionArray.view create a new pa_array.
1824 if dtype is not None:
1825 return super().view(dtype)
1826 result = type(self).__new__(type(self))
1827 result._pa_array = self._pa_array
1828 result._dtype = self._dtype
1829 result._readonly = self._readonly
1830 result._cache = {}
1831 return result
1833 def count(self) -> int:
1834 return len(self) - self._pa_array.null_count
1836 def dropna(self) -> Self:
1837 """
1838 Return ArrowExtensionArray without NA values.
1840 Returns
1841 -------
1842 ArrowExtensionArray
1843 """
1844 return self._from_pyarrow_array(pc.drop_null(self._pa_array))
1846 def _pad_or_backfill(
1847 self,
1848 *,
1849 method: FillnaOptions,
1850 limit: int | None = None,
1851 limit_area: Literal["inside", "outside"] | None = None,
1852 copy: bool = True,
1853 ) -> Self:
1854 if not self._hasna:
1855 if not copy:
1856 return self
1857 return self._from_pyarrow_array(self._pa_array)
1859 if limit is None and limit_area is None:
1860 method = missing.clean_fill_method(method)
1861 try:
1862 if method == "pad":
1863 return self._from_pyarrow_array(
1864 pc.fill_null_forward(self._pa_array)
1865 )
1866 elif method == "backfill":
1867 return self._from_pyarrow_array(
1868 pc.fill_null_backward(self._pa_array)
1869 )
1870 except pa.ArrowNotImplementedError:
1871 # ArrowNotImplementedError: Function 'coalesce' has no kernel
1872 # matching input types (duration[ns], duration[ns])
1873 # TODO: remove try/except wrapper if/when pyarrow implements
1874 # a kernel for duration types.
1875 pass
1877 return super()._pad_or_backfill(
1878 method=method, limit=limit, limit_area=limit_area, copy=copy
1879 )
1881 def fillna(
1882 self,
1883 value: object | ArrayLike,
1884 limit: int | None = None,
1885 copy: bool = True,
1886 ) -> Self:
1887 """
1888 Fill NA/NaN values using the specified method.
1890 Parameters
1891 ----------
1892 value : scalar, array-like
1893 If a scalar value is passed it is used to fill all missing values.
1894 Alternatively, an array-like "value" can be given. It's expected
1895 that the array-like have the same length as 'self'.
1896 limit : int, default None
1897 The maximum number of entries where NA values will be filled.
1898 copy : bool, default True
1899 Whether to make a copy of the data before filling. This parameter
1900 is inherited from pandas.ExtensionArray and should be left at its
1901 default value (Arrow arrays cannot be altered in place).
1903 Returns
1904 -------
1905 ArrowExtensionArray
1906 With NA/NaN filled.
1908 See Also
1909 --------
1910 api.extensions.ExtensionArray.dropna : Return ExtensionArray without
1911 NA values.
1912 api.extensions.ExtensionArray.isna : A 1-D array indicating if
1913 each value is missing.
1915 Examples
1916 --------
1917 >>> arr = pd.array(
1918 ... [np.nan, np.nan, 2, 3, np.nan, np.nan], dtype="int64[pyarrow]"
1919 ... )
1920 >>> arr.fillna(0)
1921 <ArrowExtensionArray>
1922 [0, 0, 2, 3, 0, 0]
1923 Length: 6, dtype: int64[pyarrow]
1924 """
1925 if not self._hasna:
1926 return self.copy()
1928 if isinstance(value, dict):
1929 raise TypeError(
1930 "ExtensionArray.fillna does not support filling with a dict. "
1931 "Use Series.fillna instead."
1932 )
1934 if limit is not None:
1935 return super().fillna(value=value, limit=limit, copy=copy)
1937 if isinstance(value, (np.ndarray, ExtensionArray)):
1938 # Similar to check_value_size, but we do not mask here since we may
1939 # end up passing it to the super() method.
1940 if len(value) != len(self):
1941 raise ValueError(
1942 f"Length of 'value' does not match. Got ({len(value)}) "
1943 f" expected {len(self)}"
1944 )
1946 fill_value = self._validate_setitem_value(value)
1948 try:
1949 return self._from_pyarrow_array(
1950 _safe_fill_null(self._pa_array, fill_value=fill_value)
1951 )
1952 except pa.ArrowNotImplementedError:
1953 # ArrowNotImplementedError: Function 'coalesce' has no kernel
1954 # matching input types (duration[ns], duration[ns])
1955 # TODO: remove try/except wrapper if/when pyarrow implements
1956 # a kernel for duration types.
1957 pass
1959 return super().fillna(value=value, limit=limit, copy=copy)
1961 def isin(self, values: ArrayLike) -> npt.NDArray[np.bool_]:
1962 # short-circuit to return all False array.
1963 if not len(values):
1964 return np.zeros(len(self), dtype=bool)
1966 value_set = self._box_pa(values)
1967 result = pc.is_in(self._pa_array, value_set=value_set)
1968 # pyarrow 2.0.0 returned nulls, so we explicitly specify dtype to convert nulls
1969 # to False
1970 return np.array(result, dtype=np.bool_)
1972 def _hash_pandas_object(
1973 self, *, encoding: str, hash_key: str, categorize: bool
1974 ) -> npt.NDArray[np.uint64]:
1975 """
1976 Hook for hash_pandas_object.
1978 For string/binary types with categorize=True, uses pyarrow's
1979 dictionary_encode to avoid the overhead of converting to numpy
1980 and then factorizing.
1981 """
1982 from pandas.core.util.hashing import hash_array
1984 pa_type = self._pa_array.type
1986 if categorize and (
1987 pa.types.is_string(pa_type)
1988 or pa.types.is_large_string(pa_type)
1989 or pa.types.is_binary(pa_type)
1990 or pa.types.is_large_binary(pa_type)
1991 ):
1992 # Use pyarrow's dictionary_encode instead of numpy factorize.
1993 # dictionary_encode on a ChunkedArray unifies dictionaries
1994 # across chunks, so we can work with the result directly.
1995 encoded = pc.dictionary_encode(self._pa_array)
1996 if encoded.num_chunks <= 1:
1997 if encoded.num_chunks == 0:
1998 return np.array([], dtype=np.uint64)
1999 encoded_arr = encoded.chunk(0)
2000 else:
2001 # Concatenate indices from chunks; dictionaries are unified.
2002 indices_chunks = [chunk.indices for chunk in encoded.iterchunks()]
2003 combined_indices = pa.concat_arrays(indices_chunks)
2004 encoded_arr = pa.DictionaryArray.from_arrays(
2005 combined_indices, encoded.chunk(0).dictionary
2006 )
2007 dictionary = encoded_arr.dictionary.to_numpy(zero_copy_only=False)
2009 # Hash the unique values
2010 hashed = hash_array(dictionary, encoding, hash_key, categorize=False)
2012 # Map indices to hashes, handling nulls
2013 indices = encoded_arr.indices
2014 null_mask = indices.is_null().to_numpy(zero_copy_only=False)
2015 codes = indices.fill_null(0).to_numpy()
2017 if len(hashed):
2018 result = hashed.take(codes)
2019 else:
2020 result = np.zeros(len(codes), dtype=np.uint64)
2022 if null_mask.any():
2023 result[null_mask] = lib.u8max
2025 return result
2027 return hash_array(
2028 self.to_numpy(),
2029 encoding,
2030 hash_key,
2031 categorize,
2032 )
2034 def _values_for_factorize(self) -> tuple[np.ndarray, Any]:
2035 """
2036 Return an array and missing value suitable for factorization.
2038 Returns
2039 -------
2040 values : ndarray
2041 na_value : pd.NA
2042 """
2043 values = self._pa_array.to_numpy()
2044 return values, self.dtype.na_value
2046 def factorize(
2047 self,
2048 use_na_sentinel: bool = True,
2049 ) -> tuple[np.ndarray, ExtensionArray]:
2050 """
2051 Encode the arrow array as an enumerated type.
2053 Parameters
2054 ----------
2055 use_na_sentinel : bool, default True
2056 If True, the sentinel -1 will be used for NaN values. If False,
2057 NaN values will be encoded as non-negative integers and will not drop the
2058 NaN from the uniques of the values.
2060 Returns
2061 -------
2062 codes : ndarray
2063 An integer NumPy array that's an indexer into the original
2064 ArrowArray.
2065 uniques : ArrowArray
2066 An ArrowArray containing the unique values of `self`.
2068 .. note::
2070 uniques will *not* contain an entry for the NA value of
2071 the ArrowArray if there are any missing values present
2072 in `self`.
2074 See Also
2075 --------
2076 factorize : Top-level factorize method that dispatches here.
2078 Notes
2079 -----
2080 :meth:`pandas.factorize` offers a `sort` keyword as well.
2082 Examples
2083 --------
2084 >>> arr = pd.array(
2085 ... ["Ant", "Badger", "Cobra", "Cobra", "Deer", "Ant"], dtype=str
2086 ... )
2087 >>> arr_codes, arr_uniques = arr.factorize()
2088 >>> arr_codes
2089 array([0, 1, 2, 2, 3, 0])
2090 >>> arr_uniques
2091 <ArrowStringArray>
2092 ['Ant', 'Badger', 'Cobra', 'Deer']
2093 Length: 4, dtype: str
2094 """
2095 null_encoding = "mask" if use_na_sentinel else "encode"
2097 data = self._pa_array
2099 if pa.types.is_null(data.type):
2100 if use_na_sentinel or len(self) == 0:
2101 indices = np.full(len(self), -1, dtype=np.intp)
2102 uniques = self._from_pyarrow_array(pa.chunked_array([], type=pa.null()))
2103 else:
2104 indices = np.zeros(len(self), dtype=np.intp)
2105 uniques = self._from_pyarrow_array(pa.array([None], type=pa.null()))
2106 return indices, uniques
2108 if pa.types.is_dictionary(data.type):
2109 if null_encoding == "encode":
2110 # dictionary encode does nothing if an already encoded array is given
2111 data = data.cast(data.type.value_type)
2112 encoded = data.dictionary_encode(null_encoding=null_encoding)
2113 else:
2114 encoded = data
2115 else:
2116 encoded = data.dictionary_encode(null_encoding=null_encoding)
2117 if encoded.length() == 0:
2118 indices = np.array([], dtype=np.intp)
2119 uniques = self._from_pyarrow_array(
2120 pa.chunked_array([], type=encoded.type.value_type)
2121 )
2122 else:
2123 # GH 54844
2124 combined = encoded.combine_chunks()
2125 pa_indices = combined.indices
2126 if pa_indices.null_count > 0:
2127 pa_indices = _safe_fill_null(pa_indices, -1)
2128 indices = pa_indices.to_numpy(zero_copy_only=False, writable=True).astype(
2129 np.intp, copy=False
2130 )
2131 uniques = self._from_pyarrow_array(combined.dictionary)
2133 return indices, uniques
2135 def reshape(self, *args, **kwargs):
2136 raise NotImplementedError(
2137 f"{type(self)} does not support reshape "
2138 f"as backed by a 1D pyarrow.ChunkedArray."
2139 )
2141 def round(self, decimals: int = 0, *args, **kwargs) -> Self:
2142 """
2143 Round each value in the array a to the given number of decimals.
2145 Parameters
2146 ----------
2147 decimals : int, default 0
2148 Number of decimal places to round to. If decimals is negative,
2149 it specifies the number of positions to the left of the decimal point.
2150 *args, **kwargs
2151 Additional arguments and keywords have no effect.
2153 Returns
2154 -------
2155 ArrowExtensionArray
2156 Rounded values of the ArrowExtensionArray.
2158 See Also
2159 --------
2160 DataFrame.round : Round values of a DataFrame.
2161 Series.round : Round values of a Series.
2162 """
2163 if not self.dtype._is_numeric or self.dtype._is_boolean:
2164 # pc.round raises ArrowNotImplementedError on boolean and
2165 # non-numeric types; defer to the base for the (copy / TypeError)
2166 # contract.
2167 return super().round(decimals, *args, **kwargs)
2168 result = pc.round(self._pa_array, ndigits=decimals)
2169 return self._from_pyarrow_array(result)
2171 def searchsorted(
2172 self,
2173 value: NumpyValueArrayLike | ExtensionArray,
2174 side: Literal["left", "right"] = "left",
2175 sorter: NumpySorter | None = None,
2176 ) -> npt.NDArray[np.intp] | np.intp:
2177 """
2178 Find indices where elements should be inserted to maintain order.
2180 Find the indices into a sorted array `self` (a) such that, if the
2181 corresponding elements in `value` were inserted before the indices,
2182 the order of `self` would be preserved.
2184 Assuming that `self` is sorted:
2186 ====== ================================
2187 `side` returned index `i` satisfies
2188 ====== ================================
2189 left ``self[i-1] < value <= self[i]``
2190 right ``self[i-1] <= value < self[i]``
2191 ====== ================================
2193 Parameters
2194 ----------
2195 value : array-like, list or scalar
2196 Value(s) to insert into `self`.
2197 side : {'left', 'right'}, optional
2198 If 'left', the index of the first suitable location found is given.
2199 If 'right', return the last such index. If there is no suitable
2200 index, return either 0 or N (where N is the length of `self`).
2201 sorter : 1-D array-like, optional
2202 Optional array of integer indices that sort array a into ascending
2203 order. They are typically the result of argsort.
2205 Returns
2206 -------
2207 array of ints or int
2208 If value is array-like, array of insertion points.
2209 If value is scalar, a single integer.
2211 See Also
2212 --------
2213 numpy.searchsorted : Similar method from NumPy.
2215 Examples
2216 --------
2217 >>> arr = pd.array([1, 2, 3, 5], dtype="int64[pyarrow]")
2218 >>> arr.searchsorted([4])
2219 array([3])
2220 """
2221 if self._hasna:
2222 raise ValueError(
2223 "searchsorted requires array to be sorted, which is impossible "
2224 "with NAs present."
2225 )
2226 if isinstance(value, ExtensionArray):
2227 value = value.astype(object)
2228 # Base class searchsorted would cast to object, which is *much* slower.
2229 dtype = None
2230 if (
2231 is_arrow_temporal_dtype(self.dtype)
2232 and self.dtype.pyarrow_dtype.unit == "ns"
2233 ):
2234 # np.array[datetime/timedelta].searchsorted(datetime/timedelta)
2235 # erroneously fails when numpy type resolution is nanoseconds
2236 dtype = object
2237 return self.to_numpy(dtype=dtype).searchsorted(value, side=side, sorter=sorter)
2239 def take(
2240 self,
2241 indices: TakeIndexer,
2242 allow_fill: bool = False,
2243 fill_value: Any = None,
2244 ) -> ArrowExtensionArray:
2245 """
2246 Take elements from an array.
2248 Parameters
2249 ----------
2250 indices : sequence of int or one-dimensional np.ndarray of int
2251 Indices to be taken.
2252 allow_fill : bool, default False
2253 How to handle negative values in `indices`.
2255 * False: negative values in `indices` indicate positional indices
2256 from the right (the default). This is similar to
2257 :func:`numpy.take`.
2259 * True: negative values in `indices` indicate
2260 missing values. These values are set to `fill_value`. Any other
2261 other negative values raise a ``ValueError``.
2263 fill_value : any, optional
2264 Fill value to use for NA-indices when `allow_fill` is True.
2265 This may be ``None``, in which case the default NA value for
2266 the type, ``self.dtype.na_value``, is used.
2268 For many ExtensionArrays, there will be two representations of
2269 `fill_value`: a user-facing "boxed" scalar, and a low-level
2270 physical NA value. `fill_value` should be the user-facing version,
2271 and the implementation should handle translating that to the
2272 physical version for processing the take if necessary.
2274 Returns
2275 -------
2276 ExtensionArray
2278 Raises
2279 ------
2280 IndexError
2281 When the indices are out of bounds for the array.
2282 ValueError
2283 When `indices` contains negative values other than ``-1``
2284 and `allow_fill` is True.
2286 See Also
2287 --------
2288 numpy.take
2289 api.extensions.take
2291 Notes
2292 -----
2293 ExtensionArray.take is called by ``Series.__getitem__``, ``.loc``,
2294 ``iloc``, when `indices` is a sequence of values. Additionally,
2295 it's called by :meth:`Series.reindex`, or any other method
2296 that causes realignment, with a `fill_value`.
2297 """
2298 indices_array = np.asanyarray(indices)
2300 if len(self._pa_array) == 0 and (indices_array >= 0).any():
2301 raise IndexError("cannot do a non-empty take")
2302 if indices_array.size > 0 and indices_array.max() >= len(self._pa_array):
2303 raise IndexError("out of bounds value in 'indices'.")
2305 if allow_fill:
2306 fill_mask = indices_array < 0
2307 if fill_mask.any():
2308 validate_indices(indices_array, len(self._pa_array))
2309 # TODO(ARROW-9433): Treat negative indices as NULL
2310 indices_array = pa.array(indices_array, mask=fill_mask)
2311 result = self._pa_array.take(indices_array)
2312 if isna(fill_value):
2313 return self._from_pyarrow_array(result)
2314 # TODO: ArrowNotImplementedError: Function fill_null has no
2315 # kernel matching input types (array[string], scalar[string])
2316 result = self._from_pyarrow_array(result)
2317 result[fill_mask] = fill_value
2318 return result
2319 # return type(self)(pc.fill_null(result, pa.scalar(fill_value)))
2320 else:
2321 # Nothing to fill
2322 return self._from_pyarrow_array(self._pa_array.take(indices))
2323 else: # allow_fill=False
2324 # TODO(ARROW-9432): Treat negative indices as indices from the right.
2325 if (indices_array < 0).any():
2326 # Don't modify in-place
2327 indices_array = np.copy(indices_array)
2328 indices_array[indices_array < 0] += len(self._pa_array)
2329 return self._from_pyarrow_array(self._pa_array.take(indices_array))
2331 def _maybe_convert_datelike_array(self):
2332 """Maybe convert to a datelike array."""
2333 pa_type = self._pa_array.type
2334 if pa.types.is_timestamp(pa_type):
2335 return self._to_datetimearray()
2336 elif pa.types.is_duration(pa_type):
2337 return self._to_timedeltaarray()
2338 return self
2340 def _to_datetimearray(self) -> DatetimeArray:
2341 """Convert a pyarrow timestamp typed array to a DatetimeArray."""
2342 from pandas.core.arrays.datetimes import (
2343 DatetimeArray,
2344 tz_to_dtype,
2345 )
2347 pa_type = self._pa_array.type
2348 assert pa.types.is_timestamp(pa_type)
2349 np_dtype = np.dtype(f"M8[{pa_type.unit}]")
2350 dtype = tz_to_dtype(pa_type.tz, pa_type.unit)
2351 np_array = self._pa_array.to_numpy()
2352 np_array = np_array.astype(np_dtype, copy=False)
2353 return DatetimeArray._simple_new(np_array, dtype=dtype)
2355 def _to_timedeltaarray(self) -> TimedeltaArray:
2356 """Convert a pyarrow duration typed array to a TimedeltaArray."""
2357 from pandas.core.arrays.timedeltas import TimedeltaArray
2359 pa_type = self._pa_array.type
2360 assert pa.types.is_duration(pa_type)
2361 np_dtype = np.dtype(f"m8[{pa_type.unit}]")
2362 np_array = self._pa_array.to_numpy()
2363 np_array = np_array.astype(np_dtype, copy=False)
2364 return TimedeltaArray._simple_new(np_array, dtype=np_dtype)
2366 def _values_for_json(self) -> np.ndarray:
2367 if is_numeric_dtype(self.dtype):
2368 return np.asarray(self, dtype=object)
2369 return super()._values_for_json()
2371 def to_numpy(
2372 self,
2373 dtype: npt.DTypeLike | None = None,
2374 copy: bool = False,
2375 na_value: object = lib.no_default,
2376 ) -> np.ndarray:
2377 """
2378 Convert to a NumPy ndarray.
2380 This is similar to :meth:`numpy.asarray`, but may provide additional control
2381 over how the conversion is done.
2383 Parameters
2384 ----------
2385 dtype : str or numpy.dtype, optional
2386 The dtype to pass to :meth:`numpy.asarray`.
2387 copy : bool, default False
2388 Whether to ensure that the returned value is a not a view on
2389 another array. Note that ``copy=False`` does not *ensure* that
2390 ``to_numpy()`` is no-copy. Rather, ``copy=True`` ensure that
2391 a copy is made, even if not strictly necessary.
2392 na_value : Any, optional
2393 The value to use for missing values. The default value depends
2394 on `dtype` and the type of the array.
2396 Returns
2397 -------
2398 numpy.ndarray
2399 """
2400 original_na_value = na_value
2401 dtype, na_value = to_numpy_dtype_inference(self, dtype, na_value, self._hasna)
2402 pa_type = self._pa_array.type
2403 if not self._hasna or isna(na_value) or pa.types.is_null(pa_type):
2404 data = self
2405 else:
2406 data = self.fillna(na_value)
2407 copy = False
2409 if is_arrow_temporal_dtype(self.dtype):
2410 # GH 55997
2411 if dtype != object and na_value is self.dtype.na_value:
2412 na_value = lib.no_default
2413 result = data._maybe_convert_datelike_array().to_numpy(
2414 dtype=dtype, na_value=na_value
2415 )
2416 elif pa.types.is_time(pa_type) or pa.types.is_date(pa_type):
2417 # convert to list of python datetime.time objects before
2418 # wrapping in ndarray
2419 result = np.array(list(data), dtype=dtype)
2420 if data._hasna:
2421 result[data.isna()] = na_value
2422 elif pa.types.is_null(pa_type):
2423 if dtype is not None and isna(na_value):
2424 na_value = None
2425 result = np.full(len(data), fill_value=na_value, dtype=dtype)
2426 elif not data._hasna or (
2427 pa.types.is_floating(pa_type)
2428 and (
2429 na_value is np.nan
2430 or (
2431 original_na_value is lib.no_default
2432 and is_float_dtype(dtype)
2433 and is_nan_na()
2434 )
2435 )
2436 ):
2437 result = data._pa_array.to_numpy()
2438 if dtype is not None:
2439 result = result.astype(dtype, copy=False)
2440 if copy:
2441 result = result.copy()
2442 else:
2443 if dtype is None:
2444 empty = pa.array([], type=pa_type).to_numpy(zero_copy_only=False)
2445 if can_hold_element(empty, na_value):
2446 dtype = empty.dtype
2447 else:
2448 dtype = np.object_
2449 result = np.empty(len(data), dtype=dtype)
2450 mask = data.isna()
2451 result[mask] = na_value
2452 result[~mask] = data[~mask]._pa_array.to_numpy()
2453 return result
2455 def duplicated(
2456 self, keep: Literal["first", "last", False] = "first"
2457 ) -> npt.NDArray[np.bool_]:
2458 """
2459 Return boolean ndarray denoting duplicate values.
2461 Parameters
2462 ----------
2463 keep : {'first', 'last', False}, default 'first'
2464 - ``first`` : Mark duplicates as ``True`` except for the first occurrence.
2465 - ``last`` : Mark duplicates as ``True`` except for the last occurrence.
2466 - False : Mark all duplicates as ``True``.
2468 Returns
2469 -------
2470 ndarray[bool]
2471 With true in indices where elements are duplicated and false otherwise.
2473 See Also
2474 --------
2475 DataFrame.duplicated : Return boolean Series denoting
2476 duplicate rows.
2477 Series.duplicated : Indicate duplicate Series values.
2478 api.extensions.ExtensionArray.unique : Compute the ExtensionArray
2479 of unique values.
2481 Examples
2482 --------
2483 >>> pd.array([1, 1, 2, 3, 3], dtype="int64[pyarrow]").duplicated()
2484 array([False, True, False, False, True])
2485 """
2486 pa_type = self._pa_array.type
2487 if pa.types.is_floating(pa_type) or pa.types.is_integer(pa_type):
2488 values = self.to_numpy(na_value=0)
2489 elif pa.types.is_boolean(pa_type):
2490 values = self.to_numpy(na_value=False)
2491 elif pa.types.is_temporal(pa_type):
2492 if pa_type.bit_width == 32:
2493 pa_type = pa.int32()
2494 else:
2495 pa_type = pa.int64()
2496 arr = self.astype(ArrowDtype(pa_type))
2497 values = arr.to_numpy(na_value=0)
2498 else:
2499 # factorize the values to avoid the performance penalty of
2500 # converting to object dtype
2501 values = self.factorize()[0]
2503 mask = self.isna() if self._hasna else None
2504 return algos.duplicated(values, keep=keep, mask=mask)
2506 def unique(self) -> Self:
2507 """
2508 Compute the ArrowExtensionArray of unique values.
2510 Returns
2511 -------
2512 ArrowExtensionArray
2513 """
2514 pa_result = pc.unique(self._pa_array)
2515 return self._from_pyarrow_array(pa_result)
2517 def value_counts(self, dropna: bool = True) -> Series:
2518 """
2519 Return a Series containing counts of each unique value.
2521 Parameters
2522 ----------
2523 dropna : bool, default True
2524 Don't include counts of missing values.
2526 Returns
2527 -------
2528 counts : Series
2530 See Also
2531 --------
2532 Series.value_counts
2533 """
2534 from pandas import (
2535 Index,
2536 Series,
2537 )
2539 data = self._pa_array
2540 vc = data.value_counts()
2542 values = vc.field(0)
2543 counts = vc.field(1)
2544 if dropna and data.null_count > 0:
2545 mask = values.is_valid()
2546 values = values.filter(mask)
2547 counts = counts.filter(mask)
2549 counts = ArrowExtensionArray(counts)
2551 index = Index(self._from_pyarrow_array(values), copy=False)
2553 return Series(counts, index=index, name="count", copy=False)
2555 @classmethod
2556 def _concat_same_type(cls, to_concat) -> Self:
2557 """
2558 Concatenate multiple ArrowExtensionArrays.
2560 Parameters
2561 ----------
2562 to_concat : sequence of ArrowExtensionArrays
2564 Returns
2565 -------
2566 ArrowExtensionArray
2567 """
2568 chunks = [array for ea in to_concat for array in ea._pa_array.iterchunks()]
2569 if to_concat[0].dtype == "string":
2570 # StringDtype has no attribute pyarrow_dtype
2571 pa_dtype = pa.large_string()
2572 else:
2573 pa_dtype = to_concat[0].dtype.pyarrow_dtype
2574 arr = pa.chunked_array(chunks, type=pa_dtype)
2575 return to_concat[0]._from_pyarrow_array(arr)
2577 def _accumulate(
2578 self, name: str, *, skipna: bool = True, **kwargs
2579 ) -> ArrowExtensionArray | ExtensionArray:
2580 """
2581 Return an ExtensionArray performing an accumulation operation.
2583 The underlying data type might change.
2585 Parameters
2586 ----------
2587 name : str
2588 Name of the function, supported values are:
2589 - cummin
2590 - cummax
2591 - cumsum
2592 - cumprod
2593 skipna : bool, default True
2594 If True, skip NA values.
2595 **kwargs
2596 Additional keyword arguments passed to the accumulation function.
2597 Currently, there is no supported kwarg.
2599 Returns
2600 -------
2601 array
2603 Raises
2604 ------
2605 NotImplementedError : subclass does not define accumulations
2606 """
2607 if is_string_dtype(self):
2608 return self._str_accumulate(name=name, skipna=skipna, **kwargs)
2610 pyarrow_name = {
2611 "cummax": "cumulative_max",
2612 "cummin": "cumulative_min",
2613 "cumprod": "cumulative_prod_checked",
2614 "cumsum": "cumulative_sum_checked",
2615 }.get(name, name)
2616 pyarrow_meth = getattr(pc, pyarrow_name, None)
2617 if pyarrow_meth is None:
2618 return super()._accumulate(name, skipna=skipna, **kwargs)
2620 data_to_accum = self._pa_array
2622 pa_dtype = data_to_accum.type
2624 convert_to_int = (
2625 pa.types.is_temporal(pa_dtype) and name in ["cummax", "cummin"]
2626 ) or (pa.types.is_duration(pa_dtype) and name == "cumsum")
2628 if convert_to_int:
2629 if pa_dtype.bit_width == 32:
2630 data_to_accum = data_to_accum.cast(pa.int32())
2631 else:
2632 data_to_accum = data_to_accum.cast(pa.int64())
2634 if name in ("cummax", "cummin") and pa.types.is_floating(data_to_accum.type):
2635 kwargs["start"] = float("-inf") if name == "cummax" else float("inf")
2637 try:
2638 result = pyarrow_meth(data_to_accum, skip_nulls=skipna, **kwargs)
2639 except pa.ArrowNotImplementedError as err:
2640 msg = f"operation '{name}' not supported for dtype '{self.dtype}'"
2641 raise TypeError(msg) from err
2643 if convert_to_int:
2644 result = result.cast(pa_dtype)
2646 return self._from_pyarrow_array(result)
2648 def _str_accumulate(
2649 self, name: str, *, skipna: bool = True, **kwargs
2650 ) -> ArrowExtensionArray | ExtensionArray:
2651 """
2652 Accumulate implementation for strings, see `_accumulate` docstring for details.
2654 pyarrow.compute does not implement these methods for strings.
2655 """
2656 if name == "cumprod":
2657 msg = f"operation '{name}' not supported for dtype '{self.dtype}'"
2658 raise TypeError(msg)
2660 # We may need to strip out trailing NA values
2661 tail: pa.array | None = None
2662 na_mask: pa.array | None = None
2663 pa_array = self._pa_array
2664 np_func = {
2665 "cumsum": np.cumsum,
2666 "cummin": np.minimum.accumulate,
2667 "cummax": np.maximum.accumulate,
2668 }[name]
2670 if self._hasna:
2671 na_mask = pc.is_null(pa_array)
2672 if pc.all(na_mask) == pa.scalar(True):
2673 return self._from_pyarrow_array(pa_array)
2674 if skipna:
2675 if name == "cumsum":
2676 pa_array = _safe_fill_null(pa_array, "")
2677 else:
2678 # We can retain the running min/max by forward/backward filling.
2679 pa_array = pc.fill_null_forward(pa_array)
2680 pa_array = pc.fill_null_backward(pa_array)
2681 else:
2682 # When not skipping NA values, the result should be null from
2683 # the first NA value onward.
2684 idx = pc.index(na_mask, True).as_py()
2685 tail = pa.nulls(len(pa_array) - idx, type=pa_array.type)
2686 pa_array = pa_array[:idx]
2688 # error: Cannot call function of unknown type
2689 pa_result = pa.array(np_func(pa_array), type=pa_array.type) # type: ignore[operator]
2691 if tail is not None:
2692 pa_result = pa.concat_arrays([pa_result, tail])
2693 elif na_mask is not None:
2694 pa_result = pc.if_else(na_mask, None, pa_result)
2696 result = self._from_pyarrow_array(pa_result)
2697 return result
2699 def _reduce_pyarrow(self, name: str, *, skipna: bool = True, **kwargs) -> pa.Scalar:
2700 """
2701 Return a pyarrow scalar result of performing the reduction operation.
2703 Parameters
2704 ----------
2705 name : str
2706 Name of the function, supported values are:
2707 { any, all, min, max, sum, mean, median, prod,
2708 std, var, sem, kurt, skew }.
2709 skipna : bool, default True
2710 If True, skip NaN values.
2711 **kwargs
2712 Additional keyword arguments passed to the reduction function.
2713 Currently, `ddof` is the only supported kwarg.
2715 Returns
2716 -------
2717 pyarrow scalar
2719 Raises
2720 ------
2721 TypeError : subclass does not define reductions
2722 """
2723 pa_type = self._pa_array.type
2725 data_to_reduce = self._pa_array
2727 if name == "count":
2728 return self.count()
2730 elif name in ["any", "all"] and (
2731 pa.types.is_integer(pa_type)
2732 or pa.types.is_floating(pa_type)
2733 or pa.types.is_duration(pa_type)
2734 or pa.types.is_decimal(pa_type)
2735 ):
2736 # pyarrow only supports any/all for boolean dtype, we allow
2737 # for other dtypes, matching our non-pyarrow behavior
2739 if pa.types.is_duration(pa_type):
2740 data_to_cmp = self._pa_array.cast(pa.int64())
2741 else:
2742 data_to_cmp = self._pa_array
2744 not_eq = pc.not_equal(data_to_cmp, 0)
2745 data_to_reduce = not_eq
2747 elif name in ["min", "max", "sum"] and pa.types.is_duration(pa_type):
2748 data_to_reduce = self._pa_array.cast(pa.int64())
2750 elif name in ["median", "mean", "std", "sem"] and pa.types.is_temporal(pa_type):
2751 nbits = pa_type.bit_width
2752 if name in ["std", "sem"] and pa.types.is_date32(pa_type):
2753 # compute in seconds, the unit of the result
2754 seconds = self._pa_array.cast(pa.timestamp("s"))
2755 data_to_reduce = seconds.cast(pa.int64())
2756 elif nbits == 32:
2757 data_to_reduce = self._pa_array.cast(pa.int32())
2758 else:
2759 data_to_reduce = self._pa_array.cast(pa.int64())
2761 if name == "sem":
2763 def pyarrow_meth(data, skip_nulls, **kwargs): # pyright: ignore[reportRedeclaration]
2764 numerator = pc.stddev(data, skip_nulls=skip_nulls, **kwargs)
2765 denominator = pc.sqrt_checked(pc.count(self._pa_array))
2766 return pc.divide_checked(numerator, denominator)
2768 elif name == "sum" and (
2769 pa.types.is_string(pa_type) or pa.types.is_large_string(pa_type)
2770 ):
2772 def pyarrow_meth(data, skip_nulls, min_count=0): # type: ignore[misc]
2773 mask = pc.is_null(data) if data.null_count > 0 else None
2774 if skip_nulls:
2775 if min_count > 0 and check_below_min_count(
2776 (len(data),),
2777 None if mask is None else mask.to_numpy(),
2778 min_count,
2779 ):
2780 return pa.scalar(None, type=data.type)
2781 if data.null_count > 0:
2782 # binary_join returns null if there is any null ->
2783 # have to filter out any nulls
2784 data = data.filter(pc.invert(mask))
2785 elif mask is not None or check_below_min_count(
2786 (len(data),), None, min_count
2787 ):
2788 return pa.scalar(None, type=data.type)
2790 if pa.types.is_large_string(data.type):
2791 # binary_join only supports string, not large_string
2792 data = data.cast(pa.string())
2793 data_list = pa.ListArray.from_arrays(
2794 [0, len(data)], data.combine_chunks()
2795 )[0]
2796 return pc.binary_join(data_list, "")
2797 elif name in ["argmin", "argmax"]:
2798 return super()._reduce(name, skipna=skipna, **kwargs)
2800 else:
2801 pyarrow_name = {
2802 "median": "quantile",
2803 "prod": "product",
2804 "std": "stddev",
2805 "var": "variance",
2806 "kurt": "kurtosis",
2807 }.get(name, name)
2808 # error: Incompatible types in assignment
2809 # (expression has type "Optional[Any]", variable has type
2810 # "Callable[[Any, Any, KwArg(Any)], Any]")
2811 pyarrow_meth = getattr(pc, pyarrow_name, None) # type: ignore[assignment]
2812 if pyarrow_meth is None:
2813 raise TypeError(
2814 f"'{type(self).__name__}' with dtype {self.dtype} "
2815 f"does not support operation '{name}'"
2816 )
2818 # GH51624: pyarrow defaults to min_count=1, pandas behavior is min_count=0
2819 if name in ["any", "all", "sum", "prod"] and "min_count" not in kwargs:
2820 kwargs["min_count"] = 0
2821 elif name == "median":
2822 # GH 52679: Use quantile instead of approximate_median
2823 kwargs["q"] = 0.5
2824 elif name in ["std", "var", "sem"] and "ddof" not in kwargs:
2825 # pyarrow defaults to ddof=0, pandas behavior is ddof=1
2826 kwargs["ddof"] = 1
2827 elif name in ["skew", "kurt"] and "biased" not in kwargs:
2828 kwargs["biased"] = False
2830 try:
2831 result = pyarrow_meth(data_to_reduce, skip_nulls=skipna, **kwargs)
2832 except (AttributeError, NotImplementedError, TypeError) as err:
2833 msg = (
2834 f"'{type(self).__name__}' with dtype {self.dtype} "
2835 f"does not support operation '{name}' with pyarrow "
2836 f"version {pa.__version__}. '{name}' may be supported by "
2837 f"upgrading pyarrow."
2838 )
2839 raise TypeError(msg) from err
2840 if name == "median":
2841 # GH 52679: Use quantile instead of approximate_median; returns array
2842 result = result[0]
2844 if name in ["min", "max", "sum"] and pa.types.is_duration(pa_type):
2845 result = result.cast(pa_type)
2846 if name in ["median", "mean"] and pa.types.is_temporal(pa_type):
2847 nbits = pa_type.bit_width
2848 if nbits == 32:
2849 result = result.cast(pa.int32(), safe=False)
2850 else:
2851 result = result.cast(pa.int64(), safe=False)
2852 result = result.cast(pa_type)
2853 if name in ["std", "sem"] and pa.types.is_temporal(pa_type):
2854 result = result.cast(pa.int64(), safe=False)
2855 if pa.types.is_duration(pa_type):
2856 result = result.cast(pa_type)
2857 elif pa.types.is_time(pa_type):
2858 result = result.cast(pa.duration(pa_type.unit))
2859 elif pa.types.is_date32(pa_type):
2860 # computed in seconds, the closest available unit
2861 result = result.cast(pa.duration("s"))
2862 elif pa.types.is_date64(pa_type):
2863 # the result is in milliseconds, the storage unit
2864 result = result.cast(pa.duration("ms"))
2865 else:
2866 # i.e. timestamp
2867 result = result.cast(pa.duration(pa_type.unit))
2869 return result
2871 def _reduce(
2872 self,
2873 name: str,
2874 *,
2875 skipna: bool = True,
2876 keepdims: bool = False,
2877 axis: AxisInt | None = 0,
2878 **kwargs,
2879 ):
2880 """
2881 Return a scalar result of performing the reduction operation.
2883 Parameters
2884 ----------
2885 name : str
2886 Name of the function, supported values are:
2887 { any, all, min, max, sum, mean, median, prod,
2888 std, var, sem, kurt, skew }.
2889 skipna : bool, default True
2890 If True, skip NaN values.
2891 axis : int or None, default 0
2892 The axis to reduce over; for a 1-D array this must be 0 or None.
2893 **kwargs
2894 Additional keyword arguments passed to the reduction function.
2895 Currently, `ddof` is the only supported kwarg.
2897 Returns
2898 -------
2899 scalar
2901 Raises
2902 ------
2903 TypeError : subclass does not define reductions
2904 """
2905 nv.validate_minmax_axis(axis, self.ndim)
2906 result = self._reduce_calc(name, skipna=skipna, keepdims=keepdims, **kwargs)
2907 if isinstance(result, pa.Array):
2908 return self._from_pyarrow_array(result)
2909 else:
2910 return result
2912 def _reduce_calc(
2913 self, name: str, *, skipna: bool = True, keepdims: bool = False, **kwargs
2914 ):
2915 pa_result = self._reduce_pyarrow(name, skipna=skipna, **kwargs)
2917 if keepdims:
2918 if isinstance(pa_result, pa.Scalar):
2919 result = pa.array([pa_result.as_py()], type=pa_result.type)
2920 else:
2921 result = pa.array(
2922 [pa_result],
2923 type=to_pyarrow_type(infer_dtype_from_scalar(pa_result)[0]),
2924 )
2925 return result
2927 if pc.is_null(pa_result).as_py():
2928 return self.dtype.na_value
2929 elif isinstance(pa_result, pa.Scalar):
2930 result = pa_result.as_py()
2931 pa_type = pa_result.type
2932 if pa.types.is_duration(pa_type) and pa_type.unit != "ns":
2933 return Timedelta(result).as_unit(pa_type.unit)
2934 elif pa.types.is_timestamp(pa_type) and pa_type.unit != "ns":
2935 return Timestamp(result).as_unit(pa_type.unit)
2936 return result
2937 else:
2938 return pa_result
2940 def sum(
2941 self,
2942 *,
2943 skipna: bool = True,
2944 min_count: int = 0,
2945 axis: AxisInt | None = 0,
2946 **kwargs,
2947 ):
2948 nv.validate_sum((), kwargs)
2949 return self._reduce(
2950 "sum", skipna=skipna, min_count=min_count, axis=axis, **kwargs
2951 )
2953 def min(
2954 self,
2955 *,
2956 skipna: bool = True,
2957 axis: AxisInt | None = 0,
2958 **kwargs,
2959 ):
2960 nv.validate_min((), kwargs)
2961 return self._reduce("min", skipna=skipna, axis=axis, **kwargs)
2963 def max(
2964 self,
2965 *,
2966 skipna: bool = True,
2967 axis: AxisInt | None = 0,
2968 **kwargs,
2969 ):
2970 nv.validate_max((), kwargs)
2971 return self._reduce("max", skipna=skipna, axis=axis, **kwargs)
2973 def mean(
2974 self,
2975 *,
2976 skipna: bool = True,
2977 axis: AxisInt | None = 0,
2978 **kwargs,
2979 ):
2980 nv.validate_mean((), kwargs)
2981 return self._reduce("mean", skipna=skipna, axis=axis, **kwargs)
2983 def sem(
2984 self,
2985 *,
2986 skipna: bool = True,
2987 axis: AxisInt | None = 0,
2988 ddof: int = 1,
2989 **kwargs,
2990 ):
2991 nv.validate_stat_ddof_func((), kwargs, fname="sem")
2992 return self._reduce("sem", skipna=skipna, ddof=ddof, axis=axis, **kwargs)
2994 def skew(
2995 self,
2996 *,
2997 skipna: bool = True,
2998 axis: AxisInt | None = 0,
2999 **kwargs,
3000 ):
3001 nv.validate_stat_ddof_func((), kwargs, fname="skew")
3002 return self._reduce("skew", skipna=skipna, axis=axis, **kwargs)
3004 def kurt(
3005 self,
3006 *,
3007 skipna: bool = True,
3008 axis: AxisInt | None = 0,
3009 **kwargs,
3010 ):
3011 nv.validate_stat_ddof_func((), kwargs, fname="kurt")
3012 return self._reduce("kurt", skipna=skipna, axis=axis, **kwargs)
3014 def median(
3015 self,
3016 *,
3017 skipna: bool = True,
3018 axis: AxisInt | None = 0,
3019 **kwargs,
3020 ):
3021 nv.validate_median((), kwargs)
3022 return self._reduce("median", skipna=skipna, axis=axis, **kwargs)
3024 def var(
3025 self,
3026 *,
3027 skipna: bool = True,
3028 axis: AxisInt | None = 0,
3029 ddof: int = 1,
3030 **kwargs,
3031 ):
3032 nv.validate_stat_ddof_func((), kwargs, fname="var")
3033 return self._reduce("var", skipna=skipna, ddof=ddof, axis=axis, **kwargs)
3035 def std(
3036 self,
3037 *,
3038 skipna: bool = True,
3039 axis: AxisInt | None = 0,
3040 ddof: int = 1,
3041 **kwargs,
3042 ):
3043 nv.validate_stat_ddof_func((), kwargs, fname="std")
3044 return self._reduce("std", skipna=skipna, ddof=ddof, axis=axis, **kwargs)
3046 def prod(
3047 self,
3048 *,
3049 skipna: bool = True,
3050 min_count: int = 0,
3051 axis: AxisInt | None = 0,
3052 **kwargs,
3053 ):
3054 nv.validate_prod((), kwargs)
3055 return self._reduce(
3056 "prod", skipna=skipna, min_count=min_count, axis=axis, **kwargs
3057 )
3059 def _explode(self):
3060 """
3061 See Series.explode.__doc__.
3062 """
3063 # child class explode method supports only list types; return
3064 # default implementation for non list types.
3065 if not hasattr(self.dtype, "pyarrow_dtype") or (
3066 not pa.types.is_list(self.dtype.pyarrow_dtype)
3067 and not pa.types.is_large_list(self.dtype.pyarrow_dtype)
3068 ):
3069 return super()._explode()
3070 values = self
3071 counts = pa.compute.list_value_length(values._pa_array)
3072 counts = counts.fill_null(1).to_numpy()
3073 fill_value = pa.scalar([None], type=self._pa_array.type)
3074 mask = counts == 0
3075 if mask.any():
3076 # pc.if_else here is similar to `values[mask] = fill_value`
3077 # but this avoids an object-dtype round-trip.
3078 pa_values = pc.if_else(~mask, values._pa_array, fill_value)
3079 values = self._from_pyarrow_array(pa_values)
3080 counts = counts.copy()
3081 counts[mask] = 1
3082 values = values.fillna(fill_value)
3083 values = self._from_pyarrow_array(pa.compute.list_flatten(values._pa_array))
3084 return values, counts
3086 def __setitem__(self, key, value) -> None:
3087 """Set one or more values inplace.
3089 Parameters
3090 ----------
3091 key : int, ndarray, or slice
3092 When called from, e.g. ``Series.__setitem__``, ``key`` will be
3093 one of
3095 * scalar int
3096 * ndarray of integers.
3097 * boolean ndarray
3098 * slice object
3100 value : ExtensionDtype.type, Sequence[ExtensionDtype.type], or object
3101 value or values to be set of ``key``.
3103 Returns
3104 -------
3105 None
3106 """
3107 if self._readonly:
3108 raise ValueError("Cannot modify read-only array")
3110 # GH50085: unwrap 1D indexers
3111 if isinstance(key, tuple) and len(key) == 1:
3112 key = key[0]
3114 key = check_array_indexer(self, key)
3115 value = self._validate_setitem_value(value)
3117 if com.is_null_slice(key):
3118 # fast path (GH50248)
3119 if (
3120 isinstance(value, (pa.Array, pa.ChunkedArray))
3121 and value.type == self._pa_array.type
3122 and len(value) == len(self)
3123 ):
3124 # GH#67990 this adopts ``value`` as our backing array, so copy
3125 # first if the caller may still own and mutate its buffers.
3126 if _boxing_may_borrow_memory(value.type):
3127 value = _copy_pyarrow_buffers(value)
3128 data = value
3129 else:
3130 data = self._if_else(True, value, self._pa_array)
3132 elif is_integer(key):
3133 # fast path
3134 key = cast("int", key)
3135 n = len(self)
3136 if key < 0:
3137 key += n
3138 if not 0 <= key < n:
3139 raise IndexError(
3140 f"index {key} is out of bounds for axis 0 with size {n}"
3141 )
3142 if isinstance(value, pa.Scalar):
3143 value = value.as_py()
3144 elif is_list_like(value):
3145 raise ValueError("Length of indexer and values mismatch")
3146 chunks = [
3147 *self._pa_array[:key].chunks,
3148 pa.array([value], type=self._pa_array.type, from_pandas=is_nan_na()),
3149 *self._pa_array[key + 1 :].chunks,
3150 ]
3151 data = pa.chunked_array(chunks).combine_chunks()
3153 elif is_bool_dtype(key):
3154 key = np.asarray(key, dtype=np.bool_)
3155 data = self._replace_with_mask(self._pa_array, key, value)
3157 elif is_scalar(value) or isinstance(value, pa.Scalar):
3158 mask = np.zeros(len(self), dtype=np.bool_)
3159 mask[key] = True
3160 data = self._if_else(mask, value, self._pa_array)
3162 else:
3163 indices = np.arange(len(self))[key]
3164 if len(indices) != len(value):
3165 raise ValueError("Length of indexer and values mismatch")
3166 if len(indices) == 0:
3167 return
3168 # GH#58530 wrong item assignment by repeated key
3169 _, argsort = np.unique(indices, return_index=True)
3170 indices = indices[argsort]
3171 value = value.take(argsort)
3172 mask = np.zeros(len(self), dtype=np.bool_)
3173 mask[indices] = True
3174 data = self._replace_with_mask(self._pa_array, mask, value)
3176 if isinstance(data, pa.Array):
3177 data = pa.chunked_array([data])
3178 self._pa_array = data
3179 # Invalidate any cache_readonly properties that depend on _pa_array
3180 self._cache.clear()
3182 def _rank_calc(
3183 self,
3184 *,
3185 axis: AxisInt = 0,
3186 method: RankMethod = "average",
3187 na_option: RankNaOption = "keep",
3188 ascending: bool = True,
3189 pct: bool = False,
3190 ):
3191 if axis != 0:
3192 ranked = super()._rank(
3193 axis=axis,
3194 method=method,
3195 na_option=na_option,
3196 ascending=ascending,
3197 pct=pct,
3198 )
3199 # keep dtypes consistent with the implementation below
3200 if method == "average" or pct:
3201 pa_type = pa.float64()
3202 else:
3203 pa_type = pa.uint64()
3204 result = pa.array(ranked, type=pa_type, from_pandas=is_nan_na())
3205 return result
3207 data = self._pa_array.combine_chunks()
3208 order = "ascending" if ascending else "descending"
3209 null_placement = "at_start" if na_option == "top" else "at_end"
3210 tiebreaker = "min" if method == "average" else method
3212 rank_kwargs: dict[str, Any]
3213 if pa_version_under25p0:
3214 rank_kwargs = {"sort_keys": order, "null_placement": null_placement}
3215 else:
3216 # pyarrow 25 deprecated the null_placement keyword in favor of
3217 # specifying it per sort key
3218 rank_kwargs = {"sort_keys": [("", order, null_placement)]}
3220 result = pc.rank(data, tiebreaker=tiebreaker, **rank_kwargs)
3222 if na_option == "keep":
3223 mask = pc.is_null(self._pa_array)
3224 null = pa.scalar(None, type=result.type)
3225 result = pc.if_else(mask, null, result)
3227 if method == "average":
3228 result_max = pc.rank(data, tiebreaker="max", **rank_kwargs)
3229 result_max = result_max.cast(pa.float64())
3230 result_min = result.cast(pa.float64())
3231 result = pc.divide(pc.add(result_min, result_max), 2)
3233 if pct:
3234 if not pa.types.is_floating(result.type):
3235 result = result.cast(pa.float64())
3236 if method == "dense":
3237 divisor = pc.max(result)
3238 else:
3239 divisor = pc.count(result)
3240 result = pc.divide(result, divisor)
3242 return result
3244 @property
3245 def _is_monotonic_increasing(self) -> bool:
3246 pa_array = self._pa_array
3247 if pa_array.null_count > 0:
3248 return False
3249 if len(pa_array) <= 1:
3250 return True
3251 try:
3252 return pc.all(pc.greater_equal(pa_array[1:], pa_array[:-1])).as_py()
3253 except pa.ArrowNotImplementedError:
3254 # If the type doesn't support comparison, e.g. list types
3255 return False
3257 @property
3258 def _is_monotonic_decreasing(self) -> bool:
3259 pa_array = self._pa_array
3260 if pa_array.null_count > 0:
3261 return False
3262 if len(pa_array) <= 1:
3263 return True
3264 try:
3265 return pc.all(pc.less_equal(pa_array[1:], pa_array[:-1])).as_py()
3266 except pa.ArrowNotImplementedError:
3267 # If the type doesn't support comparison, e.g. list types
3268 return False
3270 def _rank(
3271 self,
3272 *,
3273 axis: AxisInt = 0,
3274 method: RankMethod = "average",
3275 na_option: RankNaOption = "keep",
3276 ascending: bool = True,
3277 pct: bool = False,
3278 ) -> Self:
3279 """
3280 See Series.rank.__doc__.
3281 """
3282 return self._convert_rank_result(
3283 self._rank_calc(
3284 axis=axis,
3285 method=method,
3286 na_option=na_option,
3287 ascending=ascending,
3288 pct=pct,
3289 )
3290 )
3292 def _quantile(self, qs: npt.NDArray[np.float64], interpolation: str) -> Self:
3293 """
3294 Compute the quantiles of self for each quantile in `qs`.
3296 Parameters
3297 ----------
3298 qs : np.ndarray[float64]
3299 interpolation: str
3301 Returns
3302 -------
3303 same type as self
3304 """
3305 pa_dtype = self._pa_array.type
3307 data = self._pa_array
3308 if pa.types.is_temporal(pa_dtype):
3309 # https://github.com/apache/arrow/issues/33769 in these cases
3310 # we can cast to ints and back
3311 nbits = pa_dtype.bit_width
3312 if nbits == 32:
3313 data = data.cast(pa.int32())
3314 else:
3315 data = data.cast(pa.int64())
3317 result = pc.quantile(data, q=qs, interpolation=interpolation)
3319 if pa.types.is_temporal(pa_dtype):
3320 if pa.types.is_floating(result.type):
3321 result = pc.floor(result)
3322 nbits = pa_dtype.bit_width
3323 if nbits == 32:
3324 result = result.cast(pa.int32())
3325 else:
3326 result = result.cast(pa.int64())
3327 result = result.cast(pa_dtype)
3329 return self._from_pyarrow_array(result)
3331 def _mode(self, dropna: bool = True) -> Self:
3332 """
3333 Returns the mode(s) of the ExtensionArray.
3335 Always returns `ExtensionArray` even if only one value.
3337 Parameters
3338 ----------
3339 dropna : bool, default True
3340 Don't consider counts of NA values.
3342 Returns
3343 -------
3344 same type as self
3345 Sorted, if possible.
3346 """
3347 pa_type = self._pa_array.type
3348 if pa.types.is_temporal(pa_type):
3349 nbits = pa_type.bit_width
3350 if nbits == 32:
3351 data = self._pa_array.cast(pa.int32())
3352 elif nbits == 64:
3353 data = self._pa_array.cast(pa.int64())
3354 else:
3355 raise NotImplementedError(pa_type)
3356 else:
3357 data = self._pa_array
3359 if dropna:
3360 data = data.drop_null()
3362 res = pc.value_counts(data)
3363 most_common = res.field("values").filter(
3364 pc.equal(res.field("counts"), pc.max(res.field("counts")))
3365 )
3367 if pa.types.is_temporal(pa_type):
3368 most_common = most_common.cast(pa_type)
3370 most_common = most_common.take(pc.array_sort_indices(most_common))
3371 return self._from_pyarrow_array(most_common)
3373 def _validate_setitem_value(self, value):
3374 """Maybe convert value to be pyarrow compatible."""
3375 self_temporal = _is_temporal_pa_type(self._pa_array.type)
3376 if self_temporal is not None:
3377 value_temporal = _is_temporal_value(value)
3378 if value_temporal is not None and value_temporal != self_temporal:
3379 # an all-NA value carries nothing to reinterpret; only the
3380 # temporal direction is rejected, matching int64/Int64, see
3381 # test_setitem_all_na_temporal_array_still_raises
3382 if value_temporal or not _is_all_na(value):
3383 msg = f"Invalid value '{value!s}' for dtype '{self.dtype}'"
3384 raise TypeError(msg)
3385 if is_list_like(value) and getattr(value, "ndim", 1) == 1:
3386 # _box_pa would still cast, and e.g. double -> timestamp has
3387 # no cast kernel. len() is the element count only for a 1-D
3388 # value, so anything else stays with _box_pa
3389 return pa.nulls(len(value), type=self._pa_array.type)
3390 try:
3391 value = self._box_pa(value, self._pa_array.type)
3392 except pa.ArrowTypeError as err:
3393 msg = f"Invalid value '{value!s}' for dtype '{self.dtype}'"
3394 raise TypeError(msg) from err
3395 return value
3397 def interpolate(
3398 self,
3399 *,
3400 method: InterpolateOptions,
3401 axis: int,
3402 index,
3403 limit,
3404 limit_direction,
3405 limit_area,
3406 copy: bool,
3407 **kwargs,
3408 ) -> Self:
3409 """
3410 See NDFrame.interpolate.__doc__.
3411 """
3412 # NB: we return type(self) even if copy=False
3413 if not self.dtype._is_numeric:
3414 raise TypeError(f"Cannot interpolate with {self.dtype} dtype")
3416 # GH#65345: a pyarrow-native fast path for
3417 # method="linear"/limit_direction="forward" was removed here because
3418 # it only handled isolated NAs (leaving consecutive and trailing NAs
3419 # unfilled), truncated interpolated values for integer dtypes, and
3420 # did not upcast to float64 like the general path below.
3422 mask = self.isna()
3423 if self.dtype.kind == "f":
3424 data = self._pa_array.to_numpy()
3425 elif self.dtype.kind in "iu":
3426 data = self.to_numpy(dtype="f8", na_value=0.0)
3427 else:
3428 raise NotImplementedError(
3429 f"interpolate is not implemented for dtype={self.dtype}"
3430 )
3432 missing.interpolate_2d_inplace(
3433 data,
3434 method=method,
3435 axis=0,
3436 index=index,
3437 limit=limit,
3438 limit_direction=limit_direction,
3439 limit_area=limit_area,
3440 mask=mask,
3441 **kwargs,
3442 )
3443 return self._from_pyarrow_array(self._box_pa_array(pa.array(data, mask=mask)))
3445 @classmethod
3446 def _if_else(
3447 cls,
3448 cond: npt.NDArray[np.bool_] | bool,
3449 left: ArrayLike | Scalar,
3450 right: ArrayLike | Scalar,
3451 ) -> pa.Array:
3452 """
3453 Choose values based on a condition.
3455 Analogous to pyarrow.compute.if_else, with logic
3456 to fallback to numpy for unsupported types.
3458 Parameters
3459 ----------
3460 cond : npt.NDArray[np.bool_] or bool
3461 left : ArrayLike | Scalar
3462 right : ArrayLike | Scalar
3464 Returns
3465 -------
3466 pa.Array
3467 """
3469 # TODO: Remove this part when pa.if_else is fixed (GH#64320)
3470 def _maybe_combine(arr):
3471 if not isinstance(arr, pa.ChunkedArray) or not _is_varbinary_type(arr.type):
3472 return arr
3473 if not any(chunk.offset != 0 for chunk in arr.chunks):
3474 return arr
3475 try:
3476 return arr.combine_chunks()
3477 except (pa.ArrowInvalid, pa.ArrowCapacityError, MemoryError):
3478 return None
3480 left_c, right_c = _maybe_combine(left), _maybe_combine(right)
3481 if left_c is not None and right_c is not None:
3482 try:
3483 return pc.if_else(cond, left_c, right_c)
3484 except pa.ArrowNotImplementedError:
3485 pass
3486 if left_c is not None:
3487 left = left_c
3488 if right_c is not None:
3489 right = right_c
3491 def _to_numpy_and_type(value) -> tuple[np.ndarray, pa.DataType | None]:
3492 if isinstance(value, (pa.Array, pa.ChunkedArray)):
3493 pa_type = value.type
3494 elif isinstance(value, pa.Scalar):
3495 pa_type = value.type
3496 value = value.as_py()
3497 else:
3498 pa_type = None
3499 return np.array(value, dtype=object), pa_type
3501 left, left_type = _to_numpy_and_type(left)
3502 right, right_type = _to_numpy_and_type(right)
3503 pa_type = left_type or right_type
3504 result = np.where(cond, left, right)
3505 return pa.array(result, type=pa_type, from_pandas=is_nan_na())
3507 @classmethod
3508 def _replace_with_mask(
3509 cls,
3510 values: pa.Array | pa.ChunkedArray,
3511 mask: npt.NDArray[np.bool_] | bool,
3512 replacements: ArrayLike | Scalar,
3513 ) -> pa.Array | pa.ChunkedArray:
3514 """
3515 Replace items selected with a mask.
3517 Analogous to pyarrow.compute.replace_with_mask, with logic
3518 to fallback to numpy for unsupported types.
3520 Parameters
3521 ----------
3522 values : pa.Array or pa.ChunkedArray
3523 mask : npt.NDArray[np.bool_] or bool
3524 replacements : ArrayLike or Scalar
3525 Replacement value(s)
3527 Returns
3528 -------
3529 pa.Array or pa.ChunkedArray
3530 """
3531 if isinstance(replacements, pa.ChunkedArray):
3532 # replacements must be array or scalar, not ChunkedArray
3533 replacements = replacements.combine_chunks()
3534 if isinstance(values, pa.ChunkedArray) and pa.types.is_boolean(values.type):
3535 # GH#52059 replace_with_mask segfaults for chunked array
3536 # https://github.com/apache/arrow/issues/34634
3537 values = values.combine_chunks()
3538 if pa_version_under25p0 and pa.types.is_null(values.type):
3539 # GH#66703 pc.replace_with_mask aborts the process for null dtype
3540 # (apache/arrow#47447, fixed in pyarrow 25.0.0). The null type
3541 # has no non-null values to replace; return it unchanged.
3542 return values
3543 try:
3544 return pc.replace_with_mask(values, mask, replacements)
3545 except pa.ArrowNotImplementedError:
3546 pass
3547 if isinstance(replacements, pa.Array):
3548 replacements = np.array(replacements, dtype=object)
3549 elif isinstance(replacements, pa.Scalar):
3550 replacements = replacements.as_py()
3552 result = np.array(values, dtype=object)
3553 result[mask] = replacements
3554 return pa.array(result, type=values.type, from_pandas=is_nan_na())
3556 # ------------------------------------------------------------------
3557 # GroupBy Methods
3559 def _to_masked(self):
3560 pa_dtype = self._pa_array.type
3562 if pa.types.is_floating(pa_dtype) or pa.types.is_integer(pa_dtype):
3563 na_value = 1
3564 elif pa.types.is_boolean(pa_dtype):
3565 na_value = True
3566 else:
3567 raise NotImplementedError
3569 dtype = _arrow_dtype_mapping()[pa_dtype]
3570 mask = self.isna()
3571 arr = self.to_numpy(dtype=dtype.numpy_dtype, na_value=na_value)
3572 return dtype.construct_array_type()(arr, mask)
3574 # pandas groupby 'how' -> PyArrow aggregation function name
3575 _PYARROW_AGG_FUNCS: dict[str, str] = {
3576 "sum": "sum",
3577 "prod": "product",
3578 "min": "min",
3579 "max": "max",
3580 "mean": "mean",
3581 "std": "stddev",
3582 "var": "variance",
3583 "sem": "stddev", # sem = stddev / sqrt(count)
3584 }
3586 # Identity elements for operations (used to fill missing groups)
3587 _PYARROW_AGG_DEFAULTS: dict[str, int] = {
3588 "sum": 0,
3589 "prod": 1,
3590 }
3592 def _groupby_op_pyarrow(
3593 self,
3594 *,
3595 how: str,
3596 min_count: int,
3597 ngroups: int,
3598 ids: npt.NDArray[np.intp],
3599 **kwargs,
3600 ) -> Self | None:
3601 """
3602 Perform groupby aggregation using PyArrow's native Table.group_by.
3604 Returns None if not supported, caller should fall back to Cython path.
3605 """
3606 pa_agg_func = self._PYARROW_AGG_FUNCS.get(how)
3607 if pa_agg_func is None:
3608 return None
3610 # See _groupby_op for the types routed here.
3611 # PyArrow doesn't support sum/prod/mean/std/var/sem on strings.
3612 pa_type = self._pa_array.type
3613 is_str = pa.types.is_string(pa_type) or pa.types.is_large_string(pa_type)
3614 if is_str and how in ["sum", "prod", "mean", "std", "var", "sem"]:
3615 return None
3617 # Filter out NA group (ids == -1)
3618 mask = ids >= 0
3619 if not mask.all():
3620 ids = ids[mask]
3621 values = pc.filter(self._pa_array, mask)
3622 else:
3623 values = self._pa_array
3625 # Build table and run aggregation (cast ids to int64 for portability)
3626 group_id_arr = pa.array(ids, type=pa.int64())
3627 table = pa.table({"value": values, "group_id": group_id_arr})
3629 # When skipna=False, propagate nulls: any group containing a null
3630 # should aggregate to null. PyArrow skips nulls by default.
3631 skipna = kwargs.get("skipna", True)
3632 agg_option: pc.VarianceOptions | pc.ScalarAggregateOptions
3633 if how in ["std", "var", "sem"]:
3634 ddof = kwargs.get("ddof", 1)
3635 agg_option = pc.VarianceOptions(ddof=ddof, skip_nulls=skipna)
3636 else:
3637 agg_option = pc.ScalarAggregateOptions(skip_nulls=skipna)
3638 aggs: list[tuple[str, str] | tuple[str, str, pc.FunctionOptions]] = [
3639 ("value", pa_agg_func, agg_option)
3640 ]
3641 # Counts non-null values only; needed to scale sem and to apply min_count.
3642 needs_count = how == "sem" or min_count > 0
3643 if needs_count:
3644 aggs.append(("value", "count"))
3646 result_table = table.group_by("group_id").aggregate(aggs)
3647 result_group_ids = result_table.column("group_id")
3648 result_values = result_table.column(f"value_{pa_agg_func}")
3650 if how == "sem":
3651 result_values = pc.divide(
3652 result_values, pc.sqrt(result_table.column("value_count"))
3653 )
3654 elif how in ["sum", "prod"] and pa.types.is_decimal(result_values.type):
3655 # A sum or a product needs more digits than the input holds, but
3656 # hash_product keeps the input precision, and so does hash_sum
3657 # before PyArrow 21, which leaves a value too large for its own
3658 # declared type. Widen to the maximum precision, which is what
3659 # PyArrow 21 does for hash_sum and what Series.sum and Series.prod
3660 # return.
3661 if how == "prod" or pa_version_under21p0:
3662 agg_type = result_values.type
3663 if pa.types.is_decimal256(agg_type):
3664 wider_type = pa.decimal256(76, agg_type.scale)
3665 else:
3666 wider_type = pa.decimal128(38, agg_type.scale)
3667 try:
3668 result_values = result_values.cast(wider_type)
3669 except pa.ArrowInvalid:
3670 # The result needs more digits than the maximum precision,
3671 # so no Arrow type can hold it. Fall back.
3672 return None
3674 output_type = result_values.type
3675 default_value = pa.scalar(self._PYARROW_AGG_DEFAULTS.get(how), type=output_type)
3677 # Replace nulls from all-null groups with identity element.
3678 # Skipped when skipna=False, where a null must propagate as null.
3679 if (
3680 skipna
3681 and result_values.null_count > 0
3682 and how in ["sum", "prod"]
3683 and min_count == 0
3684 ):
3685 result_values = pc.if_else(
3686 pc.is_null(result_values), default_value, result_values
3687 )
3689 # Null out groups below min_count
3690 if min_count > 0:
3691 below_min_count = pc.less(
3692 result_table.column("value_count"), pa.scalar(min_count)
3693 )
3694 # _if_else works around pc.if_else on chunked strings (GH#64320)
3695 result_values = self._if_else(
3696 below_min_count.to_numpy(),
3697 pa.scalar(None, type=result_values.type),
3698 result_values,
3699 )
3701 # Place the results in group-id order: the inverse permutation takes
3702 # the row holding group i, and is null where group i had no rows.
3703 group_ids_np = result_group_ids.to_numpy(zero_copy_only=False)
3704 inverse = np.full(ngroups, -1, dtype=np.int64)
3705 inverse[group_ids_np] = np.arange(len(group_ids_np))
3706 indices = pa.array(inverse, mask=inverse < 0)
3708 if how in ["sum", "prod"] and pa.types.is_decimal(output_type):
3709 try:
3710 # take would carry an out-of-precision decimal through silently
3711 result_values.validate(full=True)
3712 except pa.ArrowInvalid:
3713 # needs more digits than the maximum precision, so let the
3714 # caller fall back to a type that can hold it
3715 return None
3717 pa_result = pc.take(result_values, indices)
3718 if default_value.as_py() is not None and min_count == 0:
3719 if result_values.null_count == 0:
3720 # every null is a group with no rows
3721 pa_result = _safe_fill_null(pa_result, default_value)
3722 else:
3723 # keep the skipna=False nulls, fill only the empty groups
3724 pa_result = pc.if_else(pc.is_null(indices), default_value, pa_result)
3725 return self._from_pyarrow_array(pa_result)
3727 def _to_groupby_compatible(self) -> ExtensionArray:
3728 """Convert to a type compatible with groupby cython ops."""
3729 pa_type = self._pa_array.type
3730 if pa.types.is_timestamp(pa_type):
3731 return self._to_datetimearray()
3732 elif pa.types.is_duration(pa_type):
3733 return self._to_timedeltaarray()
3734 else:
3735 return self._to_masked()
3737 def _groupby_result_to_arrow(self, result: ArrayLike) -> ArrayLike:
3738 """Convert a groupby result from a delegated type back to Arrow."""
3739 if isinstance(result, np.ndarray):
3740 return result
3741 elif isinstance(result, BaseMaskedArray):
3742 pa_result = result.__arrow_array__()
3743 return self._from_pyarrow_array(pa_result)
3744 else:
3745 # DatetimeArray, TimedeltaArray
3746 pa_result = pa.array(result)
3747 return self._from_pyarrow_array(pa_result)
3749 def _groupby_op(
3750 self,
3751 *,
3752 how: str,
3753 has_dropped_na: bool,
3754 min_count: int,
3755 ngroups: int,
3756 ids: npt.NDArray[np.intp],
3757 **kwargs,
3758 ):
3759 if isinstance(self.dtype, StringDtype):
3760 if how in [
3761 "prod",
3762 "mean",
3763 "median",
3764 "cumsum",
3765 "cumprod",
3766 "std",
3767 "sem",
3768 "var",
3769 "skew",
3770 ]:
3771 raise TypeError(
3772 f"dtype '{self.dtype}' does not support operation '{how}'"
3773 )
3774 # Fall through to Arrow-native path below
3776 pa_type = self._pa_array.type
3778 # Try PyArrow-native path for decimal and string types where it's faster.
3779 # For integer/float/boolean, the fallback path via _to_masked() is faster.
3780 # Date, time, binary and null types have no Cython path for min/max.
3781 if (
3782 pa.types.is_decimal(pa_type)
3783 or pa.types.is_string(pa_type)
3784 or pa.types.is_large_string(pa_type)
3785 or (
3786 how in ["min", "max"]
3787 and (
3788 pa.types.is_date(pa_type)
3789 or pa.types.is_time(pa_type)
3790 or _is_varbinary_type(pa_type)
3791 or pa.types.is_fixed_size_binary(pa_type)
3792 or pa.types.is_null(pa_type)
3793 )
3794 )
3795 ):
3796 native_result = self._groupby_op_pyarrow(
3797 how=how,
3798 min_count=min_count,
3799 ngroups=ngroups,
3800 ids=ids,
3801 **kwargs,
3802 )
3803 if native_result is not None:
3804 return native_result
3805 # For string types, fall back to parent implementation (Python path)
3806 # since _to_masked() doesn't support strings
3807 if pa.types.is_string(pa_type) or pa.types.is_large_string(pa_type):
3808 return super()._groupby_op(
3809 how=how,
3810 has_dropped_na=has_dropped_na,
3811 min_count=min_count,
3812 ngroups=ngroups,
3813 ids=ids,
3814 **kwargs,
3815 )
3817 if how in ["first", "last"]:
3818 return self._groupby_first_last(
3819 how=how,
3820 min_count=min_count,
3821 ngroups=ngroups,
3822 ids=ids,
3823 skipna=kwargs.get("skipna", True),
3824 )
3826 # Fall back to converting to a groupby-compatible delegate
3827 # (masked/datetime/timedelta array) and running the Cython path.
3828 values = self._to_groupby_compatible()
3829 result = values._groupby_op(
3830 how=how,
3831 has_dropped_na=has_dropped_na,
3832 min_count=min_count,
3833 ngroups=ngroups,
3834 ids=ids,
3835 **kwargs,
3836 )
3837 return self._groupby_result_to_arrow(result)
3839 def _groupby_quantile(
3840 self,
3841 *,
3842 qs: npt.NDArray[np.float64],
3843 interpolation: Literal["linear", "lower", "higher", "nearest", "midpoint"],
3844 ids: npt.NDArray[np.intp],
3845 ngroups: int,
3846 starts: npt.NDArray[np.int64],
3847 ends: npt.NDArray[np.int64],
3848 ) -> ArrayLike:
3849 from pandas.core.arrays.string_ import StringDtype
3851 if isinstance(self.dtype, StringDtype):
3852 raise TypeError(
3853 f"dtype '{self.dtype}' does not support operation 'quantile'"
3854 )
3856 values = self._to_groupby_compatible()
3857 result = values._groupby_quantile(
3858 qs=qs,
3859 interpolation=interpolation,
3860 ids=ids,
3861 ngroups=ngroups,
3862 starts=starts,
3863 ends=ends,
3864 )
3865 return self._groupby_result_to_arrow(result)
3867 def _apply_elementwise(self, func: Callable[..., Any]) -> list[list[Any]]:
3868 """Apply a callable to each element while maintaining the chunking structure."""
3869 return [
3870 [
3871 None if val is None else func(val)
3872 for val in chunk.to_numpy(zero_copy_only=False)
3873 ]
3874 for chunk in self._pa_array.iterchunks()
3875 ]
3877 def _convert_bool_result(self, result, na=lib.no_default, method_name=None):
3878 if na is not lib.no_default and not isna(na): # pyright: ignore [reportGeneralTypeIssues]
3879 result = result.fill_null(na)
3880 return self._from_pyarrow_array(result)
3882 def _convert_int_result(self, result):
3883 return self._from_pyarrow_array(result)
3885 def _convert_rank_result(self, result):
3886 return self._from_pyarrow_array(result)
3888 @staticmethod
3889 def _compile_re_fallback(
3890 pat: str | re.Pattern[str], case: bool = True, flags: int = 0
3891 ) -> re.Pattern[str]:
3892 # GH#66348 pyarrow's regex kernels honor no flags beyond the IGNORECASE
3893 # that `case` stands in for, so anything left over is evaluated with `re`
3894 if not case:
3895 flags |= re.IGNORECASE
3896 # as in ObjectStringArrayMixin, an already-compiled `pat` combined with
3897 # `flags` raises out of re.compile
3898 return re.compile(pat, flags=flags)
3900 def _apply_re_fallback(self, func: Callable[..., Any], pa_type: pa.DataType):
3901 return pa.chunked_array(self._apply_elementwise(func), type=pa_type)
3903 def _str_contains(
3904 self,
3905 pat: str | re.Pattern[str],
3906 case: bool = True,
3907 flags: int = 0,
3908 na: Scalar | lib.NoDefault = lib.no_default,
3909 regex: bool = True,
3910 ):
3911 if regex:
3912 pat, case, flags = self._unwrap_re_pattern(pat, case, flags)
3914 if flags:
3915 if regex:
3916 compiled = self._compile_re_fallback(pat, case, flags)
3917 func = lambda val: compiled.search(val) is not None
3918 elif case:
3919 func = lambda val: pat in val
3920 else:
3921 upper_pat = pat.upper() # type: ignore[union-attr]
3922 func = lambda val: upper_pat in val.upper()
3923 result = self._apply_re_fallback(func, pa.bool_())
3924 return self._convert_bool_result(result, na=na, method_name="contains")
3926 return ArrowStringArrayMixin._str_contains(self, pat, case, flags, na, regex)
3928 def _str_match(
3929 self,
3930 pat: str | re.Pattern[str],
3931 case: bool = True,
3932 flags: int = 0,
3933 na: Scalar | lib.NoDefault = lib.no_default,
3934 ):
3935 pat, case, flags = self._unwrap_re_pattern(pat, case, flags)
3937 if flags:
3938 compiled = self._compile_re_fallback(pat, case, flags)
3939 func = lambda val: compiled.match(val) is not None
3940 result = self._apply_re_fallback(func, pa.bool_())
3941 return self._convert_bool_result(result, na=na, method_name="match")
3943 return ArrowStringArrayMixin._str_match(self, pat, case, flags, na)
3945 def _str_fullmatch(
3946 self,
3947 pat: str | re.Pattern[str],
3948 case: bool = True,
3949 flags: int = 0,
3950 na: Scalar | lib.NoDefault = lib.no_default,
3951 ):
3952 pat, case, flags = self._unwrap_re_pattern(pat, case, flags)
3954 if flags:
3955 # anchoring the pattern and deferring to _str_match would not do:
3956 # under re.MULTILINE the added "^"/"$" match at line boundaries
3957 compiled = self._compile_re_fallback(pat, case, flags)
3958 func = lambda val: compiled.fullmatch(val) is not None
3959 result = self._apply_re_fallback(func, pa.bool_())
3960 return self._convert_bool_result(result, na=na, method_name="fullmatch")
3962 return ArrowStringArrayMixin._str_fullmatch(self, pat, case, flags, na)
3964 def _str_count(self, pat: str | re.Pattern[str], flags: int = 0) -> Self:
3965 pat, case, flags = self._unwrap_re_pattern(pat, True, flags)
3967 if flags:
3968 compiled = self._compile_re_fallback(pat, case, flags)
3969 func = lambda val: len(compiled.findall(val))
3970 # match the width pyarrow's kernel reports for this storage type
3971 pa_type = (
3972 pa.int64()
3973 if pa.types.is_large_string(self._pa_array.type)
3974 or pa.types.is_large_binary(self._pa_array.type)
3975 else pa.int32()
3976 )
3977 return self._convert_int_result(self._apply_re_fallback(func, pa_type))
3979 result = pc.count_substring_regex(self._pa_array, pat, ignore_case=not case)
3980 return self._convert_int_result(result)
3982 def _str_replace(
3983 self,
3984 pat: str | re.Pattern[str],
3985 repl: str | Callable[..., Any],
3986 n: int = -1,
3987 case: bool = True,
3988 flags: int = 0,
3989 regex: bool = True,
3990 ) -> Self:
3991 if (
3992 isinstance(pat, re.Pattern)
3993 or callable(repl)
3994 or not case
3995 or flags
3996 or (isinstance(repl, str) and r"\g<" in repl)
3997 ):
3998 # None of these are expressible with pc.replace_substring_regex; mirror
3999 # ObjectStringArrayMixin._str_replace instead, which leaves the flags
4000 # of an already-compiled `pat` alone. GH#66348
4001 if not isinstance(pat, re.Pattern):
4002 if not regex:
4003 pat = re.escape(pat)
4004 pat = self._compile_re_fallback(pat, case, flags)
4005 count = n if n >= 0 else 0
4006 func = lambda val: pat.sub(repl=repl, string=val, count=count)
4007 return self._from_pyarrow_array(
4008 self._apply_re_fallback(func, self._pa_array.type)
4009 )
4011 return ArrowStringArrayMixin._str_replace(
4012 self, pat, repl, n, case, flags, regex
4013 )
4015 def _str_repeat(self, repeats: int | Sequence[int]) -> Self:
4016 if not isinstance(repeats, int):
4017 raise NotImplementedError(
4018 f"repeat is not implemented when repeats is {type(repeats).__name__}"
4019 )
4020 return self._from_pyarrow_array(pc.binary_repeat(self._pa_array, repeats))
4022 def _str_join(self, sep: str) -> Self:
4023 pa_type = self._pa_array.type
4024 if pa.types.is_string(pa_type) or pa.types.is_large_string(pa_type):
4025 result = self._apply_elementwise(list)
4026 # pc.binary_join has no kernel for large_string values, so join as
4027 # string and restore the input's own type afterwards
4028 listed = pa.chunked_array(result, type=pa.list_(pa.string()))
4029 return self._from_pyarrow_array(pc.binary_join(listed, sep).cast(pa_type))
4030 return self._from_pyarrow_array(pc.binary_join(self._pa_array, sep))
4032 def _str_partition(self, sep: str, expand: bool):
4033 if expand:
4034 # rows of three strings, so a list array rather than Self
4035 return ArrowExtensionArray(
4036 ArrowStringArrayMixin._str_partition_expand(self, sep)
4037 )
4038 predicate = lambda val: val.partition(sep)
4039 result = self._apply_elementwise(predicate)
4040 return self._from_pyarrow_array(
4041 pa.chunked_array(result, type=pa.list_(self._pa_array.type))
4042 )
4044 def _str_rpartition(self, sep: str, expand: bool) -> Self:
4045 predicate = lambda val: val.rpartition(sep)
4046 result = self._apply_elementwise(predicate)
4047 return self._from_pyarrow_array(
4048 pa.chunked_array(result, type=pa.list_(self._pa_array.type))
4049 )
4051 def _str_casefold(self) -> Self:
4052 predicate = lambda val: val.casefold()
4053 result = self._apply_elementwise(predicate)
4054 return self._from_pyarrow_array(
4055 pa.chunked_array(result, type=self._pa_array.type)
4056 )
4058 def _str_encode(self, encoding: str, errors: str = "strict") -> Self:
4059 predicate = lambda val: val.encode(encoding, errors)
4060 result = self._apply_elementwise(predicate)
4061 if pa.types.is_large_string(self._pa_array.type):
4062 pa_type = pa.large_binary()
4063 else:
4064 pa_type = pa.binary()
4065 return self._from_pyarrow_array(pa.chunked_array(result, type=pa_type))
4067 def _str_extract(
4068 self, pat: str | re.Pattern[str], flags: int = 0, expand: bool = True
4069 ):
4070 compiled = self._compile_re_fallback(pat, flags=flags)
4072 from pandas.core.strings.accessor import get_group_names
4074 labels: list[Hashable] = list(get_group_names(compiled))
4076 if (
4077 flags
4078 or isinstance(pat, re.Pattern)
4079 # GH#67024 pc.extract_regex addresses the struct it returns by field name,
4080 # so a pattern with any unnamed group is evaluated with `re` instead
4081 or len(compiled.groupindex) != compiled.groups
4082 ):
4083 # pc.extract_regex has no `ignore_case`, so unlike elsewhere not even
4084 # IGNORECASE can be honored by the kernel
4085 matches = self._apply_elementwise(compiled.search)
4087 def extract_group(num: int):
4088 chunks = [
4089 [None if match is None else match.group(num) for match in chunk]
4090 for chunk in matches
4091 ]
4092 return self._from_pyarrow_array(
4093 pa.chunked_array(chunks, type=self._pa_array.type)
4094 )
4096 if not expand:
4097 return extract_group(1)
4098 return {label: extract_group(1 + i) for i, label in enumerate(labels)}
4100 result = pc.extract_regex(self._pa_array, pat)
4101 if expand:
4102 return {
4103 label: self._from_pyarrow_array(pc.struct_field(result, [i]))
4104 for i, label in enumerate(labels)
4105 }
4106 else:
4107 return type(self)(pc.struct_field(result, [0]))
4109 def _str_findall(self, pat: str, flags: int = 0) -> Self:
4110 regex = re.compile(pat, flags=flags)
4111 predicate = lambda val: regex.findall(val)
4112 result = self._apply_elementwise(predicate)
4113 return self._from_pyarrow_array(
4114 pa.chunked_array(result, type=pa.list_(self._pa_array.type))
4115 )
4117 def _str_get_dummies(self, sep: str = "|", dtype: NpDtype | None = None):
4118 if dtype is None:
4119 dtype = np.bool_
4120 split = pc.split_pattern(self._pa_array, sep)
4121 flattened_values = pc.list_flatten(split)
4122 uniques = flattened_values.unique()
4123 uniques_sorted = uniques.take(pa.compute.array_sort_indices(uniques))
4124 lengths = pc.list_value_length(split).fill_null(0).to_numpy()
4125 n_rows = len(self)
4126 n_cols = len(uniques)
4127 indices = pc.index_in(flattened_values, uniques_sorted).to_numpy()
4128 indices = indices + np.arange(n_rows).repeat(lengths) * n_cols
4129 _dtype = pandas_dtype(dtype)
4130 dummies_dtype: NpDtype
4131 if isinstance(_dtype, np.dtype):
4132 dummies_dtype = _dtype
4133 else:
4134 dummies_dtype = np.bool_
4135 dummies = np.zeros(n_rows * n_cols, dtype=dummies_dtype)
4136 dummies[indices] = True
4137 dummies = dummies.reshape((n_rows, n_cols))
4138 result = self._from_pyarrow_array(pa.array(list(dummies)))
4139 return result, uniques_sorted.to_pylist()
4141 def _str_index(self, sub: str, start: int = 0, end: int | None = None) -> Self:
4142 predicate = lambda val: val.index(sub, start, end)
4143 result = self._apply_elementwise(predicate)
4144 return self._from_pyarrow_array(pa.chunked_array(result, type=pa.int64()))
4146 def _str_rindex(self, sub: str, start: int = 0, end: int | None = None) -> Self:
4147 predicate = lambda val: val.rindex(sub, start, end)
4148 result = self._apply_elementwise(predicate)
4149 return self._from_pyarrow_array(pa.chunked_array(result, type=pa.int64()))
4151 def _str_rfind(self, sub: str, start: int = 0, end=None) -> Self:
4152 predicate = lambda val: val.rfind(sub, start, end)
4153 result = self._apply_elementwise(predicate)
4154 return self._from_pyarrow_array(pa.chunked_array(result, type=pa.int64()))
4156 def _str_split(
4157 self,
4158 pat: str | None = None,
4159 n: int | None = -1,
4160 expand: bool = False,
4161 regex: bool | None = None,
4162 ) -> Self:
4163 if n in {-1, 0}:
4164 n = None
4165 if pat is None:
4166 split_func = pc.utf8_split_whitespace
4167 elif regex is True:
4168 split_func = functools.partial(pc.split_pattern_regex, pattern=pat)
4169 elif regex is False:
4170 split_func = functools.partial(pc.split_pattern, pattern=pat)
4171 # GH#58321: regex is None — infer: single-char literal, multi-char regex
4172 elif len(pat) == 1:
4173 split_func = functools.partial(pc.split_pattern, pattern=pat)
4174 else:
4175 split_func = functools.partial(pc.split_pattern_regex, pattern=pat)
4176 return self._from_pyarrow_array(split_func(self._pa_array, max_splits=n))
4178 def _str_rsplit(self, pat: str | None = None, n: int | None = -1) -> Self:
4179 if pat is not None and not isinstance(pat, str):
4180 msg = f"expected a string object, not {type(pat).__name__}"
4181 raise TypeError(msg)
4182 if n in {-1, 0}:
4183 n = None
4184 if pat is None:
4185 return self._from_pyarrow_array(
4186 pc.utf8_split_whitespace(self._pa_array, max_splits=n, reverse=True)
4187 )
4188 return self._from_pyarrow_array(
4189 pc.split_pattern(self._pa_array, pat, max_splits=n, reverse=True)
4190 )
4192 def _str_translate(self, table: dict[int, str]) -> Self:
4193 predicate = lambda val: val.translate(table)
4194 result = self._apply_elementwise(predicate)
4195 return self._from_pyarrow_array(
4196 pa.chunked_array(result, type=self._pa_array.type)
4197 )
4199 def _str_wrap(self, width: int, **kwargs) -> Self:
4200 kwargs["width"] = width
4201 tw = textwrap.TextWrapper(**kwargs)
4202 predicate = lambda val: "\n".join(tw.wrap(val))
4203 result = self._apply_elementwise(predicate)
4204 return self._from_pyarrow_array(
4205 pa.chunked_array(result, type=self._pa_array.type)
4206 )
4208 def _dt_zero_or_null_int32(self) -> Self:
4209 """
4210 Return an int32 result that is 0 where the input is valid and null
4211 where the input is null.
4213 Used for duration components that are always 0 because the array's
4214 resolution is coarser than the component (e.g. nanoseconds of a
4215 second-resolution duration).
4216 """
4217 return self._from_pyarrow_array(
4218 pc.if_else(
4219 pc.is_null(self._pa_array),
4220 pa.scalar(None, type=pa.int32()),
4221 pa.scalar(0, type=pa.int32()),
4222 )
4223 )
4225 @property
4226 def _duration_unit(self) -> str:
4227 """
4228 Return the time unit for a duration-typed array.
4230 Raises ``AttributeError`` for non-duration dtypes so the duration
4231 component accessors (``_dt_days``, ``_dt_seconds``, etc.) are rejected
4232 for e.g. ``timestamp[pyarrow]`` instead of silently treating the
4233 underlying int64 as a duration. The ``dt`` dispatcher turns this into
4234 the usual "dt.<name> is not supported for <dtype>" error.
4235 """
4236 pa_type = self._pa_array.type
4237 if not pa.types.is_duration(pa_type):
4238 raise AttributeError(f"'{self.dtype}' does not have duration components")
4239 return pa_type.unit
4241 @cache_readonly
4242 def _dt_day_remainder(self) -> pa.ChunkedArray:
4243 """
4244 Return the remainder after removing full days, always non-negative.
4246 For negative durations like -22h 57m 57s (= -1 day + 1h 2m 3s),
4247 this returns the positive offset from the day boundary.
4249 This is cached because it's used by all sub-day component accessors.
4250 """
4251 unit = self._duration_unit
4252 divisor = _DURATION_DIVISORS["day"][unit]
4253 arr = self._pa_array.cast(pa.int64())
4254 days = floor_div_int64(arr, divisor)
4255 # remainder = arr - days * divisor (always non-negative)
4256 return pc.subtract(arr, pc.multiply(days, divisor))
4258 def _dt_subday_component(self, component: str, modulo: int) -> Self:
4259 """
4260 Extract a sub-day component (hours, minutes, etc.) from the day remainder.
4262 The remainder is always non-negative, so we can use integer division.
4263 """
4264 unit = self._duration_unit
4265 divisors = _DURATION_DIVISORS[component]
4266 # For units coarser than this component, the result is always 0
4267 if unit not in divisors:
4268 return self._dt_zero_or_null_int32()
4269 # Integer division is fine here since remainder is non-negative
4270 result = pc.divide(self._dt_day_remainder, divisors[unit])
4271 result = mod_int(result, modulo)
4272 return self._from_pyarrow_array(result.cast(pa.int32()))
4274 @property
4275 def _dt_days(self) -> Self:
4276 unit = self._duration_unit
4277 arr = self._pa_array.cast(pa.int64())
4278 result = floor_div_int64(arr, _DURATION_DIVISORS["day"][unit])
4279 return self._from_pyarrow_array(result.cast(pa.int32()))
4281 @property
4282 def _dt_hours(self) -> Self:
4283 return self._dt_subday_component("hour", modulo=24)
4285 @property
4286 def _dt_minutes(self) -> Self:
4287 return self._dt_subday_component("minute", modulo=60)
4289 @property
4290 def _dt_seconds(self) -> Self:
4291 # Total seconds in sub-day portion (0-86399), per Python timedelta semantics.
4292 # Every duration unit (s/ms/us/ns) is at least as fine as seconds, so the
4293 # remainder always has enough resolution and no zero-result shortcut is
4294 # needed here.
4295 unit = self._duration_unit
4296 # Convert remainder to seconds (no modulo - we want total sub-day seconds)
4297 result = pc.divide(self._dt_day_remainder, _DURATION_DIVISORS["second"][unit])
4298 return self._from_pyarrow_array(result.cast(pa.int32()))
4300 @property
4301 def _dt_milliseconds(self) -> Self:
4302 return self._dt_subday_component("millisecond", modulo=1000)
4304 @property
4305 def _dt_microseconds(self) -> Self:
4306 # Total microseconds in sub-second portion (0-999999), per Python timedelta
4307 # semantics. For units coarser than microseconds (e.g. "ms") the sub-second
4308 # portion must be scaled *up* to microseconds rather than returning 0.
4309 unit = self._duration_unit
4310 divisors = _DURATION_DIVISORS
4311 # Sub-second portion, expressed in the storage unit (always non-negative)
4312 sub_second = mod_int(self._dt_day_remainder, divisors["second"][unit])
4313 units_per_second = divisors["second"][unit]
4314 micros_per_second = divisors["second"]["us"]
4315 if units_per_second >= micros_per_second:
4316 # Unit at least as fine as microseconds: divide down (e.g. ns -> us)
4317 result = pc.divide(sub_second, units_per_second // micros_per_second)
4318 else:
4319 # Unit coarser than microseconds: scale up (e.g. ms -> us)
4320 result = pc.multiply(sub_second, micros_per_second // units_per_second)
4321 return self._from_pyarrow_array(result.cast(pa.int32()))
4323 @property
4324 def _dt_nanoseconds(self) -> Self:
4325 # Nanoseconds only (0-999), per Python timedelta semantics
4326 return self._dt_subday_component("nanosecond", modulo=1000)
4328 @property
4329 def _dt_components(self) -> dict[str, ArrowExtensionArray]:
4330 """
4331 Return all duration components.
4333 The day remainder is computed once (and cached on ``_dt_day_remainder``),
4334 so every sub-day component below reuses it. Note that the per-field
4335 semantics differ from the like-named accessors: ``seconds`` here is the
4336 0-59 seconds field (not the 0-86399 total) and ``microseconds`` is the
4337 0-999 field (not the 0-999999 total).
4338 """
4339 return {
4340 "days": self._dt_days,
4341 "hours": self._dt_hours,
4342 "minutes": self._dt_minutes,
4343 "seconds": self._dt_subday_component("second", modulo=60),
4344 "milliseconds": self._dt_milliseconds,
4345 "microseconds": self._dt_subday_component("microsecond", modulo=1000),
4346 "nanoseconds": self._dt_nanoseconds,
4347 }
4349 def _dt_to_pytimedelta(self) -> np.ndarray:
4350 data = self._pa_array.to_pylist()
4351 if self._dtype.pyarrow_dtype.unit == "ns":
4352 data = [None if ts is None else ts.to_pytimedelta() for ts in data]
4353 return np.array(data, dtype=object)
4355 def _dt_total_seconds(self) -> Self:
4356 unit = self._pa_array.type.unit
4357 unit_per_second = {"s": 1.0, "ms": 1e3, "us": 1e6, "ns": 1e9}
4358 result = pc.divide(pc.cast(self._pa_array, pa.int64()), unit_per_second[unit])
4359 return self._from_pyarrow_array(result)
4361 def _dt_as_unit(self, unit: str) -> Self:
4362 pa_type = self._pa_array.type
4363 if pa.types.is_timestamp(pa_type):
4364 target_type = pa.timestamp(unit, tz=pa_type.tz)
4365 elif pa.types.is_duration(pa_type):
4366 target_type = pa.duration(unit)
4367 else:
4368 raise NotImplementedError(f"as_unit not implemented for {pa_type}")
4370 nanos_per_unit = {"s": 1_000_000_000, "ms": 1_000_000, "us": 1_000, "ns": 1}
4371 from_nanos = nanos_per_unit[pa_type.unit]
4372 to_nanos = nanos_per_unit[unit]
4373 if to_nanos <= from_nanos:
4374 # Same or finer resolution: exact upscale. Use safe=True so that
4375 # out-of-bounds values raise instead of silently wrapping, matching
4376 # numpy/pandas as_unit.
4377 try:
4378 result = pc.cast(self._pa_array, target_type)
4379 except pa.ArrowInvalid as err:
4380 err_type = (
4381 OutOfBoundsDatetime
4382 if pa.types.is_timestamp(pa_type)
4383 else OutOfBoundsTimedelta
4384 )
4385 raise err_type(
4386 f"Cannot convert {pa_type} to {target_type} without overflow"
4387 ) from err
4388 else:
4389 # Coarser resolution: floor toward -inf like numpy/pandas. pc.cast
4390 # truncates toward zero, which is wrong for negative timedeltas and
4391 # pre-epoch timestamps (GH#63573).
4392 ratio = pa.scalar(to_nanos // from_nanos, type=pa.int64())
4393 i8 = pc.cast(self._pa_array, pa.int64())
4394 quotient = pc.divide(i8, ratio)
4395 remainder = pc.subtract(i8, pc.multiply(quotient, ratio))
4396 floored = pc.if_else(
4397 pc.less(remainder, 0), pc.subtract(quotient, 1), quotient
4398 )
4399 result = pc.cast(floored, target_type)
4400 return self._from_pyarrow_array(result)
4402 @property
4403 def _dt_year(self) -> Self:
4404 result = pc.year(self._pa_array)
4405 return self._from_pyarrow_array(result)
4407 @property
4408 def _dt_day(self) -> Self:
4409 result = pc.day(self._pa_array)
4410 return self._from_pyarrow_array(result)
4412 @property
4413 def _dt_day_of_week(self) -> Self:
4414 result = pc.day_of_week(self._pa_array)
4415 return self._from_pyarrow_array(result)
4417 @property
4418 def _dt_dayofweek(self) -> Self:
4419 warnings.warn(
4420 "ArrowDtype dayofweek accessor is deprecated and will be removed "
4421 "in a future version. Use day_of_week instead.",
4422 Pandas4Warning,
4423 stacklevel=find_stack_level(),
4424 )
4425 return self._dt_day_of_week
4427 @property
4428 def _dt_weekday(self) -> Self:
4429 # GH#12816
4430 warnings.warn(
4431 "Series.dt.weekday is deprecated and will be removed "
4432 "in a future version. Use Series.dt.day_of_week instead.",
4433 Pandas4Warning,
4434 stacklevel=find_stack_level(),
4435 )
4436 return self._dt_day_of_week
4438 @property
4439 def _dt_day_of_year(self) -> Self:
4440 result = pc.day_of_year(self._pa_array)
4441 return self._from_pyarrow_array(result)
4443 @property
4444 def _dt_dayofyear(self) -> Self:
4445 warnings.warn(
4446 "ArrowDtype dayofyear accessor is deprecated and will be removed "
4447 "in a future version. Use day_of_year instead.",
4448 Pandas4Warning,
4449 stacklevel=find_stack_level(),
4450 )
4451 return self._dt_day_of_year
4453 @property
4454 def _dt_hour(self) -> Self:
4455 result = pc.hour(self._pa_array)
4456 return self._from_pyarrow_array(result)
4458 def _dt_isocalendar(self) -> Self:
4459 result = pc.iso_calendar(self._pa_array)
4460 return self._from_pyarrow_array(result)
4462 @property
4463 def _dt_is_leap_year(self) -> Self:
4464 result = pc.is_leap_year(self._pa_array)
4465 return self._from_pyarrow_array(result)
4467 @property
4468 def _dt_is_month_start(self) -> Self:
4469 result = pc.equal(pc.day(self._pa_array), 1)
4470 return self._from_pyarrow_array(result)
4472 @property
4473 def _dt_is_month_end(self) -> Self:
4474 result = pc.equal(
4475 pc.days_between(
4476 pc.floor_temporal(self._pa_array, unit="day"),
4477 pc.ceil_temporal(self._pa_array, unit="month"),
4478 ),
4479 1,
4480 )
4481 return self._from_pyarrow_array(result)
4483 @property
4484 def _dt_is_year_start(self) -> Self:
4485 result = pc.and_(
4486 pc.equal(pc.month(self._pa_array), 1),
4487 pc.equal(pc.day(self._pa_array), 1),
4488 )
4489 return self._from_pyarrow_array(result)
4491 @property
4492 def _dt_is_year_end(self) -> Self:
4493 result = pc.and_(
4494 pc.equal(pc.month(self._pa_array), 12),
4495 pc.equal(pc.day(self._pa_array), 31),
4496 )
4497 return self._from_pyarrow_array(result)
4499 @property
4500 def _dt_is_quarter_start(self) -> Self:
4501 result = pc.equal(
4502 pc.floor_temporal(self._pa_array, unit="quarter"),
4503 pc.floor_temporal(self._pa_array, unit="day"),
4504 )
4505 return self._from_pyarrow_array(result)
4507 @property
4508 def _dt_is_quarter_end(self) -> Self:
4509 result = pc.equal(
4510 pc.days_between(
4511 pc.floor_temporal(self._pa_array, unit="day"),
4512 pc.ceil_temporal(self._pa_array, unit="quarter"),
4513 ),
4514 1,
4515 )
4516 return self._from_pyarrow_array(result)
4518 @property
4519 def _dt_days_in_month(self) -> Self:
4520 result = pc.days_between(
4521 pc.floor_temporal(self._pa_array, unit="month"),
4522 pc.ceil_temporal(self._pa_array, unit="month"),
4523 )
4524 return self._from_pyarrow_array(result)
4526 @property
4527 def _dt_daysinmonth(self) -> Self:
4528 warnings.warn(
4529 "ArrowDtype daysinmonth accessor is deprecated and will be removed "
4530 "in a future version. Use days_in_month instead.",
4531 Pandas4Warning,
4532 stacklevel=find_stack_level(),
4533 )
4534 return self._dt_days_in_month
4536 @property
4537 def _dt_microsecond(self) -> Self:
4538 # GH 59154
4539 us = pc.microsecond(self._pa_array)
4540 ms_to_us = pc.multiply(pc.millisecond(self._pa_array), 1000)
4541 result = pc.add(us, ms_to_us)
4542 return self._from_pyarrow_array(result)
4544 @property
4545 def _dt_minute(self) -> Self:
4546 result = pc.minute(self._pa_array)
4547 return self._from_pyarrow_array(result)
4549 @property
4550 def _dt_month(self) -> Self:
4551 result = pc.month(self._pa_array)
4552 return self._from_pyarrow_array(result)
4554 @property
4555 def _dt_nanosecond(self) -> Self:
4556 result = pc.nanosecond(self._pa_array)
4557 return self._from_pyarrow_array(result)
4559 @property
4560 def _dt_quarter(self) -> Self:
4561 result = pc.quarter(self._pa_array)
4562 return self._from_pyarrow_array(result)
4564 @property
4565 def _dt_second(self) -> Self:
4566 result = pc.second(self._pa_array)
4567 return self._from_pyarrow_array(result)
4569 @property
4570 def _dt_date(self) -> Self:
4571 result = self._pa_array.cast(pa.date32())
4572 return self._from_pyarrow_array(result)
4574 @property
4575 def _dt_time(self) -> Self:
4576 unit = (
4577 self.dtype.pyarrow_dtype.unit
4578 if self.dtype.pyarrow_dtype.unit in {"us", "ns"}
4579 else "ns"
4580 )
4581 result = self._pa_array.cast(pa.time64(unit))
4582 return self._from_pyarrow_array(result)
4584 @property
4585 def _dt_tz(self):
4586 return timezones.maybe_get_tz(self.dtype.pyarrow_dtype.tz)
4588 @property
4589 def _dt_unit(self):
4590 return self.dtype.pyarrow_dtype.unit
4592 def _dt_normalize(self) -> Self:
4593 result = pc.floor_temporal(self._pa_array, 1, "day")
4594 return self._from_pyarrow_array(result)
4596 def _dt_strftime(self, format: str) -> Self:
4597 result = pc.strftime(self._pa_array, format=format)
4598 return self._from_pyarrow_array(result)
4600 def _round_temporally(
4601 self,
4602 method: Literal["ceil", "floor", "round"],
4603 freq,
4604 ambiguous: TimeAmbiguous = "raise",
4605 nonexistent: TimeNonexistent = "raise",
4606 ) -> Self:
4607 if ambiguous != "raise":
4608 raise NotImplementedError("ambiguous is not supported.")
4609 if nonexistent != "raise":
4610 raise NotImplementedError("nonexistent is not supported.")
4611 offset = to_offset(freq)
4612 if offset is None:
4613 raise ValueError(f"Must specify a valid frequency: {freq}")
4614 pa_supported_unit = {
4615 "Y": "year",
4616 "YS": "year",
4617 "Q": "quarter",
4618 "QS": "quarter",
4619 "M": "month",
4620 "MS": "month",
4621 "W": "week",
4622 "D": "day",
4623 "h": "hour",
4624 "min": "minute",
4625 "s": "second",
4626 "ms": "millisecond",
4627 "us": "microsecond",
4628 "ns": "nanosecond",
4629 }
4630 unit = pa_supported_unit.get(offset._prefix, None)
4631 if unit is None:
4632 raise ValueError(f"{freq=} is not supported")
4633 multiple = offset.n
4634 rounding_method = getattr(pc, f"{method}_temporal")
4635 result = rounding_method(self._pa_array, multiple=multiple, unit=unit)
4636 return self._from_pyarrow_array(result)
4638 def _dt_ceil(
4639 self,
4640 freq,
4641 ambiguous: TimeAmbiguous = "raise",
4642 nonexistent: TimeNonexistent = "raise",
4643 ) -> Self:
4644 return self._round_temporally("ceil", freq, ambiguous, nonexistent)
4646 def _dt_floor(
4647 self,
4648 freq,
4649 ambiguous: TimeAmbiguous = "raise",
4650 nonexistent: TimeNonexistent = "raise",
4651 ) -> Self:
4652 return self._round_temporally("floor", freq, ambiguous, nonexistent)
4654 def _dt_round(
4655 self,
4656 freq,
4657 ambiguous: TimeAmbiguous = "raise",
4658 nonexistent: TimeNonexistent = "raise",
4659 ) -> Self:
4660 return self._round_temporally("round", freq, ambiguous, nonexistent)
4662 def _dt_day_name(self, locale: str | None = None) -> Self:
4663 if locale is None:
4664 locale = "C"
4665 result = pc.strftime(self._pa_array, format="%A", locale=locale)
4666 return self._from_pyarrow_array(result)
4668 def _dt_month_name(self, locale: str | None = None) -> Self:
4669 if locale is None:
4670 locale = "C"
4671 result = pc.strftime(self._pa_array, format="%B", locale=locale)
4672 return self._from_pyarrow_array(result)
4674 def _dt_to_pydatetime(self) -> Series:
4675 from pandas import Series
4677 if pa.types.is_date(self.dtype.pyarrow_dtype):
4678 raise ValueError(
4679 f"to_pydatetime cannot be called with {self.dtype.pyarrow_dtype} type. "
4680 "Convert to pyarrow timestamp type."
4681 )
4682 data = self._pa_array.to_pylist()
4683 if self._dtype.pyarrow_dtype.unit == "ns":
4684 data = [None if ts is None else ts.to_pydatetime(warn=False) for ts in data]
4685 return Series(data, dtype=object)
4687 def _dt_tz_localize(
4688 self,
4689 tz,
4690 ambiguous: TimeAmbiguous = "raise",
4691 nonexistent: TimeNonexistent = "raise",
4692 ) -> Self:
4693 if ambiguous != "raise":
4694 raise NotImplementedError(f"{ambiguous=} is not supported")
4695 nonexistent_pa = {
4696 "raise": "raise",
4697 "shift_backward": "earliest",
4698 "shift_forward": "latest",
4699 }.get(
4700 nonexistent, # type: ignore[arg-type]
4701 None,
4702 )
4703 if nonexistent_pa is None:
4704 raise NotImplementedError(f"{nonexistent=} is not supported")
4705 if tz is None:
4706 result = pc.local_timestamp(self._pa_array)
4707 else:
4708 result = pc.assume_timezone(
4709 self._pa_array, str(tz), ambiguous=ambiguous, nonexistent=nonexistent_pa
4710 )
4711 return self._from_pyarrow_array(result)
4713 def _dt_tz_convert(self, tz) -> Self:
4714 if self.dtype.pyarrow_dtype.tz is None:
4715 raise TypeError(
4716 "Cannot convert tz-naive timestamps, use tz_localize to localize"
4717 )
4718 current_unit = self.dtype.pyarrow_dtype.unit
4719 result = self._pa_array.cast(pa.timestamp(current_unit, tz))
4720 return self._from_pyarrow_array(result)
4723def transpose_homogeneous_pyarrow(
4724 arrays: Sequence[ArrowExtensionArray],
4725) -> list[ArrowExtensionArray]:
4726 """Transpose arrow extension arrays in a list, but faster.
4728 Input should be a list of arrays of equal length and all have the same
4729 dtype. The caller is responsible for ensuring validity of input data.
4730 """
4731 arrays = list(arrays)
4732 nrows, ncols = len(arrays[0]), len(arrays)
4733 indices = np.arange(nrows * ncols).reshape(ncols, nrows).T.reshape(-1)
4734 arr = pa.chunked_array([chunk for arr in arrays for chunk in arr._pa_array.chunks])
4735 arr = arr.take(indices)
4736 return [ArrowExtensionArray(arr.slice(i * ncols, ncols)) for i in range(nrows)]