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
3import operator
4import re
5from typing import (
6 TYPE_CHECKING,
7 Any,
8 Self,
9)
11import numpy as np
13from pandas._libs import (
14 lib,
15 missing as libmissing,
16)
17from pandas.compat import (
18 HAS_PYARROW,
19 PYARROW_MIN_VERSION,
20)
21from pandas.compat.numpy import function as nv
22from pandas.util._decorators import set_module
23from pandas.util._validators import validate_na_arg
25from pandas.core.dtypes.cast import construct_1d_object_array_from_listlike
26from pandas.core.dtypes.common import (
27 is_array_like_deprecate_non_pandas,
28 is_scalar,
29 pandas_dtype,
30)
31from pandas.core.dtypes.missing import isna
33from pandas.core.arrays._arrow_string_mixins import ArrowStringArrayMixin
34from pandas.core.arrays.arrow import ArrowExtensionArray
35from pandas.core.arrays.boolean import BooleanDtype
36from pandas.core.arrays.floating import Float64Dtype
37from pandas.core.arrays.integer import Int64Dtype
38from pandas.core.arrays.numeric import NumericDtype
39from pandas.core.arrays.string_ import (
40 BaseStringArray,
41 StringDtype,
42)
43from pandas.core.strings.object_array import ObjectStringArrayMixin
45if HAS_PYARROW:
46 import pyarrow as pa
47 import pyarrow.compute as pc
50if TYPE_CHECKING:
51 from collections.abc import (
52 Callable,
53 Sequence,
54 )
56 from pandas._typing import (
57 ArrayLike,
58 AxisInt,
59 Dtype,
60 NpDtype,
61 Scalar,
62 npt,
63 )
65 from pandas.core.dtypes.dtypes import ExtensionDtype
67 from pandas import Series
70def _check_pyarrow_available() -> None:
71 if not HAS_PYARROW:
72 msg = (
73 f"pyarrow>={PYARROW_MIN_VERSION} is required for PyArrow "
74 "backed ArrowExtensionArray."
75 )
76 raise ImportError(msg)
79# Matches a `\Z` that is an end-of-string assertion rather than an escaped
80# backslash followed by a literal "Z"; the captured group keeps any preceding
81# pairs of escaped backslashes intact.
82_unescaped_end_anchor = re.compile(r"(?<!\\)((?:\\\\)*)\\Z")
85# TODO: Inherit directly from BaseStringArrayMethods. Currently we inherit from
86# ObjectStringArrayMixin because we want to have the object-dtype based methods as
87# fallback for the ones that pyarrow doesn't yet support
90@set_module("pandas.arrays")
91class ArrowStringArray(ObjectStringArrayMixin, ArrowExtensionArray, BaseStringArray):
92 """
93 Extension array for string data in a ``pyarrow.ChunkedArray``.
95 .. warning::
97 ArrowStringArray is considered experimental. The implementation and
98 parts of the API may change without warning.
100 Parameters
101 ----------
102 values : pyarrow.Array or pyarrow.ChunkedArray
103 The array of data.
104 dtype : StringDtype
105 The dtype for the array.
107 Attributes
108 ----------
109 None
111 Methods
112 -------
113 None
115 See Also
116 --------
117 :func:`array`
118 The recommended function for creating an ArrowStringArray.
119 Series.str
120 The string methods are available on Series backed by
121 an ArrowStringArray.
123 Notes
124 -----
125 ArrowStringArray returns a BooleanArray for comparison methods.
127 Examples
128 --------
129 >>> pd.array(["This is", "some text", None, "data."], dtype="string[pyarrow]")
130 <ArrowStringArray>
131 ['This is', 'some text', <NA>, 'data.']
132 Length: 4, dtype: string
133 """
135 # error: Incompatible types in assignment (expression has type "StringDtype",
136 # base class "ArrowExtensionArray" defined the type as "ArrowDtype")
137 _dtype: StringDtype # type: ignore[assignment]
139 def __init__(self, values, *, dtype: StringDtype | None = None) -> None:
140 _check_pyarrow_available()
141 if isinstance(values, (pa.Array, pa.ChunkedArray)) and (
142 pa.types.is_string(values.type)
143 or pa.types.is_string_view(values.type)
144 or (
145 pa.types.is_dictionary(values.type)
146 and (
147 pa.types.is_string(values.type.value_type)
148 or pa.types.is_large_string(values.type.value_type)
149 or pa.types.is_string_view(values.type.value_type)
150 )
151 )
152 ):
153 values = pc.cast(values, pa.large_string())
155 super().__init__(values)
157 if dtype is None:
158 dtype = StringDtype(storage="pyarrow", na_value=libmissing.NA)
159 self._dtype = dtype
161 if not pa.types.is_large_string(self._pa_array.type):
162 raise ValueError(
163 "ArrowStringArray requires a PyArrow (chunked) array of "
164 "large_string type"
165 )
167 def _from_pyarrow_array(self, pa_array):
168 """
169 Construct from a pyarrow Array/ChunkedArray result of an operation.
171 Avoids full __init__ overhead (type checking, pc.cast, ArrowDtype
172 construction, etc.).
173 """
174 assert isinstance(pa_array, (pa.Array, pa.ChunkedArray))
175 if not pa.types.is_large_string(pa_array.type):
176 pa_array = pa_array.cast(pa.large_string())
177 obj = type(self).__new__(type(self))
178 if isinstance(pa_array, pa.Array):
179 pa_array = pa.chunked_array([pa_array])
180 obj._pa_array = pa_array
181 obj._dtype = self._dtype
182 obj._cache = {}
183 return obj
185 def _cast_pointwise_result(self, values) -> ArrayLike:
186 if len(values) == 0:
187 return self[:0].copy()
189 try:
190 arr = pa.array(values, from_pandas=True)
191 except (ValueError, TypeError, OverflowError):
192 # OverflowError for e.g. 2**70; safe=True so
193 # [nan, 2**64 - 1] stays object instead of rounding, GH#70233
194 values = construct_1d_object_array_from_listlike(values)
195 return lib.maybe_convert_objects(
196 values, convert_non_numeric=True, safe=True
197 )
198 if pa.types.is_string(arr.type) or pa.types.is_large_string(arr.type):
199 return self._from_pyarrow_array(arr)
200 if self.dtype.na_value is np.nan:
201 # ArrowEA has different semantics, so we return numpy-based
202 # result instead
203 values = construct_1d_object_array_from_listlike(values)
204 return lib.maybe_convert_objects(values, convert_non_numeric=True)
205 return ArrowExtensionArray(arr)
207 @classmethod
208 def _box_pa_scalar(cls, value, pa_type: pa.DataType | None = None) -> pa.Scalar:
209 pa_scalar = super()._box_pa_scalar(value, pa_type)
210 if pa.types.is_string(pa_scalar.type) and pa_type is None:
211 pa_scalar = pc.cast(pa_scalar, pa.large_string())
212 return pa_scalar
214 @classmethod
215 def _box_pa_array(
216 cls, value, pa_type: pa.DataType | None = None, copy: bool = False
217 ) -> pa.Array | pa.ChunkedArray:
218 pa_array = super()._box_pa_array(value, pa_type)
219 if pa.types.is_string(pa_array.type) and pa_type is None:
220 pa_array = pc.cast(pa_array, pa.large_string())
221 return pa_array
223 def __len__(self) -> int:
224 """
225 Length of this array.
227 Returns
228 -------
229 length : int
230 """
231 return len(self._pa_array)
233 @classmethod
234 def _from_sequence(
235 cls, scalars, *, dtype: Dtype | None = None, copy: bool = False
236 ) -> Self:
237 from pandas.core.arrays.masked import BaseMaskedArray
239 _check_pyarrow_available()
241 if dtype and not (isinstance(dtype, str) and dtype == "string"):
242 dtype = pandas_dtype(dtype)
243 assert isinstance(dtype, StringDtype) and dtype.storage == "pyarrow"
245 if isinstance(scalars, BaseMaskedArray):
246 # avoid costly conversion to object dtype in ensure_string_array and
247 # numerical issues with Float32Dtype
248 na_values = scalars._mask
249 result = scalars._data
250 if result.dtype.kind in "iub":
251 # GH#56505 Use PyArrow's native cast for integer and boolean
252 # types, avoiding the element-wise ensure_string_array loop.
253 pa_arr = pa.array(result, mask=na_values)
254 pa_arr = pc.cast(pa_arr, pa.large_string())
255 if result.dtype.kind == "b":
256 pa_arr = pc.utf8_capitalize(pa_arr)
257 else:
258 result = lib.ensure_string_array(
259 result, copy=copy, convert_na_value=False, skipna=False
260 )
261 pa_arr = pa.array(result, mask=na_values, type=pa.large_string())
262 elif isinstance(scalars, ArrowExtensionArray):
263 pa_type = scalars._pa_array.type
264 # Use PyArrow's native cast for integer, string, and boolean types.
265 # Float has different representation in PyArrow: 1.0 -> "1" instead
266 # of "1.0", and uses different scientific notation (1e+10 vs 1e10).
267 # Boolean needs capitalize (true -> True, false -> False).
268 if (
269 pa.types.is_integer(pa_type)
270 or pa.types.is_large_string(pa_type)
271 or pa.types.is_string(pa_type)
272 or pa.types.is_boolean(pa_type)
273 ):
274 pa_arr = pc.cast(scalars._pa_array, pa.large_string())
275 if pa.types.is_boolean(pa_type):
276 pa_arr = pc.utf8_capitalize(pa_arr)
277 else:
278 # Fall back for types where PyArrow's string representation
279 # differs from Python's str()
280 result = lib.ensure_string_array(scalars, copy=copy)
281 pa_arr = pa.array(result, type=pa.large_string(), from_pandas=True)
282 elif isinstance(scalars, (pa.Array, pa.ChunkedArray)):
283 pa_arr = pc.cast(scalars, pa.large_string())
284 elif isinstance(scalars, np.ndarray) and scalars.dtype.kind in "iub":
285 # GH#56505 Use PyArrow's native cast for numpy integer and boolean
286 # arrays, avoiding the element-wise ensure_string_array loop.
287 pa_arr = pa.array(scalars, from_pandas=True)
288 pa_arr = pc.cast(pa_arr, pa.large_string())
289 if scalars.dtype.kind == "b":
290 pa_arr = pc.utf8_capitalize(pa_arr)
291 else:
292 # convert non-na-likes to str
293 result = lib.ensure_string_array(scalars, copy=copy)
294 pa_arr = pa.array(result, type=pa.large_string(), from_pandas=True)
295 # error: Argument "dtype" to "ArrowStringArray" has incompatible type
296 return cls(pa_arr, dtype=dtype) # type: ignore[arg-type]
298 @classmethod
299 def _from_sequence_of_strings(
300 cls, strings, *, dtype: ExtensionDtype, copy: bool = False
301 ) -> Self:
302 return cls._from_sequence(strings, dtype=dtype, copy=copy)
304 @property
305 def dtype(self) -> StringDtype: # type: ignore[override]
306 """
307 An instance of 'string[pyarrow]'.
308 """
309 return self._dtype
311 def insert(self, loc: int, item) -> ArrowStringArray:
312 if self.dtype.na_value is np.nan and item is np.nan:
313 item = libmissing.NA
314 if not isinstance(item, str) and item is not libmissing.NA:
315 raise TypeError(
316 f"Invalid value '{item}' for dtype 'str'. Value should be a "
317 f"string or missing value, got '{type(item).__name__}' instead."
318 )
319 return super().insert(loc, item)
321 def _convert_bool_result(self, values, na=lib.no_default, method_name=None):
322 validate_na_arg(na, name="na")
323 if self.dtype.na_value is np.nan:
324 if na is lib.no_default or isna(na):
325 # NaN propagates as False
326 values = values.fill_null(False)
327 else:
328 values = values.fill_null(na)
329 return values.to_numpy()
330 elif na is not lib.no_default and not isna(na): # pyright: ignore [reportGeneralTypeIssues]
331 values = values.fill_null(na)
332 return BooleanDtype().__from_arrow__(values)
334 def _validate_setitem_value(self, value):
335 """Maybe convert value to be pyarrow compatible."""
336 if is_scalar(value):
337 if isna(value):
338 value = None
339 elif not isinstance(value, str):
340 raise TypeError(
341 f"Invalid value '{value}' for dtype 'str'. Value should be a "
342 f"string or missing value, got '{type(value).__name__}' instead."
343 )
344 elif isinstance(value, type(self)):
345 pass
346 else:
347 if not is_array_like_deprecate_non_pandas(value):
348 value = np.asarray(value, dtype=object)
349 else:
350 value = np.asarray(value)
351 if len(value) and not (
352 value.ndim == 1 and lib.is_string_array(value, skipna=True)
353 ):
354 raise TypeError(
355 "Invalid value for dtype 'str'. Value should be a "
356 "string or missing value (or array of those)."
357 )
358 return super()._validate_setitem_value(value)
360 def isin(self, values: ArrayLike) -> npt.NDArray[np.bool_]:
361 value_set = [
362 pa_scalar.as_py()
363 for pa_scalar in [pa.scalar(value, from_pandas=True) for value in values]
364 if pa_scalar.type in (pa.string(), pa.null(), pa.large_string())
365 ]
367 # short-circuit to return all False array.
368 if not value_set:
369 return np.zeros(len(self), dtype=bool)
371 result = pc.is_in(
372 self._pa_array, value_set=pa.array(value_set, type=self._pa_array.type)
373 )
374 # pyarrow 2.0.0 returned nulls, so we explicitly specify dtype to convert nulls
375 # to False
376 return np.array(result, dtype=np.bool_)
378 def astype(self, dtype, copy: bool = True):
379 dtype = pandas_dtype(dtype)
381 if dtype == self.dtype:
382 if copy:
383 return self.copy()
384 return self
385 elif isinstance(dtype, NumericDtype):
386 data = self._pa_array.cast(pa.from_numpy_dtype(dtype.numpy_dtype))
387 return dtype.__from_arrow__(data)
388 elif isinstance(dtype, np.dtype) and np.issubdtype(dtype, np.floating):
389 return self.to_numpy(dtype=dtype, na_value=np.nan)
391 return super().astype(dtype, copy=copy)
393 # ------------------------------------------------------------------------
394 # String methods interface
396 _str_isalnum = ArrowStringArrayMixin._str_isalnum
397 _str_isalpha = ArrowStringArrayMixin._str_isalpha
398 _str_isascii = ArrowStringArrayMixin._str_isascii
399 _str_isdecimal = ArrowStringArrayMixin._str_isdecimal
400 _str_isdigit = ArrowStringArrayMixin._str_isdigit
401 _str_islower = ArrowStringArrayMixin._str_islower
402 _str_isnumeric = ArrowStringArrayMixin._str_isnumeric
403 _str_isspace = ArrowStringArrayMixin._str_isspace
404 _str_istitle = ArrowStringArrayMixin._str_istitle
405 _str_isupper = ArrowStringArrayMixin._str_isupper
407 _str_map = BaseStringArray._str_map
408 _str_startswith = ArrowStringArrayMixin._str_startswith
409 _str_endswith = ArrowStringArrayMixin._str_endswith
410 _str_pad = ArrowStringArrayMixin._str_pad
411 _str_lower = ArrowStringArrayMixin._str_lower
412 _str_upper = ArrowStringArrayMixin._str_upper
413 _str_strip = ArrowStringArrayMixin._str_strip
414 _str_lstrip = ArrowStringArrayMixin._str_lstrip
415 _str_rstrip = ArrowStringArrayMixin._str_rstrip
416 _str_removesuffix = ArrowStringArrayMixin._str_removesuffix
417 _str_removeprefix = ArrowStringArrayMixin._str_removeprefix
418 _str_find = ArrowStringArrayMixin._str_find
419 _str_get = ArrowStringArrayMixin._str_get
420 _str_getitem = ArrowStringArrayMixin._str_getitem
421 _str_capitalize = ArrowStringArrayMixin._str_capitalize
422 _str_title = ArrowStringArrayMixin._str_title
423 _str_swapcase = ArrowStringArrayMixin._str_swapcase
424 _str_slice_replace = ArrowStringArrayMixin._str_slice_replace
425 _str_len = ArrowStringArrayMixin._str_len
426 _str_slice = ArrowStringArrayMixin._str_slice
427 _str_zfill = ArrowStringArrayMixin._str_zfill
428 _str_normalize = ArrowStringArrayMixin._str_normalize
430 @classmethod
431 def _preprocess_re_pattern(
432 cls, pat: str | re.Pattern[str], case: bool, flags: int
433 ) -> tuple[str, bool, int]:
434 pattern, case, flags = cls._unwrap_re_pattern(pat, case, flags)
436 # RE2 spells the end-of-string anchor `\z`; Python's `\Z` is the same
437 # assertion and may appear anywhere in the pattern, not just at the end.
438 pattern = _unescaped_end_anchor.sub(r"\1\\z", pattern)
440 return pattern, case, flags
442 def _str_contains(
443 self,
444 pat,
445 case: bool = True,
446 flags: int = 0,
447 na=lib.no_default,
448 regex: bool = True,
449 ):
450 if (
451 flags
452 or self._is_re_pattern_with_flags(pat)
453 or (regex and self._needs_python_regex(pat))
454 # a compiled pattern is not a valid literal; defer to the object path
455 or (not regex and isinstance(pat, re.Pattern))
456 ):
457 return super()._str_contains(pat, case, flags, na, regex)
459 # with regex=False `pat` is matched literally, so must not be rewritten
460 if regex:
461 pat, case, flags = self._preprocess_re_pattern(pat, case, flags)
462 return ArrowStringArrayMixin._str_contains(self, pat, case, flags, na, regex)
464 def _str_match(
465 self,
466 pat: str | re.Pattern[str],
467 case: bool = True,
468 flags: int = 0,
469 na: Scalar | lib.NoDefault = lib.no_default,
470 ):
471 if (
472 flags
473 or self._is_re_pattern_with_flags(pat)
474 or self._needs_python_regex(pat)
475 ):
476 return super()._str_match(pat, case, flags, na)
478 pat, case, flags = self._preprocess_re_pattern(pat, case, flags)
479 return ArrowStringArrayMixin._str_match(self, pat, case, flags, na)
481 def _str_fullmatch(
482 self,
483 pat: str | re.Pattern[str],
484 case: bool = True,
485 flags: int = 0,
486 na: Scalar | lib.NoDefault = lib.no_default,
487 ):
488 if (
489 flags
490 or self._is_re_pattern_with_flags(pat)
491 or self._needs_python_regex(pat)
492 ):
493 return super()._str_fullmatch(pat, case, flags, na)
495 pat, case, flags = self._preprocess_re_pattern(pat, case, flags)
496 return ArrowStringArrayMixin._str_fullmatch(self, pat, case, flags, na)
498 def _str_replace(
499 self,
500 pat: str | re.Pattern[str],
501 repl: str | Callable[..., Any],
502 n: int = -1,
503 case: bool = True,
504 flags: int = 0,
505 regex: bool = True,
506 ):
507 if (
508 isinstance(pat, re.Pattern)
509 or callable(repl)
510 or not case
511 or flags
512 or ( # substitution contains a named group pattern
513 # https://docs.python.org/3/library/re.html
514 isinstance(repl, str) and r"\g<" in repl
515 )
516 or (regex and self._needs_python_regex(pat))
517 ):
518 return super()._str_replace(pat, repl, n, case, flags, regex)
520 if regex:
521 pat, case, flags = self._preprocess_re_pattern(pat, case, flags)
523 return ArrowStringArrayMixin._str_replace(
524 self, pat, repl, n, case, flags, regex
525 )
527 def _str_repeat(self, repeats: int | Sequence[int]):
528 if not isinstance(repeats, int):
529 return super()._str_repeat(repeats)
530 else:
531 return ArrowExtensionArray._str_repeat(self, repeats=repeats)
533 def _str_count(self, pat: str | re.Pattern[str], flags: int = 0):
534 if (
535 flags
536 or self._is_re_pattern_with_flags(pat)
537 or self._needs_python_regex(pat)
538 ):
539 return super()._str_count(pat, flags)
541 pat, case, _ = self._preprocess_re_pattern(pat, True, 0)
542 result = pc.count_substring_regex(self._pa_array, pat, ignore_case=not case)
543 return self._convert_int_result(result)
545 def _str_partition(self, sep: str, expand: bool):
546 if expand:
547 # rows of three strings, so a list array rather than Self
548 return ArrowExtensionArray(
549 ArrowStringArrayMixin._str_partition_expand(self, sep)
550 )
551 # expand=False wants tuples, which the object path produces directly
552 return ObjectStringArrayMixin._str_partition(self, sep, expand)
554 def _str_get_dummies(self, sep: str = "|", dtype: NpDtype | None = None):
555 if dtype is None:
556 dtype = np.int64
557 dummies_pa, labels = ArrowExtensionArray(self._pa_array)._str_get_dummies(
558 sep, dtype
559 )
560 if len(labels) == 0:
561 return np.empty(shape=(0, 0), dtype=dtype), labels
562 dummies = np.vstack(dummies_pa.to_numpy())
563 _dtype = pandas_dtype(dtype)
564 dummies_dtype: NpDtype
565 if isinstance(_dtype, np.dtype):
566 dummies_dtype = _dtype
567 else:
568 dummies_dtype = np.bool_
569 return dummies.astype(dummies_dtype, copy=False), labels
571 def _convert_int_result(self, result):
572 if self.dtype.na_value is np.nan:
573 result = result.cast(pa.int64())
574 if isinstance(result, pa.Array):
575 result = result.to_numpy(zero_copy_only=False)
576 else:
577 result = result.to_numpy()
578 return result
580 return Int64Dtype().__from_arrow__(result)
582 def _convert_rank_result(self, result):
583 if self.dtype.na_value is np.nan:
584 if isinstance(result, pa.Array):
585 result = result.to_numpy(zero_copy_only=False)
586 else:
587 result = result.to_numpy()
588 return result.astype("float64", copy=False)
590 return Float64Dtype().__from_arrow__(result)
592 def _reduce(
593 self,
594 name: str,
595 *,
596 skipna: bool = True,
597 keepdims: bool = False,
598 axis: AxisInt | None = 0,
599 **kwargs,
600 ):
601 nv.validate_minmax_axis(axis, self.ndim)
602 if self.dtype.na_value is np.nan and name in ["any", "all"]:
603 if not skipna:
604 nas = pc.is_null(self._pa_array)
605 arr = pc.or_kleene(nas, pc.not_equal(self._pa_array, ""))
606 else:
607 arr = pc.not_equal(self._pa_array, "")
608 result = ArrowExtensionArray(arr)._reduce(
609 name, skipna=skipna, keepdims=keepdims, **kwargs
610 )
611 if keepdims:
612 # ArrowExtensionArray will return a length-1 bool[pyarrow] array
613 return result.astype(np.bool_)
614 return result
616 if name in ("count", "min", "max", "sum", "argmin", "argmax"):
617 result = self._reduce_calc(name, skipna=skipna, keepdims=keepdims, **kwargs)
618 else:
619 raise TypeError(f"Cannot perform reduction '{name}' with string dtype")
621 if name in ("argmin", "argmax") and isinstance(result, pa.Array):
622 return self._convert_int_result(result)
623 elif isinstance(result, pa.Array):
624 return type(self)(result, dtype=self.dtype)
625 else:
626 return result
628 def value_counts(self, dropna: bool = True) -> Series:
629 result = super().value_counts(dropna=dropna)
630 if self.dtype.na_value is np.nan:
631 res_values = result._values.to_numpy()
632 return result._constructor(
633 res_values, index=result.index, name=result.name, copy=False
634 )
635 return result
637 def _cmp_method(self, other, op):
638 if (
639 isinstance(other, (BaseStringArray, ArrowExtensionArray))
640 and self.dtype.na_value is not libmissing.NA
641 and other.dtype.na_value is libmissing.NA
642 ):
643 # NA has priority of NaN semantics
644 return NotImplemented
646 result = super()._cmp_method(other, op)
647 if self.dtype.na_value is np.nan:
648 if op == operator.ne:
649 return result.to_numpy(np.bool_, na_value=True)
650 else:
651 return result.to_numpy(np.bool_, na_value=False)
652 return result
654 def __pos__(self) -> Self:
655 raise TypeError(f"bad operand type for unary +: '{self.dtype}'")