o
    Zhha                     @   sx  d dl Z d dlZd dlZd dlZd dlZd dlZd dlmZ	 d dl
Zd dlZd dlm  mZ d dlmZmZmZmZ zd dlmZ d dlZd dlZW n	 eyT   Y nw ejdddd Zdd	d
Zejdddgddd Zejdg dddd Z da!dd Z"dd Z#ej$jdd Z%				dddZ&dddZ'dddZ(d d! Z)ej$jd"d# Z*ej$jd$d% Z+ej$jd&d' Z,ej$jd(d) Z-ej$jd*d+ Z.ej$jd,d- Z/ej$jd.d/ Z0ej$jd0d1 Z1ej$jd2d3 Z2d4d5 Z3d6d7 Z4ej$jd8d9 Z5ej$jd:d; Z6ej$jd<d= Z7ej$jd>d? Z8ej$jd@dA Z9ej$jdBdC Z:ej$jdDdE Z;ej$jdFdG Z<ej$jdHdI Z=ej$jdJdK Z>ej$jej$j?dLe@dMdNdO ZAej$jdPdQ ZBej$jdRdS ZCej$jdTdU ZDej$jej$jEejFjG dVdWdXdY ZHej$jdZd[ ZId\d] ZJej$jd^d_ ZKej$jej$Ld`ej$Ldadbdc ZMej$jddde ZNdfdg ZOej$jdhdi ZPej$jdjdk ZQej$jej$jRej$jSej$jTdldm ZUdndo ZVdpdq ZWej$jXej$jdrds ZYej$jZej$jdtdu Z[dvdw Z\ej$jdxdy Z]dzd{ Z^d|d} Z_d~d Z`eaejbe	cg ddd Zdej$jRdd Zeej$jdd Zfej$jdd ZgdS )    N)read_featherwrite_feather
read_tableFeatherDataset)assert_frame_equalmodule)scopec                 C   s   | d S )NZfeather )Zbase_datadirr	   r	   Q/var/www/html/lang_env/lib/python3.10/site-packages/pyarrow/tests/test_feather.pydatadir)      r   feather_c                 C   s   t j| dS )Nprefix)tempfilemktempr   r	   r	   r
   random_path.   s   r         )r   paramsc                 c   s    | j V  d S N)paramrequestr	   r	   r
   version2   s   r   )Nuncompressedlz4zstdc                 c   s8    | j dv rtj| j st| j  d | j V  d S )N)r   r   z is not available)r   paCodecis_availablepytestskipr   r	   r	   r
   compression7   s   r#   c                 C   s   g a d S r   )
TEST_FILES)r   r	   r	   r
   setup_moduleB   r   r%   c              	   C   s0   t D ]}zt| W q tjy   Y qw d S r   )r$   osremoveerror)r   pathr	   r	   r
   teardown_moduleG   s   r*   c                   C   s:   t tj td W d    d S 1 sw   Y  d S )NZtest_invalid_file)r!   raisesr   ZArrowIOErrorr   r	   r	   r	   r
   test_file_not_existO   s   
"r,   Fc           	      C   sr   |d u rt  }|d u rd}t| t| ||||d tj|s%tdt|||d}|d u r2| }t	|| d S )Nr   )r#   compression_levelr   file not writtenuse_threads)
r   r$   appendr   r&   r)   exists	Exceptionr   r   )	dfexpectedr)   columnsr0   r   r#   r-   resultr	   r	   r
   _check_pandas_roundtripU   s   
r8   c                 C   sT   |d u rt  }t| t| ||d tj|stdt|}|	| s(J d S )Nr#   r.   )
r   r$   r1   r   r&   r)   r2   r3   r   equals)tabler)   r#   r7   r	   r	   r
   _check_arrow_roundtripn   s   
r<   c                    s8   d u rt  t  fdd}t|| d S )Nc                      s   t  d d S )Nr   )r   r	   r4   r)   r   r	   r
   f   s   z!_assert_error_on_write.<locals>.f)r   r$   r1   r!   r+   )r4   excr)   r   r?   r	   r>   r
   _assert_error_on_write{   s
   
rA   c                    s   d d}dd t |D } fddt  d D }t|}t| t|D ]%\}}| d |  |d  d |  f}t||d |d  || d	 q(t| }|	|s[J d S )
Nd   rC      c                 S   s   g | ]}t  qS r	   )r   .0ir	   r	   r
   
<listcomp>   s    z test_dataset.<locals>.<listcomp>c                    s&   i | ]}d t | tj d qS )col_r   )strnprandomrandnrE   
num_valuesr	   r
   
