Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/core/arrays/_arrow_string_mixins.py: 21%

1from __future__ import annotations 

2 

3from functools import partial 

4import re 

5from typing import ( 

6 TYPE_CHECKING, 

7 Any, 

8 Literal, 

9 Self, 

10) 

11 

12import numpy as np 

13 

14from pandas._libs import lib 

15from pandas.compat import ( 

16 HAS_PYARROW, 

17 pa_version_under17p0, 

18 pa_version_under21p0, 

19) 

20 

21if HAS_PYARROW: 

22 import pyarrow as pa 

23 import pyarrow.compute as pc 

24 

25if TYPE_CHECKING: 

26 from collections.abc import Callable 

27 

28 from pandas._typing import Scalar 

29 

30 

31class ArrowStringArrayMixin: 

32 _pa_array: pa.ChunkedArray 

33 

34 def __init__(self, *args, **kwargs) -> None: 

35 raise NotImplementedError 

36 

37 def _from_pyarrow_array(self, pa_array) -> Self: 

38 raise NotImplementedError 

39 

40 def _convert_bool_result(self, result, na=lib.no_default, method_name=None): 

41 # Convert a bool-dtype result to the appropriate result type 

42 raise NotImplementedError 

43 

44 def _convert_int_result(self, result): 

45 # Convert an integer-dtype result to the appropriate result type 

46 raise NotImplementedError 

47 

48 def _apply_elementwise(self, func: Callable) -> list[list[Any]]: 

49 raise NotImplementedError 

50 

51 @staticmethod 

52 def _has_unsupported_regex(pat: str | re.Pattern) -> bool: 

53 """ 

54 Determine if regex pattern contains features not supported by RE2 / pyarrow. 

55 

56 This includes lookaround (lookahead or lookbehind) assertions and 

57 backreferences. 

58 

59 Parameters 

60 ---------- 

61 pat: str | re.Pattern 

62 Regex pattern. 

63 

64 Returns 

65 ------- 

66 bool 

67 Whether `pat` contains a lookahead or lookbehind. 

68 """ 

69 try: 

70 # error: Module "re" has no attribute "_parser" 

71 from re import _parser # type: ignore[attr-defined] 

72 

73 regex_parser = _parser.parse 

74 except Exception as err: 

75 raise type(err)( 

76 "Incompatible version for regex; you will need to upgrade pandas " 

77 "or downgrade Python" 

78 ) from err 

79 

80 def has_unsupported_code(tokens): 

81 # For certain op codes we need to recurse. 

82 for op_code, argument in tokens: 

83 if ( 

84 ( 

85 op_code == _parser.SUBPATTERN 

86 and has_unsupported_code(argument[3]) 

87 ) 

88 or ( 

89 op_code == _parser.BRANCH 

90 and any(has_unsupported_code(tokens) for tokens in argument[1]) 

91 ) 

92 or ( 

93 op_code 

94 in [_parser.ASSERT_NOT, _parser.ASSERT, _parser.GROUPREF] 

95 ) 

96 ): 

97 return True 

98 return False 

99 

100 str_pat = pat.pattern if isinstance(pat, re.Pattern) else pat 

101 try: 

102 tokens = regex_parser(str_pat) 

103 except re.error: 

104 # Pattern not valid for Python's re (e.g. RE2 syntax like \x{...} or \p) 

105 # Let the pyarrow backend handle it. 

106 return False 

107 return has_unsupported_code(tokens) 

108 

109 def _str_len(self): 

110 result = pc.utf8_length(self._pa_array) 

111 return self._convert_int_result(result) 

112 

113 def _str_lower(self) -> Self: 

114 return self._from_pyarrow_array(pc.utf8_lower(self._pa_array)) 

115 

116 def _str_upper(self) -> Self: 

117 return self._from_pyarrow_array(pc.utf8_upper(self._pa_array)) 

118 

119 def _str_strip(self, to_strip=None) -> Self: 

120 if to_strip is None: 

121 result = pc.utf8_trim_whitespace(self._pa_array) 

122 else: 

123 result = pc.utf8_trim(self._pa_array, characters=to_strip) 

124 return self._from_pyarrow_array(result) 

125 

126 def _str_lstrip(self, to_strip=None) -> Self: 

127 if to_strip is None: 

128 result = pc.utf8_ltrim_whitespace(self._pa_array) 

129 else: 

130 result = pc.utf8_ltrim(self._pa_array, characters=to_strip) 

131 return self._from_pyarrow_array(result) 

132 

