Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/io/orc.py: 25%

1"""orc compat""" 

2 

3from __future__ import annotations 

4 

5import io 

6from typing import ( 

7 TYPE_CHECKING, 

8 Any, 

9 Literal, 

10) 

11 

12from pandas._libs import lib 

13from pandas.compat._optional import import_optional_dependency 

14from pandas.util._decorators import set_module 

15from pandas.util._validators import check_dtype_backend 

16 

17from pandas.core.indexes.api import default_index 

18 

19from pandas.io._util import ( 

20 arrow_table_to_pandas, 

21 suppress_pyarrow_values_warning, 

22) 

23from pandas.io.common import ( 

24 get_handle, 

25 is_fsspec_url, 

26) 

27 

28if TYPE_CHECKING: 

29 import fsspec 

30 import pyarrow.fs 

31 

32 from pandas._typing import ( 

33 DtypeBackend, 

34 FilePath, 

35 ReadBuffer, 

36 WriteBuffer, 

37 ) 

38 

39 from pandas.core.frame import DataFrame 

40 

41 

42@set_module("pandas") 

43def read_orc( 

44 path: FilePath | ReadBuffer[bytes], 

45 columns: list[str] | None = None, 

46 dtype_backend: DtypeBackend | lib.NoDefault = lib.no_default, 

47 filesystem: pyarrow.fs.FileSystem | fsspec.spec.AbstractFileSystem | None = None, 

48 **kwargs: Any, 

49) -> DataFrame: 

50 """ 

51 Load an ORC object from the file path, returning a DataFrame. 

52 

53 This function requires the `pyarrow <https://arrow.apache.org/docs/python/>`_ 

54 library. 

55 

56 ORC is a columnar storage format 

57 that provides efficient compression and fast retrieval for analytical workloads. 

58 It allows reading specific columns, handling different filesystem 

59 types (such as local storage, cloud storage via fsspec, or pyarrow filesystem), 

60 and supports different data type backends, including `numpy_nullable` and `pyarrow`. 

61 

62 Parameters 

63 ---------- 

64 path : str, path object, or file-like object 

65 String, path object (implementing ``os.PathLike[str]``), or file-like 

66 object implementing a binary ``read()`` function. The string could be a URL. 

67 Valid URL schemes include http, ftp, s3, and file. For file URLs, a host is 

68 expected. A local file could be: 

69 ``file://localhost/path/to/table.orc``. 

70 

71 Certain URL schemes may require additional packages. For example, S3 

72 URLs require the ``s3fs`` library. See 

73 :ref:`install.optional_dependencies` for a full list. 

74 columns : list, default None 

75 If not None, only these columns will be read from the file. 

76 Output always follows the ordering of the file and not the columns list. 

77 This mirrors the original behaviour of 

78 :external+pyarrow:py:meth:`pyarrow.orc.ORCFile.read`. 

79 dtype_backend : {'numpy_nullable', 'pyarrow'} 

80 Back-end data type applied to the resultant :class:`DataFrame` 

81 (still experimental). If not specified, the default behavior 

82 is to not use nullable data types. If specified, the behavior 

83 is as follows: 

84 

85 * ``"numpy_nullable"``: returns nullable-dtype-backed :class:`DataFrame` 

86 * ``"pyarrow"``: returns pyarrow-backed nullable 

87 :class:`ArrowDtype` :class:`DataFrame` 

88 

89 .. versionadded:: 2.0 

90 

91 filesystem : fsspec or pyarrow filesystem, default None 

92 Filesystem object to use when reading the orc file. 

93 

94 .. versionadded:: 2.1.0 

95 

96 **kwargs 

97 Any additional kwargs are passed to pyarrow. 

98 

99 Returns 

100 ------- 

101 DataFrame 

102 DataFrame based on the ORC file. 

103 

104 See Also 

105 -------- 

106 read_csv : Read a comma-separated values (csv) file into a pandas DataFrame. 

107 read_excel : Read an Excel file into a pandas DataFrame. 

108 read_spss : Read an SPSS file into a pandas DataFrame. 

109 read_sas : Load a SAS file into a pandas DataFrame. 

110 read_feather : Load a feather-format object into a pandas DataFrame. 

111 

112 Notes 

113 ----- 

114 Before using this function you should read the :ref:`user guide about ORC <io.orc>` 

115 and :ref:`install optional dependencies <install.warn_orc>`. 

116 

117 If ``path`` is a URI scheme pointing to a local or remote file (e.g. "s3://"), 

118 a ``pyarrow.fs`` filesystem will be attempted to read the file. You can also pass a 

119 pyarrow or fsspec filesystem object into the filesystem keyword to override this 

120 behavior. 

121 

122 Examples 

123 -------- 

124 >>> result = pd.read_orc("example_pa.orc") # doctest: +SKIP 

125 """ 

126 # we require a newer version of pyarrow than we support for orc 

127 

128 orc = import_optional_dependency("pyarrow.orc") 

129 

130 check_dtype_backend(dtype_backend) 

131 

132 with get_handle(path, "rb", is_text=False) as handles: 