<dictcomp>   s    z test_dataset.<locals>.<dictcomp>r   r   r=   )
ranger   r;   r$   extend	enumerater   r   r   r:   )r   Z	num_filespathsdatar;   indexr)   rowsr	   rN   r
   test_dataset   s   



 rX   c                 C   sN   i }ddg}d}|D ]}t j|}||||< q
t|}t|| d d S )Nf4f8rC   r=   )rK   rL   rM   astypepd	DataFramer8   r   rU   numpy_dtypesrO   dtypevaluesr4   r	   r	   r
   test_float_no_nulls   s   
rb   c                 C   s   d}t  }t| tjjdd|d}dd tdD }tj	||}t
||| d t|}||s5J t|dd	}||sBJ t|dd	}t| | d S )
NrB   r   rC   sizec                 S      g | ]}d t | qS rI   rJ   rE   r	   r	   r
   rH          z#test_read_table.<locals>.<listcomp>r=   F)Z
memory_map)r   r$   r1   rK   rL   randintrQ   r   Tablefrom_arraysr   r   r:   r   r   	to_pandasr   rO   r)   ra   r6   r;   r7   r	   r	   r
   test_read_table   s   
rn   c                 C   s   d}t  }t| tjjdd|d}dd tdD }tj	||}t
||| d t|}t| | t|dd	}t| | t|dd	}||sOJ d S )
N)
   ro   r   ro   rc   c                 S   re   rf   rg   rE   r	   r	   r
   rH      rh   z$test_use_threads.<locals>.<listcomp>r=   Fr/   )r   r$   r1   rK   rL   ri   rQ   r   rj   rk   r   r   r   rl   r   r:   rm   r	   r	   r
   test_use_threads   s   
rp   c                 C   s   d}t  }t| tjjdd|ddk }ddg}g }g }|D ]}tj||}|tj	||d tj
||< || qtj||d	}	t|	 |	 }
t|
| d
 d S )NrC   r   ro   rc      rY   rZ   masknamesr=   )r   r$   r1   rK   rL   ri   rM   r[   r   arraynanr;   r<   rl   r8   )r   rO   r)   	null_maskZdtypesexpected_colsarraysnamera   r;   r4   r	   r	   r
   test_float_nulls   s    

r|   c           	      C   s   i g }}g d}d}|D ]}t jjdd|d}||||< ||| qt|}t|| d tj	||d}t
| d S )Ni1i2i4i8u1u2Zu4u8rC   r   rc   r=   rt   )rK   rL   ri   r[   r1   r\   r]   r8   r   r;   r<   )	r   rU   arrr_   rO   r`   ra   r4   r;   r	   r	   r
   test_integer_no_nulls   s   

r   c                 C   sR   i }dg}d}|D ]}t jjdd|d}||||< q	t|}t|| d d S )NZlonglongrC   r   rc   r=   )rK   rL   ri   r[   r\   r]   r8   r^   r	   r	   r
   test_platform_numpy_integers  s   
r   c                 C   s   t  }t| g d}d}g }tjjdd|ddk }g }|D ]$}tjjdd|d}|tj||d |d}	tj	|	|< ||	 qtj
||d	}
t|
 |
 }t|| d
 d S )Nr}   rC   r   ro   rc   rq   rr   rZ   rt   r=   )r   r$   r1   rK   rL   ri   r   rv   r[   rw   r;   r<   rl   r8   )r   r)   Z
int_dtypesrO   rz   rx   ry   r{   ra   r5   r;   r4   r	   r	   r
   test_integer_with_nulls  s"   


r   c                 C   s:   d}t jd tdt j|dki}t|| d d S )NrC   r   boolsr=   )rK   rL   seedr\   r]   rM   r8   )r   rO   r4   r	   r	   r
   test_boolean_no_nulls;  s   r   c                 C   s   t  }t| d}tjd tjjdd|ddk }tjjdd|ddk }tjtj	||dgdgd	}t
| | }t|| d
 d S )NrC   r   ro   rc   rq   rD   rr   r   rt   r=   )r   r$   r1   rK   rL   r   ri   r   r;   rv   r<   rl   r8   )r   r)   rO   rs   ra   r;   r4   r	   r	   r
   test_boolean_nullsE  s   
r   c                 C   s\   t  }t| tddD ]}tjjtjd gtt| t	 dgdgd}t
| qd S )N      )typer   rt   )r   r$   r1   rQ   r   rj   rk   rv   listZfloat64r<   )r   r)   rG   r;   r	   r	   r
   test_buffer_bounds_errorX  s   
 
r   c                 C   s6   d}t jjtjg d| tdgdgd}t| d S )NrC   )FNT)r`   r   rt   )r   rj   rk   rK   rv   objectr<   )r   repeatsr;   r	   r	   r
   test_boolean_object_nullse  s   r   c                 C   s   t jdkr
