U 饡c@spddlZddlmZddlmZddlZddlmZddlm Z ddl Z ddl Z ddl Z ddl Z ddlZddlZddlmZmZmZiaddZd d Zejd ejd ejd ejdejdejdejdej dej!dej"dej#dddej$dej%diZ&ddZ'ddZ(ddZ)d d!Z*d"d#Z+d$d%Z,d&d'Z-d(d)Z.d*d+Z/d,d-Z0d.d/Z1d0d1Z2d2d3Z3d4d5Z4dmd6d7Z5dnd:d;Z6dd?Z8d@dAZ9dodBdCZ:dDdEZ;dpdGdHZdqdOdPZ?dQdRZ@dSdTZAdUdVZBdWdXZCdYdZZDd[d\ZEdd]d]ej$ej%ejFejej#ejGd^ ZHd_d`ZIdadbZJdcddZKdedfZLdgdhZMdidjZNdkdlZOdS)rN)Sequence)futures)deepcopy) zip_longest) _pandas_apibuiltin_pickle frombytesc*Cststtjjdtjjdtjjdtjjdtjjdtjj dtjj dtjj dtjj d tjj d tjjd tjjd tjjd tjjdtjjdtjjdtjjdtjjdtjjdtjjditS)Nemptyboolint8int16int32int64uint8uint16uint32uint64Zfloat16float32float64datetimebytesunicode)_logical_type_mapupdatepalibZType_NAZ Type_BOOLZ Type_INT8Z Type_INT16Z Type_INT32Z Type_INT64Z Type_UINT8Z Type_UINT16Z Type_UINT32Z Type_UINT64ZType_HALF_FLOATZ Type_FLOATZ Type_DOUBLEZ Type_DATE32Z Type_DATE64Z Type_TIME32Z Type_TIME64Z Type_BINARYZType_FIXED_SIZE_BINARYZ Type_STRINGrr9/tmp/pip-unpacked-wheel-qoi2rb4q/pyarrow/pandas_compat.pyget_logical_type_map)sXrcCst}z ||jWStk rt|tjjr6YdSt|tjjrXdt |j YSt|tjj r||j dk rtdndYSt|tjj rYdSYdSXdS)NZ categoricalzlist[{}] datetimetzdatetimedecimalobject)ridKeyError isinstancerrZDictionaryTypeZListTypeformatget_logical_typeZ value_type TimestampTypetzZDecimal128Type) arrow_typeZlogical_type_maprrrr(Fs r(r r r r rrrrrrrz datetime64[D]rstringrcCsjzt|jjWStk rdt|jdr0YdSt|jdkrDYdSt|}|dkr\YdS|YSXdS)Nr*r datetime64[ns]r,r)_numpy_logical_type_mapdtypetyper%hasattrstrrZ infer_dtype)Zpandas_collectionresultrrrget_logical_type_from_numpyis  r4cCs|j}t|dkrJt|d|}|dk s*tt|j|jd}t|jj}n.t|drldt j |j i}d}n d}t|}||fS)Ncategorycat)Znum_categoriesorderedr*timezoner-) r/r2getattrAssertionErrorlen categoriesr7codesr1rrtzinfo_to_stringr*)columnr/ZcatsmetadataZphysical_dtyperrrget_extension_dtype_infoxs    rAcCst|}t|\}}|dkr.|j|jd}d}|dk rVt|tsVtd|t|j |dksxt|tsxt tt|||dkrdn||||dS)aConstruct the metadata for a given column Parameters ---------- column : pandas.Series or pandas.Index name : str arrow_type : pyarrow.DataType field_name : str Equivalent to `name` when `column` is a `Series`, otherwise if `column` is a pandas Index then `field_name` will not be the same as `name`. This is the name of the field in the arrow Table's schema. Returns ------- dict r") precisionscaler#Nz6Column name must be a string. Got column {} of type {}Nonename field_name pandas_type numpy_typer@) r(rArBrCr&r2 TypeErrorr'r0__name__r:)r?rFr+rGZ logical_type string_dtypeextra_metadatarrrget_column_metadatas.  rNc Cstdd|D}t|}|d||} |||d} g} t||| D]$\} } }t| | || d}| |qJg}|dk rdg}t|| |D]X\}}}t|trq|jdk rt|jts||jt|t|j||d}||qt|dkrt j d|dt d d g}t |j d |j g}t |j d |j jg}t||D]\}}t||}||qBn g}}}d t||| |dtjdtjddiS)aReturns a dictionary containing enough metadata to reconstruct a pandas DataFrame as an Arrow Table, including index columns. Parameters ---------- columns_to_convert : list[pd.Series] df : pandas.DataFrame index_levels : List[pd.Index] index_descriptors : List[Dict] preserve_index : bool types : List[pyarrow.DataType] Returns ------- dict cSsg|]}t|ts|qSr)r&dict).0descrrrr s z&construct_metadata..N)rFr+rGFrz&The DataFrame has non-str index name `z@` which will be converted to string and not roundtrip correctly. stacklevellevelsnamesspandaspyarrow)Zlibraryversion) index_columnscolumn_indexescolumnsZcreatorZpandas_versionutf8)r;ziprNappendr&rOrFr2_column_name_to_stringswarningswarn UserWarningr9r\_get_simple_index_descriptorjsondumpsr __version__rrYencode)columns_to_convertdf column_names index_levelsindex_descriptorspreserve_indextypesZnum_serialized_index_levelsZntypesZdf_typesZ index_typesZcolumn_metadatacolZsanitized_namer+r@Zindex_column_metadataZnon_str_index_nameslevelZ descriptorr[rVrWrFrrrconstruct_metadatasr         rrcCsTt|\}}t|}d|kr,tjdtdd|dkrD|rtttt|St|trRtdn |dkr^dSt|S)a!Convert a column name (or level) to either a string or a recursive collection of strings. Parameters ---------- name : str or tuple Returns ------- value : str or tuple Examples -------- >>> name = 'foo' >>> _column_name_to_strings(name) 'foo' >>> name = ('foo', 'bar') >>> _column_name_to_strings(name) "('foo', 'bar')" >>> import pandas as pd >>> name = (1, pd.Timestamp('2017-02-01 00:00:00')) >>> _column_name_to_strings(name) "('1', '2017-02-01 00:00:00')" r]z%Unsupported type for MultiIndex levelN) r&r2rdecodetuplemapr`rrJrFrrrr`"s      r`cCs,|jdk r|j|krt|jSd|SdS)zReturn the name of an index level or a default name if `index.name` is None or is already a column name. Parameters ---------- index : pandas.Index i : int Returns ------- name : str Nz__index_level_{:d}__)rFr`r')indexirkrrr_index_level_nameIs  r|cCs:t|||}|jjs(tdt|j|dk r.)r;r9r)rznrrrrsrcsP|dk r|dk rtdn2|dk r*|j}n"|dk rFfdd|D}nj}|S)NzJSchema and columns arguments are mutually exclusive, pass only one of themcsg|]}|jkr|qSr)r\rPcrjrrrR s z0_resolve_columns_of_interest..)rrWr\)rjrr\rrrr~s r~cCst|d||\}}}}}}}g} |D]} | j} t| rJtj| ddj} nxt| rt| tj j rl| dn | dd} tj| ddj} n8t | | j d\} } tj| | } | dkrtj| ddj} | | q$t||||||| }|| |fS)NT) from_pandasr)rvaluesris_categoricalrarrayr0Zis_extension_array_dtyper&rSeriesheadget_datetimetz_typer/rZ_ndarray_to_arrow_typer_rr)rjrnr\rrk_rmrZrirorrtype_r r@rrrdataframe_to_typessH     rTc! s:t||||\}}}} } } } |dkrZt|t|j} }| |dkrV|dkrVt}nd}fdddd}|dkrfddt| | D}ng}t|J}t| | D]8\}}||jr| ||q| | ||qW5QRXt |D] \}}t |tj r|||<qd d|D}|dkrxg}t||D].\}}|dk rT|nd }| t||q>t|}t| ||| | ||}|jrt|jnt}||||}d}t|d kr0zN| d d }|d kr| d d}| d d}| d d} tt||| }Wntk r.YnX|||fS)Ndrc s|dkrd}d}n |j}|j}ztj||dd}WnNtjtjtjfk r}z$|jd|j |j f7_|W5d}~XYnX|s|j dkrt dt ||j |S)NT)r0rsafez0Conversion failed for column {!s} with type {!s}rz>Field {} was non-nullable but pandas column had {} null values)Znullabler0rrZ ArrowInvalidZArrowNotImplementedErrorZArrowTypeErrorargsr'rFr/Z null_countrr2)rprZfield_nullablerr3e)rrrconvert_columnGs,  z+dataframe_to_arrays..convert_columncSs$t|tjo"|jjo"t|jjtjSr) r&npZndarrayflags contiguous issubclassr/r0integer)arrrrr_can_definitely_zero_copy]s  z6dataframe_to_arrays.._can_definitely_zero_copycsg|]\}}||qSrr)rPrf)rrrrRcsz'dataframe_to_arrays..cSsg|] }|jqSrr0rPxrrrrRrsrDrrrrrr)rr;r\r cpu_countr^rZThreadPoolExecutorrr_Zsubmitrr&ZFuturer3rrrrr@rrOr with_metadatarr)!rjrrnZnthreadsr\rrrkrrmrZrirZnrowsZncolsrZarraysexecutorrrr{Z maybe_futrofieldsrFrpandas_metadatar@Zn_rowsrrrrr)rrrdataframe_to_arrays2sz                rcCs^|jjtjkr||fSt|rB|dkrB|j}|j}t ||}n|dkrVt |j}||fSr) r/r0rZ datetime64r is_datetimetzr*unitr timestampZfrom_numpy_dtype)rr/rr*rrrrrs rcCs|j}g}dd|jD}|jD]}|j}i}t|jr^tj |j |d<t |dr|j}n"t |r|j |j|jd|j}|j |jj|d|jttkrd|d<tj|tjd|d <||q ||d S) NcSsg|]}|qSrr)rPZaxrrrrRsz0dataframe_to_serialized_dict..r8r) dictionaryr7) placementblockr#)protocolr)blocksaxes)_datarrrrrr/rrr>r*r1rrr<r7r=Zmgr_locsZas_arrayrr#rrfHIGHEST_PROTOCOLr_)frameZ block_managerrrrrZ block_datarrrdataframe_to_serialized_dicts:      rcCs>ddlmm}dd|dD}|||d}t|S)NrcSsg|] }t|qSr_reconstruct_block)rPrrrrrRsz0serialized_dict_to_dataframe..rr)pandas.core.internalscore internals BlockManagerrZ data_frame)data_intZreconstructed_blocksZ block_mgrrrrserialized_dict_to_dataframes rc Csddlmm}|dd}|d}d|krXtjj||d|dd}|j||d}nd |krt|d }|j|||j |d }nd |kr|jt ||d}npd |kr|d } t |d kst ||d} || } t| dstd| | } |j| |d}n|j||d}|S)a Construct a pandas Block from the `item` dictionary coming from pyarrow's serialization or returned by arrow::python::ConvertTableToPandas. This function takes care of converting dictionary types to pandas categorical, Timestamp-with-timezones to the proper pandas Block, and conversion to pandas ExtensionBlock Parameters ---------- item : dict For basic types, this is a dictionary in the form of {'block': np.ndarray of values, 'placement': pandas block placement}. Additional keys are present for other types (dictionary, timezone, object). columns : Column names of the table being constructed, used for extension types extension_columns : dict Dictionary of {column_name: pandas_dtype} that includes all columns and corresponding dtypes that will be converted to a pandas ExtensionBlock. Returns ------- pandas Block rNrrrr7)r<r7)rr8)rklassr/r#Zpy_arrayr__from_arrow__zGThis column does not support to be converted to a pandas ExtensionArray)rrrgetrZcategorical_typeZ from_codesZ make_blockmake_datetimetzZDatetimeTZBlockrloadsr;r:r1rr) itemr\extension_columnsrZ block_arrrr6rr/rrF pandas_dtypeZ pd_ext_arrrrrrs>       rcCstj|}tjd|dS)Nnsr*)rrstring_to_tzinforZdatetimetz_typerrrrrs rFcCsddlm}g}g}|jj}|sl|dk rl|d}|dg}|d} t||}t|| |\}} t|||} ntj |j } t|g|} t |t |||} t|||| } | | g}|| |S)Nr)rr\r[rZ)rrrrr_add_any_metadata_reconstruct_index_get_extension_dtypesrrrnum_rows'_check_data_column_metadata_consistency_deserialize_column_index_table_to_blocks)optionstabler<Zignore_metadata types_mapperr all_columnsr[rrmrzZext_columns_dtypesr\rrrrrtable_to_blockmanagers0      rcCsh|]}tt|qSr)r2rr/)rPtyprrr <srrZuintfloatr#c Csi}tjdkr|S|D]f}z |d}Wntk rB|d}YnX|d}|tkrt|}t|tjrt|dr|||<q|jD]B}|j} t| t j rz | }Wnt k rYqX|||j <q|r|jD]$}|j} || }|dk r|||j <q|S)a Based on the stored column pandas metadata and the extension types in the arrow schema, infer which columns should be converted to a pandas extension dtype. The 'numpy_type' field in the column metadata stores the string representation of the original pandas dtype (and, despite its name, not the 'pandas_type' field). Based on this string representation, a pandas/numpy dtype is constructed and then we can check if this dtype supports conversion from arrow. NrGrFrIr)rZextension_dtyper%_pandas_supported_numpy_typesrr&r1rr0rZBaseExtensionTypeZto_pandas_dtypeNotImplementedErrorrF) rZcolumns_metadatarZ ext_columnscol_metarFr/rrrrrrrCs:             rcCstdd|DstdS)Ncss.|]&}|ddkrd|kp$|ddk VqdS)rFNrGrrrrr sz:_check_data_column_metadata_consistency..)allr:)rrrrr|srcsdd|jD}|r6dd|Dfdd|D}n|}t|dkrLtjndd}|sftj|}n*tjjjt t ||d d|Dpdd }t|d krt ||}t |}|S) NcSs"g|]}t|trt|n|qSr)r&rrrrrrrRsz-_deserialize_column_index..cSs&i|]}|dt|d|dqSrGrF)rr`rrrr sz-_deserialize_column_index..csg|]}||qSrr)rPrFZcolumns_name_dictrrrRsrcSs|fSrr)rrrrz+_deserialize_column_index..cSsg|] }|dqSryr)rP col_indexrrrrRsrWr) rkr;ast literal_evalrrIndex MultiIndex from_tuplesrrx"_reconstruct_columns_from_metadata _flatten_single_level_multiindex) block_tablerr[Zcolumn_stringsZcolumns_valuesZto_pairr\rrrrs,    rc Cs,dd|D}g}g}|}|D]}t|trLt||||\}}} |dkrqnZ|ddkr|d} tjj|d|d|d| d }t|t|krqntd |d| || | qtj} t|d kr| j j ||d } nDt|d kr|d } t| | j s$| j | |d d} n | |j } || fS)NcSsi|]}|d|d|qSrrrrrrrsz&_reconstruct_index..rrrFrrr)rrFzUnrecognized index kind: {}rrrry)r&r2_extract_index_levelrrrr;rr'r_r from_arraysrr) rrmrfield_name_to_metadataZ index_arraysZ index_names result_tablerQr index_namerrzrrrrsN        rc Cs||d}t||}|j|}|dkr4|ddfStj}||}|j} t| drh| j j sh| } t |j tjjr|j jdk rt|| |j j} n|j| | jd} ||j|}|| |fS)NrFr)r/) _backwards_compatible_index_namerget_field_indexrrr? to_pandasrr1rZ writeablecopyr&r0rrr)r* make_tz_awarerr/Z remove_column) rr rGr  logical_namer r{rrprrrrrr s"       r cCs||krt|rdS|SdS)a1Compute the name of an index column that is compatible with older versions of :mod:`pyarrow`. Parameters ---------- raw_name : str logical_name : str Returns ------- result : str Notes ----- * Part of :func:`~pyarrow.pandas_compat.table_to_blockmanager` N)r)raw_namerrrrrsrcCsd}t||dk S)Nz^__index_level_\d+__$)rematch)rFpatternrrrrsrr-) rr!r rrr,rZfloatingr cCs@z t|WStk r:d|kr,tjYSt|YSXdS)a Get the numpy dtype that corresponds to a pandas type. Parameters ---------- pandas_type : str The result of a call to pandas.lib.infer_dtype. Returns ------- dtype : np.dtype The dtype that corresponds to `pandas_type`. rsN)_pandas_logical_type_mapr%robject_r/)rHrrr_pandas_type_to_numpy_type-s   rcCs,t|tjjr$t|dr|jS|jSdSdS)Nr=)r&rrrr1r=labels)mirrr_get_multiindex_codesCsrc stjt|ddp|g}t|p0fdd|D}ddt||idD}g}tdd}|D]\}}} t|} | tj kr| |}|d krt j |d d d } || } | d | }n|j| kr|| }|j| kr|d kr|| }||q^j|||jdS)a_Construct a pandas MultiIndex from `columns` and column index metadata in `column_indexes`. Parameters ---------- columns : List[pd.Index] The columns coming from a pyarrow.Table column_indexes : List[Dict[str, str]] The column index metadata deserialized from the JSON schema metadata in a :class:`~pyarrow.Table`. Returns ------- result : MultiIndex The index reconstructed using `column_indexes` metadata with levels of the correct type. Notes ----- * Part of :func:`~pyarrow.pandas_compat.table_to_blockmanager` rVNcsg|]}t|qSr)rr;)rPrqrrrrResz6_reconstruct_columns_from_metadata..cSs0g|](\}}||dt|j|ddfqS)rHrIN)rr2r/)rPrqrrrrrRjs ) fillvaluerhrur rr@r8utcr)rrr9rroperator methodcallerrrbytes_rxrrrZastype tz_localize tz_convertr/r_rrW) r\r[rVrZ levels_dtypesZ new_levelsencoderrqrZ numpy_dtyper/r*dtrrrrKs<        rcs4|jtj|||t}fdd|DS)Ncsg|]}t|qSrr)rPrr\rrrrRsz$_table_to_blocks..)rkrrZtable_to_blocksrkeys)rrr<rr3rr)rrs   rcsjtj}t||jrf|jdkrf|j\t|\}j}|jsBt d|j fdd|D||j ddS|S)NrzFound non-unique column indexcs g|]}|dkr|ndqS)rNr)rPZ_labelrVrrrRsz4_flatten_single_level_multiindex..r)r/rF) rrr&rZnlevelsrVrr/rrrrW)rzrrr/rr+rrs rcCsi}i}|j}|d}dd|D}t|}t|d|}t|dD]\}} | d} | s| d} ||krz|||} | dkrd} || } | dkrH| d d krH|| } t| jtjj sqH| d } | sqH| d }|rH|| jj krH| }tj d |d}tj j||d}t|| j||| <||| <qHt|dkrg}g}tt|jD]L}||kr|||||||n|||||j|qPtjj|t|dS|SdS)NrZcSsg|]}t|tr|qSr)r&r2)rPZidx_colrrrrRs z%_add_any_metadata..r\rGrFrDrrHr r@r8rrrr)r)rr;rrrr&r0rrr)r*rrArrayrrrFrr_ZTabler )rrZmodified_columnsZmodified_fieldsrrZZn_index_levelsZ n_columnsr{rridxrpr@Z metadata_tzZ convertedZ tz_aware_typerr\rrrrrs\        rcCs$tj|}|jdj|}|S)zB Make a datetime64 Series timezone-aware for the given tz r!)rrrr(r%r&)Zseriesr*rrrrs  r)N)rNT)NN)NFN)N)Prcollections.abcrZ concurrentrZconcurrent.futures.threadrr itertoolsrrer"rraZnumpyrrXrZ pyarrow.librrrrrr(Zbool_r r r rrrrrrrZunicode_r$r.r4rArNrrrdr`r|rrrrrrr~rrrrrrrrZsctypesrrrrrr rrZstr_rrrrrrrrrrrrrs    -V'@@    " b% >   9 (6 B A