Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/io/parquet.py: 19%
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
1"""parquet compat"""
3from __future__ import annotations
5import io
6import json
7import os
8from typing import (
9 TYPE_CHECKING,
10 Any,
11 Literal,
12)
13from warnings import (
14 catch_warnings,
15 filterwarnings,
16)
18from pandas._libs import lib
19from pandas.compat._optional import import_optional_dependency
20from pandas.errors import (
21 AbstractMethodError,
22 Pandas4Warning,
23)
24from pandas.util._decorators import set_module
25from pandas.util._validators import check_dtype_backend
27from pandas import (
28 DataFrame,
29 get_option,
30)
32from pandas.io._util import arrow_table_to_pandas
33from pandas.io.common import (
34 IOHandles,
35 get_handle,
36 is_fsspec_url,
37 is_url,
38 stringify_path,
39)
41if TYPE_CHECKING:
42 from pandas._typing import (
43 DtypeBackend,
44 FilePath,
45 ParquetCompressionOptions,
46 ReadBuffer,
47 StorageOptions,
48 WriteBuffer,
49 )
52def get_engine(engine: str) -> BaseImpl:
53 """return our implementation"""
54 if engine == "auto":
55 engine = get_option("io.parquet.engine")
57 if engine == "auto":
58 # try engines in this order
59 engine_classes = [PyArrowImpl, FastParquetImpl]
61 error_msgs = ""
62 for engine_class in engine_classes:
63 try:
64 return engine_class()
65 except ImportError as err:
66 error_msgs += "\n - " + str(err)
68 raise ImportError(
69 "Unable to find a usable engine; "
70 "tried using: 'pyarrow', 'fastparquet'.\n"
71 "A suitable version of "
72 "pyarrow or fastparquet is required for parquet "
73 "support.\n"
74 "Trying to import the above resulted in these errors:"
75 f"{error_msgs}"
76 )
78 if engine == "pyarrow":
79 return PyArrowImpl()
80 elif engine == "fastparquet":
81 return FastParquetImpl()
83 raise ValueError("engine must be one of 'pyarrow', 'fastparquet'")
86def _get_path_or_handle(
87 path: FilePath | ReadBuffer[bytes] | WriteBuffer[bytes],
88 fs: Any,
89 storage_options: StorageOptions | None = None,
90 mode: str = "rb",
91 is_dir: bool = False,
92) -> tuple[
93 FilePath | ReadBuffer[bytes] | WriteBuffer[bytes], IOHandles[bytes] | None, Any
94]:
95 """File handling for PyArrow."""
96 path_or_handle = stringify_path(path)
97 if fs is not None:
98 pa_fs = import_optional_dependency("pyarrow.fs", errors="ignore")
99 fsspec = import_optional_dependency("fsspec", errors="ignore")
100 if pa_fs is not None and isinstance(fs, pa_fs.FileSystem):
101 if storage_options:
102 raise NotImplementedError(
103 "storage_options not supported with a pyarrow FileSystem."
104 )
105 elif fsspec is not None and isinstance(fs, fsspec.spec.AbstractFileSystem):
106 pass
107 else:
108 raise ValueError(
109 f"filesystem must be a pyarrow or fsspec FileSystem, "
110 f"not a {type(fs).__name__}"
111 )
112 if is_fsspec_url(path_or_handle) and fs is None:
113 if storage_options is None:
114 pa = import_optional_dependency("pyarrow")
115 pa_fs = import_optional_dependency("pyarrow.fs")
117 try:
118 fs, path_or_handle = pa_fs.FileSystem.from_uri(path)
119 except (TypeError, pa.ArrowInvalid):
120 pass
121 if fs is None:
122 fsspec = import_optional_dependency("fsspec")
123 fs, path_or_handle = fsspec.core.url_to_fs(
124 path_or_handle, **(storage_options or {})
125 )
126 elif storage_options and (not is_url(path_or_handle) or mode != "rb"):
127 # can't write to a remote url
128 # without making use of fsspec at the moment
129 raise ValueError("storage_options passed with buffer, or non-supported URL")
131 handles = None
132 if (
133 not fs
134 and not is_dir
135 and isinstance(path_or_handle, str)
136 and not os.path.isdir(path_or_handle)
137 ):
138 # use get_handle only when we are very certain that it is not a directory
139 # fsspec resources can also point to directories
140 # this branch is used for example when reading from non-fsspec URLs
141 handles = get_handle(
142 path_or_handle, mode, is_text=False, storage_options=storage_options
143 )
144 fs = None
145 path_or_handle = handles.handle
146 return path_or_handle, handles, fs
149class BaseImpl:
150 @staticmethod
151 def validate_dataframe(df: DataFrame) -> None:
152 if not isinstance(df, DataFrame):
153 raise ValueError("to_parquet only supports IO with DataFrames")
155 def write(self, df: DataFrame, path, compression, **kwargs) -> None:
156 raise AbstractMethodError(self)
158 def read(self, path, columns=None, **kwargs) -> DataFrame:
159 raise AbstractMethodError(self)
162class PyArrowImpl(BaseImpl):
163 def __init__(self) -> None:
164 import_optional_dependency(
165 "pyarrow", extra="pyarrow is required for parquet support."
166 )
167 import pyarrow.parquet
169 # import utils to register the pyarrow extension types
170 import pandas.core.arrays.arrow.extension_types # pyright: ignore[reportUnusedImport] # noqa: F401
172 self.api = pyarrow
174 def write(
175 self,
176 df: DataFrame,
177 path: FilePath | WriteBuffer[bytes],
178 compression: ParquetCompressionOptions = "snappy",
179 index: bool | None = None,
180 storage_options: StorageOptions | None = None,
181 partition_cols: list[str] | None = None,
182 filesystem=None,
183 **kwargs,
184 ) -> None:
185 self.validate_dataframe(df)
187 from_pandas_kwargs: dict[str, Any] = {"schema": kwargs.pop("schema", None)}
188 if index is not None:
189 from_pandas_kwargs["preserve_index"] = index
191 table = self.api.Table.from_pandas(df, **from_pandas_kwargs)
193 if df.attrs:
194 df_metadata = {"PANDAS_ATTRS": json.dumps(df.attrs)}
195 existing_metadata = table.schema.metadata
196 merged_metadata = {**existing_metadata, **df_metadata}
197 table = table.replace_schema_metadata(merged_metadata)
199 path_or_handle, handles, filesystem = _get_path_or_handle(
200 path,
201 filesystem,
202 storage_options=storage_options,
203 mode="wb",
204 is_dir=partition_cols is not None,
205 )
206 if (
207 isinstance(path_or_handle, io.BufferedWriter)
208 and hasattr(path_or_handle, "name")
209 and isinstance(path_or_handle.name, (str, bytes))
210 ):
211 if isinstance(path_or_handle.name, bytes):
212 path_or_handle = path_or_handle.name.decode()
213 else:
214 path_or_handle = path_or_handle.name
216 try:
217 if partition_cols is not None:
218 # writes to multiple files under the given path
219 self.api.parquet.write_to_dataset(
220 table,
221 path_or_handle,
222 compression=compression,
223 partition_cols=partition_cols,
224 filesystem=filesystem,
225 **kwargs,
226 )
227 else:
228 # write to single output file
229 self.api.parquet.write_table(
230 table,
231 path_or_handle,
232 compression=compression,
233 filesystem=filesystem,
234 **kwargs,
235 )
236 finally:
237 if handles is not None:
238 handles.close()
240 def read(
241 self,
242 path,
243 columns=None,
244 filters=None,
245 dtype_backend: DtypeBackend | lib.NoDefault = lib.no_default,
246 storage_options: StorageOptions | None = None,
247 filesystem=None,
248 to_pandas_kwargs: dict[str, Any] | None = None,
249 **kwargs,
250 ) -> DataFrame:
251 kwargs["use_pandas_metadata"] = True
253 path_or_handle, handles, filesystem = _get_path_or_handle(
254 path,
255 filesystem,
256 storage_options=storage_options,
257 mode="rb",
258 )
259 try:
260 pa_table = self.api.parquet.read_table(
261 path_or_handle,
262 columns=columns,
263 filesystem=filesystem,
264 filters=filters,
265 **kwargs,
266 )
267 with catch_warnings():
268 filterwarnings(
269 "ignore",
270 "make_block is deprecated",
271 Pandas4Warning,
272 )
273 result = arrow_table_to_pandas(
274 pa_table,
275 dtype_backend=dtype_backend,
276 to_pandas_kwargs=to_pandas_kwargs,
277 )
279 if pa_table.schema.metadata:
280 if b"PANDAS_ATTRS" in pa_table.schema.metadata:
281 df_metadata = pa_table.schema.metadata[b"PANDAS_ATTRS"]
282 result.attrs = json.loads(df_metadata)
283 return result
284 finally:
285 if handles is not None:
286 handles.close()
289class FastParquetImpl(BaseImpl):
290 def __init__(self) -> None:
291 # since pandas is a dependency of fastparquet
292 # we need to import on first use
293 fastparquet = import_optional_dependency(
294 "fastparquet", extra="fastparquet is required for parquet support."
295 )
296 self.api = fastparquet
298 def write(
299 self,
300 df: DataFrame,
301 path,
302 compression: Literal["snappy", "gzip", "brotli"] | None = "snappy",
303 index=None,
304 partition_cols=None,
305 storage_options: StorageOptions | None = None,
306 filesystem=None,
307 **kwargs,
308 ) -> None:
309 self.validate_dataframe(df)
311 if "partition_on" in kwargs and partition_cols is not None:
312 raise ValueError(
313 "Cannot use both partition_on and "
314 "partition_cols. Use partition_cols for partitioning data"
315 )
316 if "partition_on" in kwargs:
317 partition_cols = kwargs.pop("partition_on")
319 if partition_cols is not None:
320 kwargs["file_scheme"] = "hive"
322 if filesystem is not None:
323 raise NotImplementedError(
324 "filesystem is not implemented for the fastparquet engine."
325 )
327 # cannot use get_handle as write() does not accept file buffers
328 path = stringify_path(path)
329 if is_fsspec_url(path):
330 fsspec = import_optional_dependency("fsspec")
332 # if filesystem is provided by fsspec, file must be opened in 'wb' mode.
333 kwargs["open_with"] = lambda path, _: fsspec.open(
334 path, "wb", **(storage_options or {})
335 ).open()
336 elif storage_options:
337 raise ValueError(
338 "storage_options passed with file object or non-fsspec file path"
339 )
341 with catch_warnings(record=True):
342 self.api.write(
343 path,
344 df,
345 compression=compression,
346 write_index=index,
347 partition_on=partition_cols,
348 **kwargs,
349 )
351 def read(
352 self,
353 path,
354 columns=None,
355 filters=None,
356 storage_options: StorageOptions | None = None,
357 filesystem=None,
358 to_pandas_kwargs: dict | None = None,
359 **kwargs,
360 ) -> DataFrame:
361 parquet_kwargs: dict[str, Any] = {}
362 dtype_backend = kwargs.pop("dtype_backend", lib.no_default)
363 # We are disabling nullable dtypes for fastparquet pending discussion
364 parquet_kwargs["pandas_nulls"] = False
365 if dtype_backend is not lib.no_default:
366 raise ValueError(
367 "The 'dtype_backend' argument is not supported for the "
368 "fastparquet engine"
369 )
370 if filesystem is not None:
371 raise NotImplementedError(
372 "filesystem is not implemented for the fastparquet engine."
373 )
374 if to_pandas_kwargs is not None:
375 raise NotImplementedError(
376 "to_pandas_kwargs is not implemented for the fastparquet engine."
377 )
378 path = stringify_path(path)
379 handles = None
380 if is_fsspec_url(path):
381 fsspec = import_optional_dependency("fsspec")
383 parquet_kwargs["fs"] = fsspec.open(path, "rb", **(storage_options or {})).fs
384 elif isinstance(path, str) and not os.path.isdir(path):
385 # use get_handle only when we are very certain that it is not a directory
386 # fsspec resources can also point to directories
387 # this branch is used for example when reading from non-fsspec URLs
388 handles = get_handle(
389 path, "rb", is_text=False, storage_options=storage_options
390 )
391 path = handles.handle
393 try:
394 parquet_file = self.api.ParquetFile(path, **parquet_kwargs)
395 with catch_warnings():
396 filterwarnings(
397 "ignore",
398 "make_block is deprecated",
399 Pandas4Warning,
400 )
401 return parquet_file.to_pandas(
402 columns=columns, filters=filters, **kwargs
403 )
404 finally:
405 if handles is not None:
406 handles.close()
409def to_parquet(
410 df: DataFrame,
411 path: FilePath | WriteBuffer[bytes] | None = None,
412 engine: str = "auto",
413 compression: ParquetCompressionOptions = "snappy",
414 index: bool | None = None,
415 storage_options: StorageOptions | None = None,
416 partition_cols: list[str] | None = None,
417 filesystem: Any = None,
418 **kwargs,
419) -> bytes | None:
420 """
421 Write a DataFrame to the parquet format.
423 Parameters
424 ----------
425 df : DataFrame
426 path : str, path object, file-like object, or None, default None
427 String, path object (implementing ``os.PathLike[str]``), or file-like
428 object implementing a binary ``write()`` function. If None, the result
429 is returned as bytes. If a string, it will be used as Root Directory
430 path when writing a partitioned dataset. The engine fastparquet does
431 not accept file-like objects.
432 engine : {'auto', 'pyarrow', 'fastparquet'}, default 'auto'
433 Parquet library to use. If 'auto', then the option
434 ``io.parquet.engine`` is used. The default ``io.parquet.engine``
435 behavior is to try 'pyarrow', falling back to 'fastparquet' if
436 'pyarrow' is unavailable.
438 When using the ``'pyarrow'`` engine and no storage options are provided
439 and a filesystem is implemented by both ``pyarrow.fs`` and ``fsspec``
440 (e.g. "s3://"), then the ``pyarrow.fs`` filesystem is attempted first.
441 Use the filesystem keyword with an instantiated fsspec filesystem
442 if you wish to use its implementation.
443 compression : {'snappy', 'gzip', 'brotli', 'lz4', 'zstd', None},
444 default 'snappy'. Name of the compression to use. Use ``None``
445 for no compression.
446 index : bool, default None
447 If ``True``, include the dataframe's index(es) in the file output. If
448 ``False``, they will not be written to the file.
449 If ``None``, similar to ``True`` the dataframe's index(es)
450 will be saved. However, instead of being saved as values,
451 the RangeIndex will be stored as a range in the metadata so it
452 doesn't require much space and is faster. Other indexes will
453 be included as columns in the file output.
454 partition_cols : str or list, optional, default None
455 Column names by which to partition the dataset.
456 Columns are partitioned in the order they are given.
457 Must be None if path is not a string.
458 storage_options : dict, optional
459 Extra options that make sense for a particular storage connection, e.g.
460 host, port, username, password, etc. For HTTP(S) URLs the key-value
461 pairs are forwarded to ``urllib.request.Request`` as header options.
462 For other URLs (e.g. starting with "s3://", and "gcs://") the
463 key-value pairs are forwarded to ``fsspec.open``. Please see ``fsspec``
464 and ``urllib`` for more details, and for more examples on storage
465 options refer `here <https://pandas.pydata.org/docs/user_guide/io.html?
466 highlight=storage_options#reading-writing-remote-files>`_.
467 filesystem : fsspec or pyarrow filesystem, default None
468 Filesystem object to use when reading the parquet file. Only implemented
469 for ``engine="pyarrow"``.
471 .. versionadded:: 2.1.0
473 **kwargs
474 Additional keyword arguments passed to the engine:
476 * For ``engine="pyarrow"``: passed to :func:`pyarrow.parquet.write_table`
477 or :func:`pyarrow.parquet.write_to_dataset` (when using partition_cols)
478 * For ``engine="fastparquet"``: passed to :func:`fastparquet.write`
480 Returns
481 -------
482 bytes if no path argument is provided else None
483 """
484 if isinstance(partition_cols, str):
485 partition_cols = [partition_cols]
486 impl = get_engine(engine)
488 path_or_buf: FilePath | WriteBuffer[bytes] = io.BytesIO() if path is None else path
490 impl.write(
491 df,
492 path_or_buf,
493 compression=compression,
494 index=index,
495 partition_cols=partition_cols,
496 storage_options=storage_options,
497 filesystem=filesystem,
498 **kwargs,
499 )
501 if path is None:
502 assert isinstance(path_or_buf, io.BytesIO)
503 return path_or_buf.getvalue()
504 else:
505 return None
508@set_module("pandas")
509def read_parquet(
510 path: FilePath | ReadBuffer[bytes],
511 engine: str = "auto",
512 columns: list[str] | None = None,
513 storage_options: StorageOptions | None = None,
514 dtype_backend: DtypeBackend | lib.NoDefault = lib.no_default,
515 filesystem: Any = None,
516 filters: list[tuple] | list[list[tuple]] | None = None,
517 to_pandas_kwargs: dict | None = None,
518 **kwargs,
519) -> DataFrame:
520 """
521 Load a parquet object from the file path, returning a DataFrame.
523 The function automatically handles reading the data from a parquet file
524 and creates a DataFrame with the appropriate structure.
526 Parameters
527 ----------
528 path : str, path object or file-like object
529 String, path object (implementing ``os.PathLike[str]``), or file-like
530 object implementing a binary ``read()`` function.
531 The string could be a URL. Valid URL schemes include http, ftp, s3,
532 gs, and file. For file URLs, a host is expected. A local file could be:
533 ``file://localhost/path/to/table.parquet``.
534 A file URL can also be a path to a directory that contains multiple
535 partitioned parquet files. Both pyarrow and fastparquet support
536 paths to directories as well as file URLs. A directory path could be:
537 ``file://localhost/path/to/tables`` or ``s3://bucket/partition_dir``.
538 engine : {'auto', 'pyarrow', 'fastparquet'}, default 'auto'
539 Parquet library to use. If 'auto', then the option
540 ``io.parquet.engine`` is used. The default ``io.parquet.engine``
541 behavior is to try 'pyarrow', falling back to 'fastparquet' if
542 'pyarrow' is unavailable.
544 When using the ``'pyarrow'`` engine and no storage options are provided
545 and a filesystem is implemented by both ``pyarrow.fs`` and ``fsspec``
546 (e.g. "s3://"), then the ``pyarrow.fs`` filesystem is attempted first.
547 Use the filesystem keyword with an instantiated fsspec filesystem
548 if you wish to use its implementation.
549 columns : list, default=None
550 If not None, only these columns will be read from the file.
551 storage_options : dict, optional
552 Extra options that make sense for a particular storage connection, e.g.
553 host, port, username, password, etc. For HTTP(S) URLs the key-value
554 pairs are forwarded to ``urllib.request.Request`` as header options.
555 For other URLs (e.g. starting with "s3://", and "gcs://") the
556 key-value pairs are forwarded to ``fsspec.open``. Please see ``fsspec``
557 and ``urllib`` for more details, and for more examples on storage
558 options refer `here <https://pandas.pydata.org/docs/user_guide/io.html?
559 highlight=storage_options#reading-writing-remote-files>`_.
560 dtype_backend : {'numpy_nullable', 'pyarrow'}
561 Back-end data type applied to the resultant :class:`DataFrame`
562 (still experimental). If not specified, the default behavior
563 is to not use nullable data types. If specified, the behavior
564 is as follows:
566 * ``"numpy_nullable"``: returns nullable-dtype-backed :class:`DataFrame`
567 * ``"pyarrow"``: returns pyarrow-backed nullable
568 :class:`ArrowDtype` :class:`DataFrame`
570 .. versionadded:: 2.0
572 filesystem : fsspec or pyarrow filesystem, default None
573 Filesystem object to use when reading the parquet file. Only implemented
574 for ``engine="pyarrow"``.
576 .. versionadded:: 2.1.0
578 filters : List[Tuple] or List[List[Tuple]], default None
579 To filter out data.
580 Filter syntax: [[(column, op, val), ...],...]
581 where op is [==, =, >, >=, <, <=, !=, in, not in]
582 The innermost tuples are transposed into a set of filters applied
583 through an `AND` operation.
584 The outer list combines these sets of filters through an `OR`
585 operation.
586 A single list of tuples can also be used, meaning that no `OR`
587 operation between set of filters is to be conducted.
589 Using this argument will NOT result in row-wise filtering of the final
590 partitions unless ``engine="pyarrow"`` is also specified. For
591 other engines, filtering is only performed at the partition level, that is,
592 to prevent the loading of some row-groups and/or files.
594 .. versionadded:: 2.1.0
596 to_pandas_kwargs : dict | None, default None
597 Keyword arguments to pass through to :func:`pyarrow.Table.to_pandas`
598 when ``engine="pyarrow"``.
600 .. versionadded:: 3.0.0
602 **kwargs
603 Additional keyword arguments passed to the engine:
605 * For ``engine="pyarrow"``: passed to :func:`pyarrow.parquet.read_table`
606 * For ``engine="fastparquet"``: passed to
607 :meth:`fastparquet.ParquetFile.to_pandas`
609 Returns
610 -------
611 DataFrame
612 DataFrame based on parquet file.
614 See Also
615 --------
616 DataFrame.to_parquet : Create a parquet object that serializes a DataFrame.
618 Examples
619 --------
620 >>> original_df = pd.DataFrame({"foo": range(5), "bar": range(5, 10)})
621 >>> original_df
622 foo bar
623 0 0 5
624 1 1 6
625 2 2 7
626 3 3 8
627 4 4 9
628 >>> df_parquet_bytes = original_df.to_parquet()
629 >>> from io import BytesIO
630 >>> restored_df = pd.read_parquet(BytesIO(df_parquet_bytes))
631 >>> restored_df
632 foo bar
633 0 0 5
634 1 1 6
635 2 2 7
636 3 3 8
637 4 4 9
638 >>> restored_df.equals(original_df)
639 True
640 >>> restored_bar = pd.read_parquet(BytesIO(df_parquet_bytes), columns=["bar"])
641 >>> restored_bar
642 bar
643 0 5
644 1 6
645 2 7
646 3 8
647 4 9
648 >>> restored_bar.equals(original_df[["bar"]])
649 True
651 The function uses `kwargs` that are passed directly to the engine.
652 In the following example, we use the `filters` argument of the pyarrow
653 engine to filter the rows of the DataFrame.
655 Since `pyarrow` is the default engine, we can omit the `engine` argument.
656 Note that the `filters` argument is implemented by the `pyarrow` engine,
657 which can benefit from multithreading and also potentially be more
658 economical in terms of memory.
660 >>> sel = [("foo", ">", 2)]
661 >>> restored_part = pd.read_parquet(BytesIO(df_parquet_bytes), filters=sel)
662 >>> restored_part
663 foo bar
664 0 3 8
665 1 4 9
666 """
668 impl = get_engine(engine)
669 check_dtype_backend(dtype_backend)
671 return impl.read(
672 path,
673 columns=columns,
674 filters=filters,
675 storage_options=storage_options,
676 dtype_backend=dtype_backend,
677 filesystem=filesystem,
678 to_pandas_kwargs=to_pandas_kwargs,
679 **kwargs,
680 )