133 def _str_rstrip(self, to_strip=None) -> Self: 

134 if to_strip is None: 

135 result = pc.utf8_rtrim_whitespace(self._pa_array) 

136 else: 

137 result = pc.utf8_rtrim(self._pa_array, characters=to_strip) 

138 return self._from_pyarrow_array(result) 

139 

140 def _str_pad( 

141 self, 

142 width: int, 

143 side: Literal["left", "right", "both"] = "left", 

144 fillchar: str = " ", 

145 ) -> Self: 

146 if side == "left": 

147 pa_pad = pc.utf8_lpad 

148 elif side == "right": 

149 pa_pad = pc.utf8_rpad 

150 elif side == "both": 

151 if pa_version_under17p0: 

152 # GH#59624 fall back to object dtype 

153 from pandas import array 

154 

155 obj_arr = self.astype(object, copy=False) # type: ignore[attr-defined] 

156 obj = array(obj_arr, dtype=object) 

157 result = obj._str_pad(width, side, fillchar) # type: ignore[attr-defined] 

158 return type(self)._from_sequence(result, dtype=self.dtype) # type: ignore[attr-defined] 

159 else: 

160 # GH#54792 

161 # https://github.com/apache/arrow/issues/15053#issuecomment-2317032347 

162 lean_left = (width % 2) == 0 

163 pa_pad = partial(pc.utf8_center, lean_left_on_odd_padding=lean_left) 

164 else: 

165 raise ValueError( 

166 f"Invalid side: {side}. Side must be one of 'left', 'right', 'both'" 

167 ) 

168 return self._from_pyarrow_array( 

169 pa_pad(self._pa_array, width=width, padding=fillchar) 

170 ) 

171 

172 def _str_get(self, i: int) -> Self: 

173 lengths = pc.utf8_length(self._pa_array) 

174 if i >= 0: 

175 out_of_bounds = pc.greater_equal(i, lengths) 

176 start = i 

177 stop = i + 1 

178 step = 1 

179 else: 

180 out_of_bounds = pc.greater(-i, lengths) 

181 start = i 

182 stop = i - 1 

183 step = -1 

184 not_out_of_bounds = pc.invert(out_of_bounds.fill_null(True)) 

185 selected = pc.utf8_slice_codeunits( 

186 self._pa_array, start=start, stop=stop, step=step 

187 ) 

188 null_value = pa.scalar(None, type=self._pa_array.type) 

189 result = pc.if_else(not_out_of_bounds, selected, null_value) 

190 return self._from_pyarrow_array(result) 

191 

192 def _str_slice( 

193 self, start: int | None = None, stop: int | None = None, step: int | None = None 

194 ) -> Self: 

195 if start is None: 

196 if step is not None and step < 0: 

197 # GH#59710 

198 start = -1 

199 else: 

200 start = 0 

201 if step is None: 

202 step = 1 

203 return self._from_pyarrow_array( 

204 pc.utf8_slice_codeunits(self._pa_array, start=start, stop=stop, step=step) 

205 ) 

206 

207 def _str_getitem(self, key: slice | int) -> Self: 

208 if isinstance(key, slice): 

209 return self._str_slice(start=key.start, stop=key.stop, step=key.step) 

210 else: 

211 return self._str_get(key) 

212 

213 def _str_slice_replace( 

214 self, start: int | None = None, stop: int | None = None, repl: str | None = None 

215 ) -> Self: 

216 if repl is None: 

217 repl = "" 

218 if start is None: 

219 start = 0 

220 if stop is None: 

221 stop = np.iinfo(np.int64).max 

222 return self._from_pyarrow_array( 

223 pc.utf8_replace_slice(self._pa_array, start, stop, repl) 

224 ) 

225 

226 def _str_replace( 

227 self, 

228 pat: str | re.Pattern, 

229 repl: str | Callable, 

230 n: int = -1, 

231 case: bool = True, 

232 flags: int = 0, 

233 regex: bool = True, 

234 ) -> Self: 

235 if ( 

236 isinstance(pat, re.Pattern) 

237 or callable(repl) 

238 or not case 

239 or flags 

240 or (isinstance(repl, str) and r"\g<" in repl) 

241 ): 

242 raise NotImplementedError( 

243 "replace is not supported with a re.Pattern, callable repl, " 

244 "case=False, flags!=0, or when the replacement string contains " 

245 "named group references (\\g<...>)" 

246 ) 

247 

248 if pat == "": 

249 # pyarrow hangs for empty patterns 

