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 Self,
8)
10import numpy as np
12from pandas._libs import (
13 lib,
14 missing as libmissing,
15)
16from pandas.compat import (
17 HAS_PYARROW,
18 PYARROW_MIN_VERSION,
19 pa_version_under16p0,
20)
21from pandas.util._decorators import set_module
22from pandas.util._validators import validate_na_arg
24from pandas.core.dtypes.common import (
25 is_scalar,
26 pandas_dtype,
27)
28from pandas.core.dtypes.inference import is_array_like
29from pandas.core.dtypes.missing import isna
31from pandas.core.arrays._arrow_string_mixins import ArrowStringArrayMixin
32from pandas.core.arrays.arrow import ArrowExtensionArray
33from pandas.core.arrays.boolean import BooleanDtype
34from pandas.core.arrays.floating import Float64Dtype
35from pandas.core.arrays.integer import Int64Dtype
36from pandas.core.arrays.numeric import NumericDtype
37from pandas.core.arrays.string_ import (
38 BaseStringArray,
39 StringDtype,
40)
41from pandas.core.strings.object_array import ObjectStringArrayMixin
43if HAS_PYARROW:
44 import pyarrow as pa
45 import pyarrow.compute as pc
48if TYPE_CHECKING:
49 from collections.abc import (
50 Callable,
51 Sequence,
52 )
54 from pandas._typing import (
55 ArrayLike,
56 Dtype,
57 NpDtype,
58 Scalar,
59 npt,
60 )
62 from pandas.core.dtypes.dtypes import ExtensionDtype
64 from pandas import Series
67def _check_pyarrow_available() -> None:
68 if not HAS_PYARROW:
69 msg = (
70 f"pyarrow>={PYARROW_MIN_VERSION} is required for PyArrow "
71 "backed ArrowExtensionArray."
72 )
73 raise ImportError(msg)
76def _is_string_view(typ):
77 return not pa_version_under16p0 and pa.types.is_string_view(typ)
80# TODO: Inherit directly from BaseStringArrayMethods. Currently we inherit from
81# ObjectStringArrayMixin because we want to have the object-dtype based methods as
82# fallback for the ones that pyarrow doesn't yet support
85@set_module("pandas.arrays")
86class ArrowStringArray(ObjectStringArrayMixin, ArrowExtensionArray, BaseStringArray):
87 """
88 Extension array for string data in a ``pyarrow.ChunkedArray``.
90 .. warning::
92 ArrowStringArray is considered experimental. The implementation and
93 parts of the API may change without warning.
95 Parameters
96 ----------
97 values : pyarrow.Array or pyarrow.ChunkedArray
98 The array of data.
99 dtype : StringDtype
100 The dtype for the array.
102 Attributes
103 ----------
104 None
106 Methods
107 -------
108 None
110 See Also
111 --------
112 :func:`array`
113 The recommended function for creating an ArrowStringArray.
114 Series.str
115 The string methods are available on Series backed by
116 an ArrowStringArray.
118 Notes
119 -----
120 ArrowStringArray returns a BooleanArray for comparison methods.
122 Examples
123 --------
124 >>> pd.array(["This is", "some text", None, "data."], dtype="string[pyarrow]")
125 <ArrowStringArray>
126 ['This is', 'some text', <NA>, 'data.']
127 Length: 4, dtype: string
128 """
130 # error: Incompatible types in assignment (expression has type "StringDtype",
131 # base class "ArrowExtensionArray" defined the type as "ArrowDtype")
132 _dtype: StringDtype # type: ignore[assignment]
134 def __init__(self, values, *, dtype: StringDtype | None = None) -> None:
135 _check_pyarrow_available()
136 if isinstance(values, (pa.Array, pa.ChunkedArray)) and (
137 pa.types.is_string(values.type)
138 or _is_string_view(values.type)
139 or (
140 pa.types.is_dictionary(values.type)
141 and (
142 pa.types.is_string(values.type.value_type)
143 or pa.types.is_large_string(values.type.value_type)
144 or _is_string_view(values.type.value_type)
145 )
146 )
147 ):
148 values = pc.cast(values, pa.large_string())
150 super().__init__(values)
152 if dtype is None:
153 dtype = StringDtype(storage="pyarrow", na_value=libmissing.NA)
154 self._dtype = dtype
156 if not pa.types.is_large_string(self._pa_array.type):
157 raise ValueError(
158 "ArrowStringArray requires a PyArrow (chunked) array of "
159 "large_string type"
160 )
162 def _from_pyarrow_array(self, pa_array):
163 """
164 Construct from the pyarrow array result of an operation, retaining
165 self.dtype.na_value.
166 """
167 return type(self)(pa_array, dtype=self.dtype)
169 @classmethod
170 def _box_pa_scalar(cls, value, pa_type: pa.DataType | None = None) -> pa.Scalar:
171 pa_scalar = super()._box_pa_scalar(value, pa_type)
172 if pa.types.is_string(pa_scalar.type) and pa_type is None:
173 pa_scalar = pc.cast(pa_scalar, pa.large_string())
174 return pa_scalar
176 @classmethod
177 def _box_pa_array(
178 cls, value, pa_type: pa.DataType | None = None, copy: bool = False
179 ) -> pa.Array | pa.ChunkedArray:
180 pa_array = super()._box_pa_array(value, pa_type)
181 if pa.types.is_string(pa_array.type) and pa_type is None:
182 pa_array = pc.cast(pa_array, pa.large_string())
183 return pa_array
185 def __len__(self) -> int:
186 """
187 Length of this array.
189 Returns
190 -------
191 length : int
192 """
193 return len(self._pa_array)
195 @classmethod
196 def _from_sequence(
197 cls, scalars, *, dtype: Dtype | None = None, copy: bool = False
198 ) -> Self:
199 from pandas.core.arrays.masked import BaseMaskedArray
201 _check_pyarrow_available()
203 if dtype and not (isinstance(dtype, str) and dtype == "string"):
204 dtype = pandas_dtype(dtype)
205 assert isinstance(dtype, StringDtype) and dtype.storage == "pyarrow"
207 if isinstance(scalars, BaseMaskedArray):
208 # avoid costly conversion to object dtype in ensure_string_array and
209 # numerical issues with Float32Dtype
210 na_values = scalars._mask
211 result = scalars._data
212 result = lib.ensure_string_array(
213 result, copy=copy, convert_na_value=False, skipna=False
214 )
215 pa_arr = pa.array(result, mask=na_values, type=pa.large_string())
216 elif isinstance(scalars, ArrowExtensionArray):
217 pa_type = scalars._pa_array.type
218 # Use PyArrow's native cast for integer, string, and boolean types.
219 # Float has different representation in PyArrow: 1.0 -> "1" instead
220 # of "1.0", and uses different scientific notation (1e+10 vs 1e10).
221 # Boolean needs capitalize (true -> True, false -> False).
222 if (
223 pa.types.is_integer(pa_type)
224 or pa.types.is_large_string(pa_type)
225 or pa.types.is_string(pa_type)
226 or pa.types.is_boolean(pa_type)
227 ):
228 pa_arr = pc.cast(scalars._pa_array, pa.large_string())
229 if pa.types.is_boolean(pa_type):
230 pa_arr = pc.utf8_capitalize(pa_arr)
231 else:
232 # Fall back for types where PyArrow's string representation
233 # differs from Python's str()
234 result = lib.ensure_string_array(scalars, copy=copy)
235 pa_arr = pa.array(result, type=pa.large_string(), from_pandas=True)
236 elif isinstance(scalars, (pa.Array, pa.ChunkedArray)):
237 pa_arr = pc.cast(scalars, pa.large_string())
238 else:
239 # convert non-na-likes to str
240 result = lib.ensure_string_array(scalars, copy=copy)
241 pa_arr = pa.array(result, type=pa.large_string(), from_pandas=True)
242 # error: Argument "dtype" to "ArrowStringArray" has incompatible type
243 return cls(pa_arr, dtype=dtype) # type: ignore[arg-type]
245 @classmethod
246 def _from_sequence_of_strings(
247 cls, strings, *, dtype: ExtensionDtype, copy: bool = False
248 ) -> Self:
249 return cls._from_sequence(strings, dtype=dtype, copy=copy)
251 @property
252 def dtype(self) -> StringDtype: # type: ignore[override]
253 """
254 An instance of 'string[pyarrow]'.
255 """
256 return self._dtype
258 def insert(self, loc: int, item) -> ArrowStringArray:
259 if self.dtype.na_value is np.nan and item is np.nan:
260 item = libmissing.NA
261 if not isinstance(item, str) and item is not libmissing.NA:
262 raise TypeError(
263 f"Invalid value '{item}' for dtype 'str'. Value should be a "
264 f"string or missing value, got '{type(item).__name__}' instead."
265 )
266 return super().insert(loc, item)
268 def _convert_bool_result(self, values, na=lib.no_default, method_name=None):
269 validate_na_arg(na, name="na")
270 if self.dtype.na_value is np.nan:
271 if na is lib.no_default or isna(na):
272 # NaN propagates as False
273 values = values.fill_null(False)
274 else:
275 values = values.fill_null(na)
276 return values.to_numpy()
277 elif na is not lib.no_default and not isna(na): # pyright: ignore [reportGeneralTypeIssues]
278 values = values.fill_null(na)
279 return BooleanDtype().__from_arrow__(values)
281 def _maybe_convert_setitem_value(self, value):
282 """Maybe convert value to be pyarrow compatible."""
283 if is_scalar(value):
284 if isna(value):
285 value = None
286 elif not isinstance(value, str):
287 raise TypeError(
288 f"Invalid value '{value}' for dtype 'str'. Value should be a "
289 f"string or missing value, got '{type(value).__name__}' instead."
290 )
291 elif isinstance(value, type(self)):
292 pass
293 else:
294 if not is_array_like(value):
295 value = np.asarray(value, dtype=object)
296 else:
297 value = np.asarray(value)
298 if len(value) and not (
299 value.ndim == 1 and lib.is_string_array(value, skipna=True)
300 ):
301 raise TypeError(
302 "Invalid value for dtype 'str'. Value should be a "
303 "string or missing value (or array of those)."
304 )
305 return super()._maybe_convert_setitem_value(value)
307 def isin(self, values: ArrayLike) -> npt.NDArray[np.bool_]:
308 value_set = [
309 pa_scalar.as_py()
310 for pa_scalar in [pa.scalar(value, from_pandas=True) for value in values]
311 if pa_scalar.type in (pa.string(), pa.null(), pa.large_string())
312 ]
314 # short-circuit to return all False array.
315 if not value_set:
316 return np.zeros(len(self), dtype=bool)
318 result = pc.is_in(
319 self._pa_array, value_set=pa.array(value_set, type=self._pa_array.type)
320 )
321 # pyarrow 2.0.0 returned nulls, so we explicitly specify dtype to convert nulls
322 # to False
323 return np.array(result, dtype=np.bool_)
325 def astype(self, dtype, copy: bool = True):
326 dtype = pandas_dtype(dtype)
328 if dtype == self.dtype:
329 if copy:
330 return self.copy()
331 return self
332 elif isinstance(dtype, NumericDtype):
333 data = self._pa_array.cast(pa.from_numpy_dtype(dtype.numpy_dtype))
334 return dtype.__from_arrow__(data)
335 elif isinstance(dtype, np.dtype) and np.issubdtype(dtype, np.floating):
336 return self.to_numpy(dtype=dtype, na_value=np.nan)
338 return super().astype(dtype, copy=copy)
340 # ------------------------------------------------------------------------
341 # String methods interface
343 _str_isalnum = ArrowStringArrayMixin._str_isalnum
344 _str_isalpha = ArrowStringArrayMixin._str_isalpha
345 _str_isdecimal = ArrowStringArrayMixin._str_isdecimal
346 _str_isdigit = ArrowStringArrayMixin._str_isdigit
347 _str_islower = ArrowStringArrayMixin._str_islower
348 _str_isnumeric = ArrowStringArrayMixin._str_isnumeric
349 _str_isspace = ArrowStringArrayMixin._str_isspace
350 _str_istitle = ArrowStringArrayMixin._str_istitle
351 _str_isupper = ArrowStringArrayMixin._str_isupper
353 _str_map = BaseStringArray._str_map
354 _str_startswith = ArrowStringArrayMixin._str_startswith
355 _str_endswith = ArrowStringArrayMixin._str_endswith
356 _str_pad = ArrowStringArrayMixin._str_pad
357 _str_lower = ArrowStringArrayMixin._str_lower
358 _str_upper = ArrowStringArrayMixin._str_upper
359 _str_strip = ArrowStringArrayMixin._str_strip
360 _str_lstrip = ArrowStringArrayMixin._str_lstrip
361 _str_rstrip = ArrowStringArrayMixin._str_rstrip
362 _str_removesuffix = ArrowStringArrayMixin._str_removesuffix
363 _str_removeprefix = ArrowStringArrayMixin._str_removeprefix
364 _str_find = ArrowStringArrayMixin._str_find
365 _str_get = ArrowStringArrayMixin._str_get
366 _str_getitem = ArrowStringArrayMixin._str_getitem
367 _str_capitalize = ArrowStringArrayMixin._str_capitalize
368 _str_title = ArrowStringArrayMixin._str_title
369 _str_swapcase = ArrowStringArrayMixin._str_swapcase
370 _str_slice_replace = ArrowStringArrayMixin._str_slice_replace
371 _str_len = ArrowStringArrayMixin._str_len
372 _str_slice = ArrowStringArrayMixin._str_slice
374 @staticmethod
375 def _is_re_pattern_with_flags(pat: str | re.Pattern) -> bool:
376 # check if `pat` is a compiled regex pattern with flags that are not
377 # supported by pyarrow
378 return (
379 isinstance(pat, re.Pattern)
380 and (pat.flags & ~(re.IGNORECASE | re.UNICODE)) != 0
381 )
383 @staticmethod
384 def _preprocess_re_pattern(
385 pat: str | re.Pattern, case: bool, flags: int
386 ) -> tuple[str, bool, int]:
387 if isinstance(pat, re.Pattern):
388 pattern = pat.pattern
389 # TODO flags passed separately by user are ignored
390 flags = pat.flags
391 # flags is not supported by pyarrow, but `case` is -> extract and remove
392 if flags & re.IGNORECASE:
393 case = False
394 flags = flags & ~re.IGNORECASE
395 # when creating a pattern with re.compile and a string, it automatically
396 # gets a UNICODE flag, while pyarrow assumes unicode for strings anyway
397 flags = flags & ~re.UNICODE
398 else:
399 pattern = pat
401 if (
402 pattern.endswith("\\Z")
403 # Second condition counts the number of `\` that patterns ends with
404 # prior to Z -> needs to be odd to end with an unescaped \Z
405 and (len(pattern) - len(pattern[:-1].rstrip("\\")) + 1) % 2 == 1
406 ):
407 pattern = pattern[:-2] + "\\z"
409 return pattern, case, flags
411 def _str_contains(
412 self,
413 pat,
414 case: bool = True,
415 flags: int = 0,
416 na=lib.no_default,
417 regex: bool = True,
418 ):
419 if (
420 flags
421 or self._is_re_pattern_with_flags(pat)
422 or (regex and self._has_unsupported_regex(pat))
423 ):
424 return super()._str_contains(pat, case, flags, na, regex)
426 pat, case, flags = self._preprocess_re_pattern(pat, case, flags)
427 return ArrowStringArrayMixin._str_contains(self, pat, case, flags, na, regex)
429 def _str_match(
430 self,
431 pat: str | re.Pattern,
432 case: bool = True,
433 flags: int = 0,
434 na: Scalar | lib.NoDefault = lib.no_default,
435 ):
436 if (
437 flags
438 or self._is_re_pattern_with_flags(pat)
439 or self._has_unsupported_regex(pat)
440 ):
441 return super()._str_match(pat, case, flags, na)
443 pat, case, flags = self._preprocess_re_pattern(pat, case, flags)
444 return ArrowStringArrayMixin._str_match(self, pat, case, flags, na)
446 def _str_fullmatch(
447 self,
448 pat: str | re.Pattern,
449 case: bool = True,
450 flags: int = 0,
451 na: Scalar | lib.NoDefault = lib.no_default,
452 ):
453 if (
454 flags
455 or self._is_re_pattern_with_flags(pat)
456 or self._has_unsupported_regex(pat)
457 ):
458 return super()._str_fullmatch(pat, case, flags, na)
460 pat, case, flags = self._preprocess_re_pattern(pat, case, flags)
461 return ArrowStringArrayMixin._str_fullmatch(self, pat, case, flags, na)
463 def _str_replace(
464 self,
465 pat: str | re.Pattern,
466 repl: str | Callable,
467 n: int = -1,
468 case: bool = True,
469 flags: int = 0,
470 regex: bool = True,
471 ):
472 if (
473 isinstance(pat, re.Pattern)
474 or callable(repl)
475 or not case
476 or flags
477 or ( # substitution contains a named group pattern
478 # https://docs.python.org/3/library/re.html
479 isinstance(repl, str) and r"\g<" in repl
480 )
481 or (regex and self._has_unsupported_regex(pat))
482 ):
483 return super()._str_replace(pat, repl, n, case, flags, regex)
485 if regex:
486 pat, case, flags = self._preprocess_re_pattern(pat, case, flags)
488 return ArrowStringArrayMixin._str_replace(
489 self, pat, repl, n, case, flags, regex
490 )
492 def _str_repeat(self, repeats: int | Sequence[int]):
493 if not isinstance(repeats, int):
494 return super()._str_repeat(repeats)
495 else:
496 return ArrowExtensionArray._str_repeat(self, repeats=repeats)
498 def _str_count(self, pat: str, flags: int = 0):
499 if flags or self._has_unsupported_regex(pat):
500 return super()._str_count(pat, flags)
502 pat, _, _ = self._preprocess_re_pattern(pat, True, 0)
503 result = pc.count_substring_regex(self._pa_array, pat)
504 return self._convert_int_result(result)
506 def _str_get_dummies(self, sep: str = "|", dtype: NpDtype | None = None):
507 if dtype is None:
508 dtype = np.int64
509 dummies_pa, labels = ArrowExtensionArray(self._pa_array)._str_get_dummies(
510 sep, dtype
511 )
512 if len(labels) == 0:
513 return np.empty(shape=(0, 0), dtype=dtype), labels
514 dummies = np.vstack(dummies_pa.to_numpy())
515 _dtype = pandas_dtype(dtype)
516 dummies_dtype: NpDtype
517 if isinstance(_dtype, np.dtype):
518 dummies_dtype = _dtype
519 else:
520 dummies_dtype = np.bool_
521 return dummies.astype(dummies_dtype, copy=False), labels
523 def _convert_int_result(self, result):
524 if self.dtype.na_value is np.nan:
525 result = result.cast(pa.int64())
526 if isinstance(result, pa.Array):
527 result = result.to_numpy(zero_copy_only=False)
528 else:
529 result = result.to_numpy()
530 return result
532 return Int64Dtype().__from_arrow__(result)
534 def _convert_rank_result(self, result):
535 if self.dtype.na_value is np.nan:
536 if isinstance(result, pa.Array):
537 result = result.to_numpy(zero_copy_only=False)
538 else:
539 result = result.to_numpy()
540 return result.astype("float64", copy=False)
542 return Float64Dtype().__from_arrow__(result)
544 def _reduce(
545 self, name: str, *, skipna: bool = True, keepdims: bool = False, **kwargs
546 ):
547 if self.dtype.na_value is np.nan and name in ["any", "all"]:
548 if not skipna:
549 nas = pc.is_null(self._pa_array)
550 arr = pc.or_kleene(nas, pc.not_equal(self._pa_array, ""))
551 else:
552 arr = pc.not_equal(self._pa_array, "")
553 result = ArrowExtensionArray(arr)._reduce(
554 name, skipna=skipna, keepdims=keepdims, **kwargs
555 )
556 if keepdims:
557 # ArrowExtensionArray will return a length-1 bool[pyarrow] array
558 return result.astype(np.bool_)
559 return result
561 if name in ("min", "max", "sum", "argmin", "argmax"):
562 result = self._reduce_calc(name, skipna=skipna, keepdims=keepdims, **kwargs)
563 else:
564 raise TypeError(f"Cannot perform reduction '{name}' with string dtype")
566 if name in ("argmin", "argmax") and isinstance(result, pa.Array):
567 return self._convert_int_result(result)
568 elif isinstance(result, pa.Array):
569 return type(self)(result, dtype=self.dtype)
570 else:
571 return result
573 def value_counts(self, dropna: bool = True) -> Series:
574 result = super().value_counts(dropna=dropna)
575 if self.dtype.na_value is np.nan:
576 res_values = result._values.to_numpy()
577 return result._constructor(
578 res_values, index=result.index, name=result.name, copy=False
579 )
580 return result
582 def _cmp_method(self, other, op):
583 if (
584 isinstance(other, (BaseStringArray, ArrowExtensionArray))
585 and self.dtype.na_value is not libmissing.NA
586 and other.dtype.na_value is libmissing.NA
587 ):
588 # NA has priority of NaN semantics
589 return NotImplemented
591 result = super()._cmp_method(other, op)
592 if self.dtype.na_value is np.nan:
593 if op == operator.ne:
594 return result.to_numpy(np.bool_, na_value=True)
595 else:
596 return result.to_numpy(np.bool_, na_value=False)
597 return result
599 def __pos__(self) -> Self:
600 raise TypeError(f"bad operand type for unary +: '{self.dtype}'")