td G dd d}tjtddd d| tjgdd	d
gd}t	 }z	t
||| d W n	 ty;   Y nw tj|rDJ d S )Nwin32z/Windows hangs on to file handle for some reasonc                   @      e Zd ZdS )z6test_delete_partial_file_on_error.<locals>.CustomClassN__name__
__module____qualname__r	   r	   r	   r
   CustomClasss      r   rD      foobar)numbersstringsr   r   r6   r=   )sysplatformr!   r"   r\   r]   rQ   rK   rw   r   r   r3   r&   r)   r2   )r   r   r4   r)   r	   r	   r
   !test_delete_partial_file_on_errorn  s    

r   c                 C   s   d}dd ddt jg}td|| i}g d}td|| i}t||| d g d}td|| i}td|| i}t||| d d	d ddt jg}td|| i}g d}td|| i}t||| d d S )
N  r   r   quxr   )r   Ns   bars   quxNr=   fooNr   r   Nr   )rK   rw   r\   r]   r8   )r   r   ra   r4   Z	ex_valuesr5   r	   r	   r
   test_strings  s   r   c                 C   s$   t ddgd i}t|| d d S )Nr    ro   r=   r\   r]   r8   r   r4   r	   r	   r
   test_empty_strings     r   c                 C   s$   t dd gd i}t|| d d S )NZall_nonero   r=   r   r   r	   r	   r
   test_all_none  r   r   c                 C   s4   t ddd}|j|jdd}t|| d d S )Nr   r   rq   )NNNABcategory)r   r=   )r\   r]   Zassignr   r[   r8   r   r	   r	   r
   test_all_null_category  s   r   c                 C   s.   dd t dD }t|}t|d| d d S )Nc                 S   s   i | ]}d  |dgd qS )zc{}r   ro   )formatrE   r	   r	   r
   rP     s    z+test_multithreaded_read.<locals>.<dictcomp>rC   T)r0   r   )rQ   r\   r]   r8   )r   rU   r4   r	   r	   r
   test_multithreaded_read  s
   
r   c                 C   s<   t dt jt jd dgd }td|i}t|| d d S )Nr   r   r   ro   r   r=   )rK   rv   rw   r\   r]   r8   )r   ra   r4   r	   r	   r
   test_nan_as_null  s   r   c                 C   sj   d}dd ddt jg}td|| i}|d d|d< g d}tdt|| i}t||| d d S )	Nr   r   r   r   r   r   r   r=   )rK   rw   r\   r]   r[   ZCategoricalr8   )r   r   ra   r4   r5   r	   r	   r
   test_category  s   r   c                 C   sB   t dt jdddi}|jjdjd|d< t|| d d S )	Nnaivez
2016-03-28ro   )ZperiodsutczAmerica/Los_Angeleswith_tzr=   )r\   r]   Z
date_ranger   dttz_localizeZ
tz_convertr8   r   r	   r	   r
   test_timestamp  s
   r   c              
   C   sJ   t dt dddd t dddgi}|jjd|d< t|| d d S )Ntesti  r   rq   r   r   r=   )r\   r]   	Timestampr   r   r   r8   r   r	   r	   r
   test_timestamp_with_nulls  s   r   znot supported)reasonr+   c                  C   0   t dt dd t dgi} t| dd d S )Nr   1 day3 dayr   r=   r\   r]   Z	Timedeltar8   r4   r	   r	   r
   test_timedelta_with_nulls_v1  s
   r   c                  C   r   )Nr   r   r   r   r=   r   r   r	   r	   r
   test_timedelta_with_nulls  s
   r   c                 C   s:   t dt g di}|jjd|d< t|| d d S )Nr   )l    ^cQ)! Nl   6"Ir   r   r=   )r\   r]   ZDatetimeIndexr   r   r   r8   r   r	   r	   r
   (test_out_of_float64_timestamp_with_nulls  s
   r   c                 C   sB   t g dg dd}|}| dkr|jtd}t||| d d S )Nr   r   rq      )TFTF)r   r   r   r   r=   )r\   r]   renamerJ   r8   r   r4   r5   r	   r	   r
   test_non_string_columns  s   
