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 pa_version_under25p0 = _palv < Version("25.0.0") 

27 pa_version_under26p0 = _palv < Version("26.0.0") 

28 PYARROW_INSTALLED = True 

29 HAS_PYARROW = _palv >= Version(PYARROW_MIN_VERSION) 

30except ImportError: 

31 pa_version_under14p0 = True 

32 pa_version_under14p1 = True 

33 pa_version_under15p0 = True 

34 pa_version_under16p0 = True 

35 pa_version_under17p0 = True 

36 pa_version_under18p0 = True 

37 pa_version_under19p0 = True 

38 pa_version_under20p0 = True 

39 pa_version_under21p0 = True 

40 pa_version_under22p0 = True 

41 pa_version_under23p0 = True 

42 pa_version_under25p0 = True 

43 pa_version_under26p0 = True 

44 PYARROW_INSTALLED = False 

45 HAS_PYARROW = False 

46 

47 

48def _safe_fill_null( 

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

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

51 """ 

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

53 

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

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

56 the standard pyarrow.compute.fill_null. 

57 

58 Parameters 

59 ---------- 

60 arr : pyarrow.Array | pyarrow.ChunkedArray 

61 Input array with potential null values. 

62 fill_value : Any 

63 Value to fill nulls with. 

64 

65 Returns 

66 ------- 

67 pyarrow.Array | pyarrow.ChunkedArray 

68 Array with nulls filled with fill_value. 

69 """ 

70 import pyarrow.compute as pc 

71 

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

73 use_fallback = ( 

74 HAS_PYARROW and is_windows and not pa_version_under21p0 and pa_version_under22p0 

75 ) 

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

77 return pc.fill_null(arr, fill_value) 

78 

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

80 

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

82 

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

84 mask = pc.is_null(arr) 

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

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

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

88 

89 if isinstance(arr, pa.ChunkedArray): 

90 return pa.chunked_array( 

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

92 ) 

93 return fill_null_duration(arr, fill_scalar) 

94 

95 if isinstance(arr, pa.ChunkedArray): 

96 return pa.chunked_array( 

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

98 ) 

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