250 # (https://github.com/apache/arrow/issues/39149) 

251 # use same func definition as ObjectStringArrayMixin._str_replace 

252 if regex: 

253 count = n if n >= 0 else 0 

254 func = lambda val: re.sub(pat, repl, val, count=count) 

255 else: 

256 func = lambda val: val.replace(pat, repl, n) 

257 

258 result = self._apply_elementwise(func) 

259 return self._from_pyarrow_array( 

260 pa.chunked_array(result, type=self._pa_array.type) 

261 ) 

262 

263 func = pc.replace_substring_regex if regex else pc.replace_substring 

264 # https://github.com/apache/arrow/issues/39149 

265 # GH 56404, unexpected behavior with negative max_replacements with pyarrow. 

266 pa_max_replacements = None if n < 0 else n 

267 result = func( 

268 self._pa_array, 

269 pattern=pat, 

270 replacement=repl, 

271 max_replacements=pa_max_replacements, 

272 ) 

273 return self._from_pyarrow_array(result) 

274 

275 def _str_capitalize(self) -> Self: 

276 return self._from_pyarrow_array(pc.utf8_capitalize(self._pa_array)) 

277 

278 def _str_title(self) -> Self: 

279 return self._from_pyarrow_array(pc.utf8_title(self._pa_array)) 

280 

281 def _str_swapcase(self) -> Self: 

282 return self._from_pyarrow_array(pc.utf8_swapcase(self._pa_array)) 

283 

284 def _str_removeprefix(self, prefix: str): 

285 if prefix == "": 

286 return self._from_pyarrow_array(self._pa_array) 

287 starts_with = pc.starts_with(self._pa_array, pattern=prefix) 

288 removed = pc.utf8_slice_codeunits(self._pa_array, len(prefix)) 

289 result = pc.if_else(starts_with, removed, self._pa_array) 

290 return self._from_pyarrow_array(result) 

291 

292 def _str_removesuffix(self, suffix: str): 

293 if suffix == "": 

294 return self._from_pyarrow_array(self._pa_array) 

295 ends_with = pc.ends_with(self._pa_array, pattern=suffix) 

296 removed = pc.utf8_slice_codeunits(self._pa_array, 0, stop=-len(suffix)) 

297 result = pc.if_else(ends_with, removed, self._pa_array) 

298 return self._from_pyarrow_array(result) 

299 

300 def _str_startswith( 

301 self, pat: str | tuple[str, ...], na: Scalar | lib.NoDefault = lib.no_default 

302 ): 

303 if isinstance(pat, str): 

304 result = pc.starts_with(self._pa_array, pattern=pat) 

305 elif len(pat) == 0: 

306 # For empty tuple we return null for missing values and False 

307 # for valid values. 

308 result = pc.if_else(pc.is_null(self._pa_array), None, False) 

309 else: 

310 result = pc.starts_with(self._pa_array, pattern=pat[0]) 

311 

312 for p in pat[1:]: 

313 result = pc.or_(result, pc.starts_with(self._pa_array, pattern=p)) 

314 return self._convert_bool_result(result, na=na, method_name="startswith") 

315 

316 def _str_endswith( 

317 self, pat: str | tuple[str, ...], na: Scalar | lib.NoDefault = lib.no_default 

318 ): 

319 if isinstance(pat, str): 

320 result = pc.ends_with(self._pa_array, pattern=pat) 

321 elif len(pat) == 0: 

322 # For empty tuple we return null for missing values and False 

323 # for valid values. 

324 result = pc.if_else(pc.is_null(self._pa_array), None, False) 

325 else: 

326 result = pc.ends_with(self._pa_array, pattern=pat[0]) 

327 

328 for p in pat[1:]: 

329 result = pc.or_(result, pc.ends_with(self._pa_array, pattern=p)) 

330 return self._convert_bool_result(result, na=na, method_name="endswith") 

331 

332 def _str_isalnum(self): 

333 result = pc.utf8_is_alnum(self._pa_array) 

334 return self._convert_bool_result(result) 

335 

336 def _str_isalpha(self): 

337 result = pc.utf8_is_alpha(self._pa_array) 

338 return self._convert_bool_result(result) 

339 

340 def _str_isascii(self): 

341 result = pc.string_is_ascii(self._pa_array) 

342 return self._convert_bool_result(result) 

343 

344 def _str_isdecimal(self): 

345 result = pc.utf8_is_decimal(self._pa_array) 

346 return self._convert_bool_result(result) 

347 

348 def _str_isdigit(self): 