r   zunicode filenames not supported)r   c                 C   s4   d d}tdg di}t|t|d| d d S )Ns   Besa_Kavajë.featherzutf-8r   r   r   )r)   r   )decoder\   r]   r8   r   )r   r{   r4   r	   r	   r
   test_unicode_filename  s
   

r   c                 C   sB   t g dg dg dd}|ddg }t||| ddgd d S )Nr   )rD            )r   rq   rD   r   )r   boowoor   r   )r   r6   r   r   r	   r	   r
   test_read_columns  s   
r   c                 C   st   t  }t| d}tjd tjjdd|d}td|i}t	|| td|d|d  i}t
||d d S )	NrC   r   ro   rc   ZintsZ	more_intsr   )r)   )r   r$   r1   rK   rL   r   ri   r   r;   r   r<   )r   r)   rO   ra   r;   r	   r	   r
   test_overwritten_file  s   

r   c                 C   sZ   t  }tjtdddg dd }t||| d |	d t
|}t|| d S )N   r   rq   abcr   r=   r   )ioBytesIOr\   r]   rK   arangereshapecopyr   seekr   r   )r   bufr4   r7   r	   r	   r
   test_filelike_objects/  s   
r   zignore:Sparse:FutureWarningz(ignore:DataFrame.to_sparse:FutureWarningc                 C   sR   t jjjs
td g dg dd}t|jdd}|	 }t
||| d d S )Nz2version of pandas does not support SparseDataFramer   r   r   )r   r   r   r   r   )Z
fill_valuer=   )r   Zpandas_compatZ_pandas_apiZ
has_sparser!   r"   r\   r]   Z	to_sparseZto_denser8   )r   rU   r4   r5   r	   r	   r
   test_sparse_dataframe>  s   

r   c                  C   s2   t jtdddtdd } t| t d S )Nr   r   rq   Zaaar   )	r\   r]   rK   r   r   r   r   rA   
ValueErrorr   r	   r	   r
   test_duplicate_columns_pandasL  s   r   c                  C   s>   t jg dg dg dgg dd} t|  t| tdd d S )Nr   r   rD   r   r   r   	   )r   r   r   rt   r   r=   )r   r;   r<   rA   r   )r;   r	   r	   r
   test_duplicate_columnsV  s   $r   c                  C   sN   G dd d} t d|  |  gi}t|t t dg di}t|t d S )Nc                   @   r   )ztest_unsupported.<locals>.ANr   r	   r	   r	   r
   r   c  r   r   r   )r   r   g       @)r\   r]   rA   r   	TypeError)r   r4   r	   r	   r
   test_unsupported]  s
   
r   c                  C   sx   t dtdi} t| }t }t||ddd |	 }tj
t|}|jdks/J t|ddks:J d S )Nr   r      r   	chunksizer   r   r   )r\   r]   rK   r   r   r;   r   r   r   getvalueZipc	open_fileBufferReaderZnum_record_batcheslenZ	get_batch)r4   r;   r   r7   Zipc_filer	   r	   r
   test_v2_set_chunksizen  s   
r   c                  C   s   t dtdi} g d}|D ]\}}t| ||d qt }tjt	dd t
| |ddd	 W d    n1 s:w   Y  tjt	d
d t
| |ddd W d    n1 sYw   Y  tjt	dd t
| |dd W d    d S 1 sxw   Y  d S )Nr   r   ))r   N)r   N)r   r   )r   r   )r   r   )r   ro   )r#   r-   z2Feather V1 files do not support compression optionmatchr   r   )r#   r   z0Feather V1 files do not support chunksize optioni   r   z"compression="snappy" not supportedsnappyr9   )r\   r]   rK   r   r8   r   r   r!   r+   r   r   )r4   casesr#   r-   r   r	   r	   r
   test_v2_compression_options}  s0   
"r   c                  C   s   t jdstd t jtddgdgd} t	 }t
| | | }t	 }t
| |dd | }t|t|k s>J d S )	NZ	lz4_framez+LZ4 compression support is not built in C++r   i f0rt   r   r9   )r   r   r    r!   r"   r;   rK   repeatr   r   r   r   r   )tr   Zdefault_resultZuncompressed_resultr	   r	   r
   test_v2_lz4_default_compression  s   

r  c                  C   sl   t jt g dg d ggdgd} t }tjtdd t| |dd W d    d S 1 s/w   Y  d S )Nr   r   rt   z[Unsupported Feather V1 type: list<item: int64>. Use V2 format to serialize all Arrow types.r   r   r=   )	r   r;   rv   r   r   r!   r+   r   r   )r;   r   r	   r	   r
   test_v1_unsupported_types  s   ""r  c                 C   s$   t dtdi}t|| d d S )Nr   i r=   )r\   r]   rK   r   r8   r   r	   r	   r
   test_large_dataframe  s   r  c                  C   s   dgdd gd d  } t d| i}t }t||dd tt| }t	|| t