133 source = handles.handle 

134 if is_fsspec_url(path) and filesystem is None: 

135 pa = import_optional_dependency("pyarrow") 

136 pa_fs = import_optional_dependency("pyarrow.fs") 

137 try: 

138 filesystem, source = pa_fs.FileSystem.from_uri(path) 

139 except (TypeError, pa.ArrowInvalid): 

140 pass 

141 

142 pa_table = orc.read_table( 

143 source=source, columns=columns, filesystem=filesystem, **kwargs 

144 ) 

145 return arrow_table_to_pandas(pa_table, dtype_backend=dtype_backend) 

146 

147 

148def to_orc( 

149 df: DataFrame, 

150 path: FilePath | WriteBuffer[bytes] | None = None, 

151 *, 

152 engine: Literal["pyarrow"] = "pyarrow", 

153 index: bool | None = None, 

154 engine_kwargs: dict[str, Any] | None = None, 

155) -> bytes | None: 

156 """ 

157 Write a DataFrame to the ORC format. 

158 

159 Parameters 

160 ---------- 

161 df : DataFrame 

162 The dataframe to be written to ORC. Raises NotImplementedError 

163 if dtype of one or more columns is category, unsigned integers, 

164 intervals, periods or sparse. 

165 path : str, file-like object or None, default None 

166 If a string, it will be used as the root directory path 

167 when writing a partitioned dataset. By file-like object, 

168 we refer to objects with a write() method, such as a file handle 

169 (e.g. via builtin open function). If path is None, 

170 a bytes object is returned. 

171 

172 The string could be a URL. Valid URL schemes include http, ftp, s3, 

173 gs, and file. For file URLs, a host is expected. A local file could be: 

174 ``file://localhost/path/to/table.orc``. A remote example could be: 

175 ``s3://bucket/path/to/table.orc``. 

176 

177 Certain URL schemes may require additional packages. For example, S3 

178 URLs require the ``s3fs`` library. See 

179 :ref:`install.optional_dependencies` for a full list. 

180 engine : str, default 'pyarrow' 

181 ORC library to use. 

182 index : bool, optional 

183 If ``True``, include the dataframe's index(es) in the file output. If 

184 ``False``, they will not be written to the file. 

185 If ``None``, similar to ``infer`` the dataframe's index(es) 

186 will be saved. However, instead of being saved as values, 

187 the RangeIndex will be stored as a range in the metadata so it 

188 doesn't require much space and is faster. Other indexes will 

189 be included as columns in the file output. 

190 engine_kwargs : dict[str, Any] or None, default None 

191 Additional keyword arguments passed to :func:`pyarrow.orc.write_table`. 

192 

193 Returns 

194 ------- 

195 bytes if no path argument is provided else None 

196 

197 Raises 

198 ------ 

199 NotImplementedError 

200 Dtype of one or more columns is category, unsigned integers, interval, 

201 period or sparse. 

202 ValueError 

203 engine is not pyarrow. 

204 

205 Notes 

206 ----- 

207 * Before using this function you should read the 

208 :ref:`user guide about ORC <io.orc>` and 

209 :ref:`install optional dependencies <install.warn_orc>`. 

210 * This function requires `pyarrow <https://arrow.apache.org/docs/python/>`_ 

211 library. 

212 * For supported dtypes please refer to `supported ORC features in Arrow 

213 <https://arrow.apache.org/docs/cpp/orc.html#data-types>`__. 

214 * Currently timezones in datetime columns are not preserved when a 

215 dataframe is converted into ORC files. 

216 """ 

217 if index is None: 

218 index = df.index.names[0] is not None 

219 if engine_kwargs is None: 

220 engine_kwargs = {} 

221 

222 # validate index 

223 # -------------- 

224 

225 # validate that we have only a default index 

226 # raise on anything else as we don't serialize the index 

227 

228 if not df.index.equals(default_index(len(df))): 

229 raise ValueError( 

230 "orc does not support serializing a non-default index for the index; " 

231 "you can .reset_index() to make the index into column(s)" 

232 ) 

233 

234 if df.index.name is not None: 

235 raise ValueError("orc does not serialize index meta-data on a default index") 

236 

237 if engine != "pyarrow": 

238 raise ValueError("engine must be 'pyarrow'") 

239 pa = import_optional_dependency("pyarrow") 

240 orc = import_optional_dependency("pyarrow.orc") 

241 

242 was_none = path is None 

243 if was_none: 

244 path = io.BytesIO() 

245 assert path is not None # For mypy 

246 with get_handle(path, "wb", is_text=False) as handles: 

247 try: 

248 with suppress_pyarrow_values_warning(): 

249 table = pa.Table.from_pandas(df, preserve_index=index) 

250 orc.write_table(table, handles.handle, **engine_kwargs) 

251 except (TypeError, pa.ArrowNotImplementedError) as e: 

252 raise NotImplementedError( 

253 "The dtype of one or more columns is not supported yet." 

254 ) from e 

255 

256 if was_none: 

257 assert isinstance(path, io.BytesIO) # For mypy 

258 return path.getvalue() 

259 return None