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 functools import partial
4import re
5from typing import (
6 TYPE_CHECKING,
7 Any,
8 Literal,
9 Self,
10)
11import unicodedata
13import numpy as np
15from pandas._libs import lib
16from pandas.compat import (
17 HAS_PYARROW,
18 pa_version_under17p0,
19 pa_version_under21p0,
20)
22if HAS_PYARROW:
23 import pyarrow as pa
24 import pyarrow.compute as pc
26if TYPE_CHECKING:
27 from collections.abc import Callable
28 from types import ModuleType
30 from pandas._typing import Scalar
33# Matches a `\uxxxx` (unicode escape sequence) with an odd number of preceding "\"
34# (otherwise the \\u is escaped itself). Also checks for `\U` and `\N{`.
35_escaped_unicode_char = re.compile(r"(?<!\\)(?:\\\\)*\\(?:[uU]|N{)")
38class ArrowStringArrayMixin:
39 _pa_array: pa.ChunkedArray
41 def __init__(self, *args, **kwargs) -> None:
42 raise NotImplementedError
44 def _from_pyarrow_array(self, pa_array: pa.Array | pa.ChunkedArray) -> Self:
45 raise NotImplementedError
47 def _convert_bool_result(self, result, na=lib.no_default, method_name=None):
48 # Convert a bool-dtype result to the appropriate result type
49 raise NotImplementedError
51 def _convert_int_result(self, result):
52 # Convert an integer-dtype result to the appropriate result type
53 raise NotImplementedError
55 def _apply_elementwise(self, func: Callable[..., Any]) -> list[list[Any]]:
56 raise NotImplementedError
58 @staticmethod
59 def _re_parser() -> ModuleType:
60 try:
61 # error: Module "re" has no attribute "_parser"
62 from re import _parser # type: ignore[attr-defined]
63 except Exception as err:
64 raise type(err)(
65 "Incompatible version for regex; you will need to upgrade pandas "
66 "or downgrade Python"
67 ) from err
68 return _parser
70 @staticmethod
71 def _regex_nodes(pat: str) -> list[tuple[Any, Any]]:
72 """
73 Parse a regex pattern into a flat list of its (op, arg) nodes.
75 Nodes nested in groups, branches, repeats and character classes are
76 included.
78 Parameters
79 ----------
80 pat: str | re.Pattern
81 Regex pattern.
83 Returns
84 -------
85 list[tuple[Any, Any]]
86 The nodes of `pat`, or an empty list if Python's ``re`` cannot
87 parse it.
88 """
89 _parser = ArrowStringArrayMixin._re_parser()
90 repeats = {_parser.MAX_REPEAT, _parser.MIN_REPEAT, _parser.POSSESSIVE_REPEAT}
92 def walk(tokens):
93 for op_code, argument in tokens:
94 yield op_code, argument
95 if op_code == _parser.SUBPATTERN:
96 yield from walk(argument[3])
97 elif op_code == _parser.BRANCH:
98 for branch in argument[1]:
99 yield from walk(branch)
100 elif op_code in repeats:
101 yield from walk(argument[2])
102 elif op_code in (_parser.ATOMIC_GROUP, _parser.IN):
103 yield from walk(argument)
105 try:
106 tokens = _parser.parse(pat)
107 except re.error:
108 # Pattern not valid for Python's re (e.g. RE2 syntax like \x{...} or \p)
109 # Let the pyarrow backend handle it.
110 return []
111 return list(walk(tokens))
113 def _needs_python_regex(self, pat: str | re.Pattern[str]) -> bool:
114 """
115 Whether a regex operation must use Python's ``re`` instead of pyarrow.
117 This is the case for features not supported by RE2 / pyarrow or with
118 different behaviour:
119 - Not supported: lookaround assertions and backreferences, unicode escape
120 sequences like ``\\u``, ``\\U`` and ``\\N{}``
121 - ``\\s`` and ``\\S``, since RE2's ``\\s`` misses whitespace such as
122 ``"\\x0b"`` even in ASCII data
123 - ``\\w``, ``\\d`` and ``\\b`` (or their negations) on non-ASCII data, since
124 RE2 only matches them against ASCII characters while Python's ``re`` is
125 Unicode-aware
126 - At-end-of-string ``$`` assertion, which Python matches at the end or just
127 before a trailing newline, while pyarrow only matches at the end
128 (only falling back if newlines are present)
129 """
130 str_pat = pat.pattern if isinstance(pat, re.Pattern) else pat
131 if re.search(_escaped_unicode_char, str_pat) is not None:
132 return True
134 _parser = self._re_parser()
135 unsupported_ops = {_parser.ASSERT, _parser.ASSERT_NOT, _parser.GROUPREF}
136 whitespace = {
137 # RE2's \s misses whitespace such as "\x0b", even in ASCII data
138 (_parser.CATEGORY, _parser.CATEGORY_SPACE),
139 (_parser.CATEGORY, _parser.CATEGORY_NOT_SPACE),
140 }
141 unicode_sensitive = {
142 # unicode-sensitive categories and assertions
143 (_parser.CATEGORY, _parser.CATEGORY_WORD),
144 (_parser.CATEGORY, _parser.CATEGORY_NOT_WORD),
145 (_parser.CATEGORY, _parser.CATEGORY_DIGIT),
146 (_parser.CATEGORY, _parser.CATEGORY_NOT_DIGIT),
147 (_parser.AT, _parser.AT_BOUNDARY),
148 (_parser.AT, _parser.AT_NON_BOUNDARY),
149 }
150 newline_sensitive = {
151 # newline-sensitive behaviour of $
152 (_parser.AT, _parser.AT_END),
153 }
155 nodes = self._regex_nodes(str_pat)
156 if any(op_code in unsupported_ops for op_code, _ in nodes):
157 return True
158 if any(
159 (op_code, argument) in whitespace
160 for op_code, argument in nodes
161 if op_code == _parser.CATEGORY
162 ):
163 return True
164 if any(
165 (op_code, argument) in unicode_sensitive
166 for op_code, argument in nodes
167 if op_code in (_parser.CATEGORY, _parser.AT)
168 ):
169 # all-ASCII data keeps the faster pyarrow path
170 if self._has_non_ascii():
171 return True
172 if any(
173 (op_code, argument) in newline_sensitive
174 for op_code, argument in nodes
175 if op_code == _parser.AT
176 ):
177 if pc.any(pc.match_substring(self._pa_array, "\n"), min_count=0).as_py():
178 return True
180 return False
182 def _has_non_ascii(self) -> bool:
183 # min_count=0 so that an empty or all-null array reports False
184 return not pc.all(pc.string_is_ascii(self._pa_array), min_count=0).as_py()
186 @staticmethod
187 def _is_re_pattern_with_flags(pat: str | re.Pattern[str]) -> bool:
188 # check if `pat` is a compiled regex pattern with flags that are not
189 # supported by pyarrow
190 return (
191 isinstance(pat, re.Pattern)
192 and (pat.flags & ~(re.IGNORECASE | re.UNICODE)) != 0
193 )
195 @staticmethod
196 def _unwrap_re_pattern(
197 pat: str | re.Pattern[str], case: bool, flags: int
198 ) -> tuple[str, bool, int]:
199 """
200 Reduce `pat` to the (pattern, case, flags) triple the pyarrow kernels take.
202 IGNORECASE is the only flag they support and they spell it as `case`, so
203 it is folded into `case`; whatever `flags` remain afterwards cannot be
204 honored by a kernel.
205 """
206 if isinstance(pat, re.Pattern):
207 # a compiled str pattern always carries re.UNICODE, which pyarrow
208 # assumes for strings anyway
209 flags |= pat.flags & ~re.UNICODE
210 pattern = pat.pattern
211 else:
212 pattern = pat
213 if flags & re.IGNORECASE:
214 case = False
215 flags &= ~re.IGNORECASE
216 return pattern, case, flags
218 def _str_len(self):
219 result = pc.utf8_length(self._pa_array)
220 return self._convert_int_result(result)
222 def _str_lower(self) -> Self:
223 return self._from_pyarrow_array(pc.utf8_lower(self._pa_array))
225 def _str_upper(self) -> Self:
226 return self._from_pyarrow_array(pc.utf8_upper(self._pa_array))
228 def _str_strip(self, to_strip=None) -> Self:
229 if to_strip is None:
230 result = pc.utf8_trim_whitespace(self._pa_array)
231 else:
232 result = pc.utf8_trim(self._pa_array, characters=to_strip)
233 return self._from_pyarrow_array(result)
235 def _str_lstrip(self, to_strip=None) -> Self:
236 if to_strip is None:
237 result = pc.utf8_ltrim_whitespace(self._pa_array)
238 else:
239 result = pc.utf8_ltrim(self._pa_array, characters=to_strip)
240 return self._from_pyarrow_array(result)
242 def _str_rstrip(self, to_strip=None) -> Self:
243 if to_strip is None:
244 result = pc.utf8_rtrim_whitespace(self._pa_array)
245 else:
246 result = pc.utf8_rtrim(self._pa_array, characters=to_strip)
247 return self._from_pyarrow_array(result)
249 def _str_pad(
250 self,
251 width: int,
252 side: Literal["left", "right", "both"] = "left",
253 fillchar: str = " ",
254 ) -> Self:
255 if width < 0:
256 width = 0
257 if side == "left":
258 pa_pad = pc.utf8_lpad
259 elif side == "right":
260 pa_pad = pc.utf8_rpad
261 elif side == "both":
262 if pa_version_under17p0:
263 # GH#59624 fall back to object dtype
264 from pandas import array
266 obj_arr = self.astype(object, copy=False) # type: ignore[attr-defined]
267 obj = array(obj_arr, dtype=object)
268 result = obj._str_pad(width, side, fillchar) # type: ignore[attr-defined]
269 return type(self)._from_sequence(result, dtype=self.dtype) # type: ignore[attr-defined]
270 else:
271 # GH#54792
272 # https://github.com/apache/arrow/issues/15053#issuecomment-2317032347
273 lean_left = (width % 2) == 0
274 pa_pad = partial(pc.utf8_center, lean_left_on_odd_padding=lean_left)
275 else:
276 raise ValueError(
277 f"Invalid side: {side}. Side must be one of 'left', 'right', 'both'"
278 )
279 return self._from_pyarrow_array(
280 pa_pad(self._pa_array, width=width, padding=fillchar)
281 )
283 def _str_zfill(self, width: int) -> Self:
284 if pa_version_under21p0:
285 predicate = lambda val: val.zfill(width)
286 result = self._apply_elementwise(predicate)
287 return self._from_pyarrow_array(
288 pa.chunked_array(result, type=self._pa_array.type)
289 )
290 # pc.utf8_zfill raises on a negative width, while str.zfill returns the
291 # string unchanged -> clamping to zero to get that behaviour (GH#69486)
292 width = max(width, 0)
293 return self._from_pyarrow_array(pc.utf8_zfill(self._pa_array, width))
295 def _str_normalize(self, form: Literal["NFC", "NFD", "NFKC", "NFKD"]) -> Self:
296 if form not in ("NFC", "NFD", "NFKC", "NFKD"):
297 raise ValueError("invalid normalization form")
298 if form in ("NFC", "NFKC"):
299 # GH#64359 pc.utf8_normalize only decomposes; it skips the canonical
300 # composition step, so for the composing forms it returns decomposed
301 # output. Fall back to unicodedata for these.
302 predicate = lambda val: unicodedata.normalize(form, val)
303 result = self._apply_elementwise(predicate)
304 return self._from_pyarrow_array(
305 pa.chunked_array(result, type=self._pa_array.type)
306 )
307 return self._from_pyarrow_array(pc.utf8_normalize(self._pa_array, form=form))
309 def _str_get(self, i: int) -> Self:
310 lengths = pc.utf8_length(self._pa_array)
311 if i >= 0:
312 out_of_bounds = pc.greater_equal(i, lengths)
313 start = i
314 stop = i + 1
315 step = 1
316 else:
317 out_of_bounds = pc.greater(-i, lengths)
318 start = i
319 stop = i - 1
320 step = -1
321 not_out_of_bounds = pc.invert(out_of_bounds.fill_null(True))
322 selected = pc.utf8_slice_codeunits(
323 self._pa_array, start=start, stop=stop, step=step
324 )
325 null_value = pa.scalar(None, type=self._pa_array.type)
326 result = pc.if_else(not_out_of_bounds, selected, null_value)
327 return self._from_pyarrow_array(result)
329 def _str_slice(
330 self, start: int | None = None, stop: int | None = None, step: int | None = None
331 ) -> Self:
332 if start is None:
333 if step is not None and step < 0:
334 # GH#59710
335 start = -1
336 else:
337 start = 0
338 if step is None:
339 step = 1
340 return self._from_pyarrow_array(
341 pc.utf8_slice_codeunits(self._pa_array, start=start, stop=stop, step=step)
342 )
344 def _str_getitem(self, key: slice | int) -> Self:
345 if isinstance(key, slice):
346 return self._str_slice(start=key.start, stop=key.stop, step=key.step)
347 else:
348 return self._str_get(key)
350 def _str_slice_replace(
351 self, start: int | None = None, stop: int | None = None, repl: str | None = None
352 ) -> Self:
353 if repl is None:
354 repl = ""
355 if start is None:
356 start = 0
357 if stop is None:
358 stop = np.iinfo(np.int64).max
359 return self._from_pyarrow_array(
360 pc.utf8_replace_slice(self._pa_array, start, stop, repl)
361 )
363 def _str_replace(
364 self,
365 pat: str | re.Pattern[str],
366 repl: str | Callable[..., Any],
367 n: int = -1,
368 case: bool = True,
369 flags: int = 0,
370 regex: bool = True,
371 ) -> Self:
372 if (
373 isinstance(pat, re.Pattern)
374 or callable(repl)
375 or not case
376 or flags
377 or (isinstance(repl, str) and r"\g<" in repl)
378 ):
379 raise NotImplementedError(
380 "replace is not supported with a re.Pattern, callable repl, "
381 "case=False, flags!=0, or when the replacement string contains "
382 "named group references (\\g<...>)"
383 )
385 if pat == "":
386 # pyarrow hangs for empty patterns
387 # (https://github.com/apache/arrow/issues/39149)
388 # use same func definition as ObjectStringArrayMixin._str_replace
389 if regex:
390 count = n if n >= 0 else 0
391 func = lambda val: re.sub(pat, repl, val, count=count)
392 else:
393 func = lambda val: val.replace(pat, repl, n)
395 result = self._apply_elementwise(func)
396 return self._from_pyarrow_array(
397 pa.chunked_array(result, type=self._pa_array.type)
398 )
400 func = pc.replace_substring_regex if regex else pc.replace_substring
401 # https://github.com/apache/arrow/issues/39149
402 # GH 56404, unexpected behavior with negative max_replacements with pyarrow.
403 pa_max_replacements = None if n < 0 else n
404 result = func(
405 self._pa_array,
406 pattern=pat,
407 replacement=repl,
408 max_replacements=pa_max_replacements,
409 )
410 return self._from_pyarrow_array(result)
412 def _str_capitalize(self) -> Self:
413 return self._from_pyarrow_array(pc.utf8_capitalize(self._pa_array))
415 def _str_title(self) -> Self:
416 return self._from_pyarrow_array(pc.utf8_title(self._pa_array))
418 def _str_swapcase(self) -> Self:
419 return self._from_pyarrow_array(pc.utf8_swapcase(self._pa_array))
421 def _str_removeprefix(self, prefix: str):
422 if prefix == "":
423 return self._from_pyarrow_array(self._pa_array)
424 starts_with = pc.starts_with(self._pa_array, pattern=prefix)
425 removed = pc.utf8_slice_codeunits(self._pa_array, len(prefix))
426 result = pc.if_else(starts_with, removed, self._pa_array)
427 return self._from_pyarrow_array(result)
429 def _str_removesuffix(self, suffix: str):
430 if suffix == "":
431 return self._from_pyarrow_array(self._pa_array)
432 ends_with = pc.ends_with(self._pa_array, pattern=suffix)
433 removed = pc.utf8_slice_codeunits(self._pa_array, 0, stop=-len(suffix))
434 result = pc.if_else(ends_with, removed, self._pa_array)
435 return self._from_pyarrow_array(result)
437 def _str_startswith(
438 self, pat: str | tuple[str, ...], na: Scalar | lib.NoDefault = lib.no_default
439 ):
440 if isinstance(pat, str):
441 result = pc.starts_with(self._pa_array, pattern=pat)
442 elif len(pat) == 0:
443 # For empty tuple we return null for missing values and False
444 # for valid values.
445 result = pc.if_else(pc.is_null(self._pa_array), None, False)
446 else:
447 result = pc.starts_with(self._pa_array, pattern=pat[0])
449 for p in pat[1:]:
450 result = pc.or_(result, pc.starts_with(self._pa_array, pattern=p))
451 return self._convert_bool_result(result, na=na, method_name="startswith")
453 def _str_endswith(
454 self, pat: str | tuple[str, ...], na: Scalar | lib.NoDefault = lib.no_default
455 ):
456 if isinstance(pat, str):
457 result = pc.ends_with(self._pa_array, pattern=pat)
458 elif len(pat) == 0:
459 # For empty tuple we return null for missing values and False
460 # for valid values.
461 result = pc.if_else(pc.is_null(self._pa_array), None, False)
462 else:
463 result = pc.ends_with(self._pa_array, pattern=pat[0])
465 for p in pat[1:]:
466 result = pc.or_(result, pc.ends_with(self._pa_array, pattern=p))
467 return self._convert_bool_result(result, na=na, method_name="endswith")
469 def _str_isalnum(self):
470 result = pc.utf8_is_alnum(self._pa_array)
471 return self._convert_bool_result(result)
473 def _str_isalpha(self):
474 result = pc.utf8_is_alpha(self._pa_array)
475 return self._convert_bool_result(result)
477 def _str_isascii(self):
478 result = pc.string_is_ascii(self._pa_array)
479 return self._convert_bool_result(result)
481 def _str_isdecimal(self):
482 result = pc.utf8_is_decimal(self._pa_array)
483 return self._convert_bool_result(result)
485 def _str_isdigit(self):
486 if pa_version_under21p0:
487 # https://github.com/pandas-dev/pandas/issues/61466
488 res_list = self._apply_elementwise(str.isdigit)
489 return self._convert_bool_result(
490 pa.chunked_array(res_list, type=pa.bool_())
491 )
492 result = pc.utf8_is_digit(self._pa_array)
493 return self._convert_bool_result(result)
495 def _str_islower(self):
496 result = pc.utf8_is_lower(self._pa_array)
497 return self._convert_bool_result(result)
499 def _str_isnumeric(self):
500 result = pc.utf8_is_numeric(self._pa_array)
501 return self._convert_bool_result(result)
503 def _str_isspace(self):
504 result = pc.utf8_is_space(self._pa_array)
505 return self._convert_bool_result(result)
507 def _str_istitle(self):
508 result = pc.utf8_is_title(self._pa_array)
509 return self._convert_bool_result(result)
511 def _str_isupper(self):
512 result = pc.utf8_is_upper(self._pa_array)
513 return self._convert_bool_result(result)
515 def _str_contains(
516 self,
517 pat,
518 case: bool = True,
519 flags: int = 0,
520 na: Scalar | lib.NoDefault = lib.no_default,
521 regex: bool = True,
522 ):
523 if regex:
524 # GH#66348 a compiled `pat` reaches us from ArrowExtensionArray, which
525 # unlike ArrowStringArray does not unwrap it before dispatching here.
526 pat, case, flags = self._unwrap_re_pattern(pat, case, flags)
528 if flags:
529 raise NotImplementedError(f"contains not implemented with {flags=}")
531 if regex:
532 pa_contains = pc.match_substring_regex
533 else:
534 pa_contains = pc.match_substring
535 result = pa_contains(self._pa_array, pat, ignore_case=not case)
536 return self._convert_bool_result(result, na=na, method_name="contains")
538 def _str_match(
539 self,
540 pat: str | re.Pattern[str],
541 case: bool = True,
542 flags: int = 0,
543 na: Scalar | lib.NoDefault = lib.no_default,
544 ):
545 # GH#63108 the accessor pre-compiles `pat` whenever the user passes `flags`
546 pattern, case, flags = self._unwrap_re_pattern(pat, case, flags)
548 if pattern.startswith("^"):
549 pattern = pattern[1:]
550 pattern = f"^({pattern})"
551 return ArrowStringArrayMixin._str_contains(
552 self, pattern, case, flags, na, regex=True
553 )
555 def _str_fullmatch(
556 self,
557 pat: str | re.Pattern[str],
558 case: bool = True,
559 flags: int = 0,
560 na: Scalar | lib.NoDefault = lib.no_default,
561 ):
562 pat, case, flags = self._unwrap_re_pattern(pat, case, flags)
564 if (not pat.endswith("$") or pat.endswith("\\$")) and not pat.startswith("^"):
565 pat = f"^({pat})$"
566 elif not pat.endswith("$") or pat.endswith("\\$"):
567 pat = f"^({pat[1:]})$"
568 elif not pat.startswith("^"):
569 pat = f"^({pat[0:-1]})$"
570 return ArrowStringArrayMixin._str_match(self, pat, case, flags, na)
572 def _str_find(self, sub: str, start: int = 0, end: int | None = None):
573 if self._has_non_ascii():
574 # GH#64123 - pc.find_substring returns byte offsets instead of
575 # character offsets for multi-byte UTF-8 characters, so we fall back
576 # to Python str.find which correctly returns character offsets.
577 res_list = self._apply_elementwise(lambda val: val.find(sub, start, end))
578 return self._convert_int_result(pa.chunked_array(res_list, type=pa.int64()))
580 if (start == 0 or start is None) and end is None:
581 result = pc.find_substring(self._pa_array, sub)
582 else:
583 if sub == "":
584 # GH#56792
585 res_list = self._apply_elementwise(
586 lambda val: val.find(sub, start, end)
587 )
588 return self._convert_int_result(
589 pa.chunked_array(res_list, type=pa.int64())
590 )
591 if start is None:
592 start_offset = 0
593 start = 0
594 elif start < 0:
595 start_offset = pc.add(start, pc.utf8_length(self._pa_array))
596 start_offset = pc.if_else(pc.less(start_offset, 0), 0, start_offset)
597 else:
598 start_offset = start
599 slices = pc.utf8_slice_codeunits(self._pa_array, start, stop=end)
600 result = pc.find_substring(slices, sub)
601 found = pc.not_equal(result, pa.scalar(-1, type=result.type))
602 offset_result = pc.add(result, start_offset)
603 result = pc.if_else(found, offset_result, -1)
604 result = result.cast(pa.int64())
605 return self._convert_int_result(result)
607 def _str_partition_expand(self, sep: str) -> pa.ChunkedArray:
608 """
609 Split each string on the first occurrence of ``sep``.
611 Returns a ``list<string>`` array holding one three-element list per row
612 -- the part before the separator, the separator, and the part after --
613 which ``StringMethods._wrap_result`` expands into three columns. Rows
614 without ``sep`` get two empty strings, matching ``str.partition``.
616 The caller wraps this in an :class:`ArrowExtensionArray`; the rows are
617 lists, so it is not of the calling array's own type.
618 """
619 if not sep:
620 # pyarrow reports this as "Empty separator"; keep str.partition's
621 # wording so every dtype raises the same way
622 raise ValueError("empty separator")
624 str_type = self._pa_array.type
625 chunks = [
626 self._partition_chunk(chunk, sep, str_type)
627 for chunk in self._pa_array.chunks
628 ]
629 return pa.chunked_array(chunks, type=pa.list_(str_type))
631 @staticmethod
632 def _partition_chunk(chunk: pa.Array, sep: str, str_type: pa.DataType) -> pa.Array:
633 """
634 Build the ``list<string>`` rows for one chunk of :meth:`_str_partition_expand`.
636 Working a chunk at a time keeps the concatenation below within the
637 offset width of ``str_type``, which matters for columns near the 2 GiB
638 limit of 32-bit ``string``.
639 """
640 # max_splits=1 gives [before] when sep is absent, else [before, after];
641 # padding to a fixed two elements makes the tail null in the first case
642 split = pc.split_pattern(chunk, sep, max_splits=1)
643 pieces = pc.list_slice(split, 0, 2, return_fixed_size_list=True)
645 before = pc.list_element(pieces, 0)
646 tail = pc.list_element(pieces, 1)
647 after = pc.fill_null(tail, pa.scalar("", type=str_type))
648 middle = pc.if_else(
649 pc.is_valid(tail),
650 pa.scalar(sep, type=str_type),
651 pa.scalar("", type=str_type),
652 )
654 n = len(chunk)
655 # Interleave the three columns into [before[0], sep[0], after[0],
656 # before[1], ...]. Arrow has no kernel that zips arrays into rows, so
657 # take from their concatenation. Both ranges below are plain index
658 # arithmetic: NumPy builds them in one strided pass, while pyarrow has
659 # no arange at all before 21.0 and would need four kernels for the
660 # interleave, so they stay NumPy at every supported version.
661 values = pa.concat_arrays([before, middle, after])
662 indices = np.arange(3 * n, dtype=np.int64).reshape(3, n).T.reshape(-1)
663 values = values.take(pa.array(indices))
664 # int32 rather than int64 offsets, but built wide so that a chunk with
665 # more than 2**31 / 3 rows raises instead of wrapping around
666 offsets = pa.array(np.arange(0, 3 * n + 1, 3, dtype=np.int64), type=pa.int32())
667 # a null string partitions to a null row, not to a row of nulls
668 mask = pc.is_null(before) if before.null_count else None
669 return pa.ListArray.from_arrays(offsets, values, mask=mask)