jtdd t|t d	d W d    d S 1 sHw   Y  d S )
N   xi   r   i   Zbyte_colr=   zp'byte_col' exceeds 2GB maximum capacity of a Feather binary column. This restriction may be lifted in the futurer   r   )r\   r]   r   r   r   r   r   r   r   r   r!   r+   r   )ra   r4   r   r7   r	   r	   r
   !test_chunked_binary_error_message  s   
"r  c                 C   s6   t jt g dgdgd}t| d }t|| d S )Nr   r   rt   data.feather)r   r;   rv   rJ   r<   )tempdirr   r;   r)   r	   r	   r
   test_feather_without_pandas  s   r  c                 C   s   t jtdddg dd}t|ddg|ddg | d t|d	d
g|ddg | d t|ddg|ddg | d t|dd	g|ddg | d d S )Nr   r   rq   r   r   r   r   )r6   r5   r   r   r   r   r   )r\   r]   rK   r   r   r8   r   r	   r	   r
   test_read_column_selection  s    
r	  c                 C   s   t jg dg dg dgg dd}t| d }t|||d t jg dg dg dgg dd}g dg d	fD ]}t||d
}||sHJ q9d S )Nr   r   r   r   rt   r  r=   r   r   r   )r   r   r   r   )r   r;   rJ   r   r   r:   )r  r   r;   r)   r5   Zcol_selectionr7   r	   r	   r
   %test_read_column_duplicated_selection  s   $r  c                 C   s   t jg dg dg dgg dd}t| d }t||dd t|}||s*J t|d	dgd
}|jddgks;J tt	 t|ddgd
 W d    d S 1 sTw   Y  d S )Nr   r   r   r
  rt   r  r   r=   r   r   r   r   )
r   r;   rJ   r   r   r:   Zcolumn_namesr!   r+   r   )r  r;   r)   r7   r	   r	   r
   #test_read_column_duplicated_in_file
  s   $"r  c                 C   s   t dt jjg dg dgddgdi}t|| d t dt dd	gd
dggi}t|| d t dt dd	gd
dggddgd ggi}t|| d d S )Ncolr   r   f1f2rt   r9   r   r   rq   r   rD   r   )r   r;   ZStructArrayrk   r<   rv   )r#   r;   r	   r	   r
   test_nested_types  s   
 ,r  )r   r   r   c                 C   s   t | |d d S )Nr9   )r<   )r;   r#   r	   r	   r
   test_roundtrip*  s   r  c                 C   s0   t dtdi}t| d }||sJ d S )Nr   rD   z)v0.17.0.version.2-compression.lz4.feather)r   r;   rQ   r   r:   )r   r5   r7   r	   r	   r
   Atest_feather_v017_experimental_compression_backward_compatibility/  s   r  c                 C   sN   t jdg dig dd}| dkr|jddjtd}n|}t||| d	 d S )
Nr   r   r   )rV   r   T)Zdropr   r=   )r\   r]   Zreset_indexr   rJ   r8   r   r	   r	   r
   test_preserve_index_pandas@  s
   r  c                 C   s   ddl m } td|d|dgi}t|| d  |d}|d}t| d dd}||d d ks8J ||d d	 ksBJ d S )
Nr   )datetimedatez
1654-01-01z
1920-01-01ztest_resolution.featherT)Ztimestamp_as_objectr   )r  r\   r]   fromisoformatr   r   )r  r  r4   Z
expected_0Z
expected_1r7   r	   r	   r
   0test_feather_datetime_resolution_arrow_to_pandasL  s   

r  )r   )NNNFNNN)NN)Nr   )hr   r&   r   r   r!   Z
hypothesishZhypothesis.strategiesZ
strategiesstnumpyrK   Zpyarrowr   Zpyarrow.tests.strategiesZtestsZpastZpyarrow.featherr   r   r   r   Zpandas.testingr   Zpandasr\   Zpyarrow.pandas_compatImportErrorZfixturer   r   r   r#   r$   r%   r*   markr,   r8   r<   rA   rX   rb   rn   rp   r|   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   Zxfailr   r   r   r   r   Zskipifr)   supports_unicode_filenamesr   r   r   r   filterwarningsr   r   r   r   r   r   r   r   r   r  r  Zslowr  Zlarge_memoryr  r  r	  r  r  r  givenZ
all_tablesZsampled_fromr  r  r  r  r	   r	   r	   r
   <module>   s  
















	
	









	








	

%