349 if pa_version_under21p0: 

350 # https://github.com/pandas-dev/pandas/issues/61466 

351 res_list = self._apply_elementwise(str.isdigit) 

352 return self._convert_bool_result( 

353 pa.chunked_array(res_list, type=pa.bool_()) 

354 ) 

355 result = pc.utf8_is_digit(self._pa_array) 

356 return self._convert_bool_result(result) 

357 

358 def _str_islower(self): 

359 result = pc.utf8_is_lower(self._pa_array) 

360 return self._convert_bool_result(result) 

361 

362 def _str_isnumeric(self): 

363 result = pc.utf8_is_numeric(self._pa_array) 

364 return self._convert_bool_result(result) 

365 

366 def _str_isspace(self): 

367 result = pc.utf8_is_space(self._pa_array) 

368 return self._convert_bool_result(result) 

369 

370 def _str_istitle(self): 

371 result = pc.utf8_is_title(self._pa_array) 

372 return self._convert_bool_result(result) 

373 

374 def _str_isupper(self): 

375 result = pc.utf8_is_upper(self._pa_array) 

376 return self._convert_bool_result(result) 

377 

378 def _str_contains( 

379 self, 

380 pat, 

381 case: bool = True, 

382 flags: int = 0, 

383 na: Scalar | lib.NoDefault = lib.no_default, 

384 regex: bool = True, 

385 ): 

386 if flags: 

387 raise NotImplementedError(f"contains not implemented with {flags=}") 

388 

389 if regex: 

390 pa_contains = pc.match_substring_regex 

391 else: 

392 pa_contains = pc.match_substring 

393 result = pa_contains(self._pa_array, pat, ignore_case=not case) 

394 return self._convert_bool_result(result, na=na, method_name="contains") 

395 

396 def _str_match( 

397 self, 

398 pat: str, 

399 case: bool = True, 

400 flags: int = 0, 

401 na: Scalar | lib.NoDefault = lib.no_default, 

402 ): 

403 if pat.startswith("^"): 

404 pat = pat[1:] 

405 pat = f"^({pat})" 

406 return ArrowStringArrayMixin._str_contains( 

407 self, pat, case, flags, na, regex=True 

408 ) 

409 

410 def _str_fullmatch( 

411 self, 

412 pat: str, 

413 case: bool = True, 

414 flags: int = 0, 

415 na: Scalar | lib.NoDefault = lib.no_default, 

416 ): 

417 if (not pat.endswith("$") or pat.endswith("\\$")) and not pat.startswith("^"): 

418 pat = f"^({pat})$" 

419 elif not pat.endswith("$") or pat.endswith("\\$"): 

420 pat = f"^({pat[1:]})$" 

421 elif not pat.startswith("^"): 

422 pat = f"^({pat[0:-1]})$" 

423 return ArrowStringArrayMixin._str_match(self, pat, case, flags, na) 

424 

425 def _str_find(self, sub: str, start: int = 0, end: int | None = None): 

426 if not pc.all(pc.string_is_ascii(self._pa_array)).as_py(): 

427 # GH#64123 - pc.find_substring returns byte offsets instead of 

428 # character offsets for multi-byte UTF-8 characters, so we fall back 

429 # to Python str.find which correctly returns character offsets. 

430 res_list = self._apply_elementwise(lambda val: val.find(sub, start, end)) 

431 return self._convert_int_result(pa.chunked_array(res_list)) 

432 

433 if (start == 0 or start is None) and end is None: 

434 result = pc.find_substring(self._pa_array, sub) 

435 else: 

436 if sub == "": 

437 # GH#56792 

438 res_list = self._apply_elementwise( 

439 lambda val: val.find(sub, start, end) 

440 ) 

441 return self._convert_int_result(pa.chunked_array(res_list)) 

442 if start is None: 

443 start_offset = 0 

444 start = 0 

445 elif start < 0: 

446 start_offset = pc.add(start, pc.utf8_length(self._pa_array)) 

447 start_offset = pc.if_else(pc.less(start_offset, 0), 0, start_offset) 

448 else: 

449 start_offset = start 

450 slices = pc.utf8_slice_codeunits(self._pa_array, start, stop=end) 

451 result = pc.find_substring(slices, sub) 

452 found = pc.not_equal(result, pa.scalar(-1, type=result.type)) 

453 offset_result = pc.add(result, start_offset) 

454 result = pc.if_else(found, offset_result, -1) 

455 result = result.cast(pa.int64()) 

456 return self._convert_int_result(result)