B x\%@sdZddlmZddlmZddlmZddlmZddl m Z m Z ddl m Z mZdd ZGd d d eZGd d d eZGdddeZdddZdddZdS)z parquet compat ) LooseVersion)catch_warnings) string_types)AbstractMethodError) DataFrame get_option)get_filepath_or_buffer is_s3_urlcCs|dkrtd}|dkrXytStk r2YnXytStk rNYnXtd|dkrhtd|dkrvtS|dkrtSdS) z return our implementation autozio.parquet.enginez}Unable to find a usable engine; tried using: 'pyarrow', 'fastparquet'. pyarrow or fastparquet is required for parquet support)pyarrow fastparquetz.engine must be one of 'pyarrow', 'fastparquet'r r N)r PyArrowImpl ImportErrorFastParquetImpl ValueError)enginer0lib/python3.7/site-packages/pandas/io/parquet.py get_engines$rc@s.eZdZdZeddZddZdddZdS) BaseImplNcCsLt|tstd|jjdkr&tdtdd|jjD}|sHtddS)Nz+to_parquet only supports IO with DataFrames>unicodestringz%parquet must have string column namescss |]}|dk rt|tVqdS)N) isinstancer).0namerrr >sz.BaseImpl.validate_dataframe..z!Index level names must be strings)rrrcolumnsZ inferred_typeallindexnames)dfZ valid_namesrrrvalidate_dataframe2s  zBaseImpl.validate_dataframecKs t|dS)N)r)selfr path compressionkwargsrrrwriteEszBaseImpl.writecKs t|dS)N)r)r"r#rr%rrrreadHsz BaseImpl.read)N)__name__ __module__ __qualname__api staticmethodr!r&r'rrrrr.s rc@s(eZdZddZd ddZd dd ZdS) r cCsRyddl}ddl}Wntk r0tdYnXt|jdkrHtd||_dS)Nrzpyarrow is required for parquet support you can install via conda conda install pyarrow -c conda-forge or via pip pip install -U pyarrow z0.9.0zpyarrow >= 0.9.0 is required for parquet support you can install via conda conda install pyarrow -c conda-forge or via pip pip install -U pyarrow )r Zpyarrow.parquetrr __version__r+)r"r rrr__init__Ns  zPyArrowImpl.__init__snappymsNc Ks||t|dd\}}}}|dkr,i} nd|i} |jjj|f| } |dk rp|jjj| |f|||d|n|jjj| |f||d|dS)Nwb)modeZpreserve_index)r$coerce_timestampspartition_cols)r$r3)r!rr+ZTableZ from_pandasparquetZwrite_to_datasetZ write_table) r"r r#r$r3rr4r%_Zfrom_pandas_kwargstablerrrr&gs zPyArrowImpl.writecKsXt|\}}}}d|d<|jjj|fd|i|}|rTy |Wn YnX|S)NTZuse_pandas_metadatar)rr+r5Z read_table to_pandasclose)r"r#rr%r6 should_closeresultrrrr'|s  zPyArrowImpl.read)r/r0NN)N)r(r)r*r.r&r'rrrrr Ls r c@s(eZdZddZd ddZd ddZdS) rcCsJy ddl}Wntk r(tdYnXt|jdkr@td||_dS)Nrzfastparquet is required for parquet support you can install via conda conda install fastparquet -c conda-forge or via pip pip install -U fastparquetz0.2.1zfastparquet >= 0.2.1 is required for parquet support you can install via conda conda install fastparquet -c conda-forge or via pip pip install -U fastparquet)r rrr-r+)r"r rrrr.s  zFastParquetImpl.__init__r/Nc Ks||d|kr$|dk r$tdnd|kr6|d}|dk rFd|d<t|rpt|dd\}}}}dd|d <nt|\}}}}td d $|jj||f|||d |WdQRXdS) N partition_onzYCannot use both partition_on and partition_cols. Use partition_cols for partitioning dataZhiveZ file_schemer1)r2cSs|S)Nr)r#r6rrrsz'FastParquetImpl.write.. open_withT)record)r$Z write_indexr<)r!rpopr rrr+r&)r"r r#r$rr4r%r6rrrr&s    zFastParquetImpl.writecKsnt|r>t|\}}}}z|jj||jjd}Wd|Xnt|\}}}}|j|}|jfd|i|S)N)r>r)r rr+Z ParquetFiles3openr9r8)r"r#rr%rAr6r:Z parquet_filerrrr's  zFastParquetImpl.read)r/NN)N)r(r)r*r.r&r'rrrrrs rr r/NcKs$t|}|j||f|||d|S)a Write a DataFrame to the parquet format. Parameters ---------- path : str File path or Root Directory path. Will be used as Root Directory path while writing a partitioned dataset. .. versionchanged:: 0.24.0 engine : {'auto', 'pyarrow', 'fastparquet'}, default 'auto' Parquet library to use. If 'auto', then the option ``io.parquet.engine`` is used. The default ``io.parquet.engine`` behavior is to try 'pyarrow', falling back to 'fastparquet' if 'pyarrow' is unavailable. compression : {'snappy', 'gzip', 'brotli', None}, default 'snappy' Name of the compression to use. Use ``None`` for no compression. index : bool, default None If ``True``, include the dataframe's index(es) in the file output. If ``False``, they will not be written to the file. If ``None``, the engine's default behavior will be used. .. versionadded 0.24.0 partition_cols : list, optional, default None Column names by which to partition the dataset Columns are partitioned in the order they are given .. versionadded:: 0.24.0 kwargs Additional keyword arguments passed to the engine )r$rr4)rr&)r r#rr$rr4r%implrrr to_parquets$rDcKst|}|j|fd|i|S)a Load a parquet object from the file path, returning a DataFrame. .. versionadded 0.21.0 Parameters ---------- path : string File path columns : list, default=None If not None, only these columns will be read from the file. .. versionadded 0.21.1 engine : {'auto', 'pyarrow', 'fastparquet'}, default 'auto' Parquet library to use. If 'auto', then the option ``io.parquet.engine`` is used. The default ``io.parquet.engine`` behavior is to try 'pyarrow', falling back to 'fastparquet' if 'pyarrow' is unavailable. kwargs are passed to the engine Returns ------- DataFrame r)rr')r#rrr%rCrrr read_parquetsrE)r r/NN)r N)__doc__Zdistutils.versionrwarningsrZ pandas.compatrZ pandas.errorsrZpandasrrZpandas.io.commonrr robjectrr rrDrErrrrs     ?K (