Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/compat/pyarrow.py: 42%

1"""support pyarrow compatibility across versions""" 

2 

3from __future__ import annotations 

4 

5import sys 

6from typing import Any 

7 

8from pandas.util.version import Version 

9 

10PYARROW_MIN_VERSION = "13.0.0" 

11try: 

12 import pyarrow as pa 

13 

14 _palv = Version(Version(pa.__version__).base_version) 

15 pa_version_under14p0 = _palv < Version("14.0.0") 

16 pa_version_under14p1 = _palv < Version("14.0.1") 

17 pa_version_under15p0 = _palv < Version("15.0.0") 

18 pa_version_under16p0 = _palv < Version("16.0.0") 

19 pa_version_under17p0 = _palv < Version("17.0.0") 

20 pa_version_under18p0 = _palv < Version("18.0.0") 

21 pa_version_under19p0 = _palv < Version("19.0.0") 

22 pa_version_under20p0 = _palv < Version("20.0.0") 

23 pa_version_under21p0 = _palv < Version("21.0.0") 

24 pa_version_under22p0 = _palv < Version("22.0.0") 

25 pa_version_under23p0 = _palv < Version("23.0.0") 

26 PYARROW_INSTALLED = True 

27 HAS_PYARROW = _palv >= Version(PYARROW_MIN_VERSION) 

28except ImportError: 

29 pa_version_under14p0 = True 

30 pa_version_under14p1 = True 

31 pa_version_under15p0 = True 

32 pa_version_under16p0 = True 

33 pa_version_under17p0 = True 

34 pa_version_under18p0 = True 

35 pa_version_under19p0 = True 

36 pa_version_under20p0 = True 

37 pa_version_under21p0 = True 

38 pa_version_under22p0 = True 

39 pa_version_under23p0 = True 

40 PYARROW_INSTALLED = False 

41 HAS_PYARROW = False 

42 

43 

44def _safe_fill_null( 

45 arr: pa.Array | pa.ChunkedArray, fill_value: Any 

46) -> pa.Array | pa.ChunkedArray: 

47 """ 

48 Safe wrapper for pyarrow.compute.fill_null with fallback for Windows + pyarrow 21. 

49 

50 pyarrow 21.0.0 on Windows has a bug in fill_null that incorrectly fills null values. 

51 This function uses a fallback implementation for that specific case, otherwise uses 

52 the standard pyarrow.compute.fill_null. 

53 

54 Parameters 

55 ---------- 

56 arr : pyarrow.Array | pyarrow.ChunkedArray 

57 Input array with potential null values. 

58 fill_value : Any 

59 Value to fill nulls with. 

60 

61 Returns 

62 ------- 

63 pyarrow.Array | pyarrow.ChunkedArray 

64 Array with nulls filled with fill_value. 

65 """ 

66 import pyarrow.compute as pc 

67 

68 is_windows = sys.platform in ["win32", "cygwin"] 

69 use_fallback = ( 

70 HAS_PYARROW and is_windows and not pa_version_under21p0 and pa_version_under22p0 

71 ) 

72 if not use_fallback or isinstance(fill_value, (pa.Array, pa.ChunkedArray)): 

73 return pc.fill_null(arr, fill_value) 

74 

75 fill_scalar = pa.scalar(fill_value, type=arr.type) 

76 

77 if pa.types.is_duration(arr.type): 

78 

79 def fill_null_duration(arr: pa.Array, fill_scalar: pa.Scalar) -> pa.Array: 

80 mask = pc.is_null(arr) 

81 zero_duration = pa.scalar(0, type=arr.type) 

82 arr_zeroed = pc.if_else(mask, zero_duration, arr) 

83 return pc.if_else(mask, fill_scalar, arr_zeroed) 

84 

85 if isinstance(arr, pa.ChunkedArray): 

86 return pa.chunked_array( 

87 [fill_null_duration(chunk, fill_scalar) for chunk in arr.chunks] 

88 ) 

89 return fill_null_duration(arr, fill_scalar) 

90 

91 if isinstance(arr, pa.ChunkedArray): 

92 return pa.chunked_array( 

93 [pc.if_else(pc.is_null(chunk), fill_scalar, chunk) for chunk in arr.chunks] 

94 ) 

95 return pc.if_else(pc.is_null(arr), fill_scalar, arr)