o
    թZhe                     @   s  d dl mZ d dlm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mZ d dlmZ eded fdd	Zd
d Ze	jdejejfejejfejejfejejfej ej!fej"ej#fej$ej%fej&ej'fej(ej)fej*ej+fej,ej-fej.ej/fej0ej1fej2ej3fgG dd dZ4G dd dZ5G dd dZ6dd Z7dd Z8dd Z9dd Z:e	jde;dddd Z<e	jdej ej!fej(ej)fejejfej"ej#fgG d d! d!Z=d"d# Z>d$d% Z?d&d' Z@e	jd(ejejejejej!ej#ej%ej'ej)ej+ej-ej/ej1ej3gG d)d* d*ZAd+d, ZBd-d. ZCd/d0 ZDe	jd(ej!ej)ejej#gG d1d2 d2ZEd3d4 ZFd5d6 ZGdS )7    )	Generator)contextmanagerN)	hashtable)isinreturn)NNNc                   c   s*    t   z
d V  W t   d S t   w N)tracemallocstartstop r   r   W/var/www/html/lang_env/lib/python3.10/site-packages/pandas/tests/libs/test_hashtable.pyactivated_tracemalloc   s
   r   c                  C   s4   t  } | t dt f} tdd | jD S )NTc                 s   s    | ]}|j V  qd S r   )size).0xr   r   r   	<genexpr>   s    z-get_allocated_khash_memory.<locals>.<genexpr>)r   Ztake_snapshotZfilter_tracesZDomainFilterhtZget_hashtable_trace_domainsumZtraces)Zsnapshotr   r   r   get_allocated_khash_memory   s
   r   ztable_type, dtypec                   @   s   e Zd Zdd Zdd Zdd Zdd Zd	d
 Zdd Zdd Z	dd Z
dd Zdd Zdd Zdd Zejdedddd ZdS )TestHashTablec                 C   s  d}|d}t |dksJ ||vsJ ||d t |dks"J ||v s(J ||dks1J ||d d ||v s?J |d |v sGJ t |dksOJ ||dksXJ ||d dkscJ ||d ||v soJ |d |v swJ t |dksJ ||dksJ ||d dksJ |d |vsJ ||d d ||v sJ |d |v sJ t |dksJ ||dksJ ||d dksJ tjtt|d d	 ||d  W d    d S 1 sw   Y  d S )
N   7   r   *      )         match)lenset_itemget_itempytestraisesKeyErrorstrself
table_typedtypeindextabler   r   r   test_get_set_contains_len6   s<   "z'TestHashTable.test_get_set_contains_lenc                 C   s  |t jkr
td d}|ddd}t|dksJ ||vs J ||d t|dks.J ||v s4J ||dks=J tjtd	d
 |	  W d    n1 sSw   Y  ||d d |
d tj|v slJ ||v srJ |d |v szJ t|dksJ ||dksJ ||d dksJ |	 dksJ |
d ||v sJ |d |v sJ t|dksJ ||d dksJ |	 dksJ |d |vsJ tjtt|d d
 ||d  W d    d S 1 sw   Y  d S )NMask not supported for objectr   r   TZ	uses_maskr   r   r   NAr   r      r   r   )r   PyObjectHashTabler"   skipr   r    r!   r#   r$   get_naZset_napdr/   r%   r&   r   r   r   test_get_set_contains_len_maskZ   s@   




"z,TestHashTable.test_get_set_contains_len_maskc           	      C   s   |t jkr?d}| }t||}t|tj| }||j_||j_||| t	|D ]}|
|| || ks>J q/d S d S )NM   )r   Int64HashTablenparangeastypeint64flags	writeableZmap_keys_to_valuesranger!   )	r'   r(   r)   writableNr+   keysvalsir   r   r   test_map_keys_to_values~   s   
z%TestHashTable.test_map_keys_to_valuesc                 C   sX   d}| }t || |}||j_|| t|D ]}||| |ks)J qd S )N   )r8   r9   r:   r<   r=   map_locationsr>   r!   r'   r(   r)   r?   r@   r+   rA   rC   r   r   r   test_map_locations   s   
z TestHashTable.test_map_locationsc                 C   s   |t jkr
td d}|dd}t|| |}||j_|	|t
g d t|d D ]}||| |ks=J q0tjttt||d  d |||d   W d    n1 sbw   Y  | dksoJ d S )	Nr-   r0   Tr.   )FFTr   r   r   )r   r1   r"   r2   r8   r9   r:   r<   r=   rF   arrayr>   r!   r#   r$   reescaper%   r3   rG   r   r   r   test_map_locations_mask   s   


"z%TestHashTable.test_map_locations_maskc           	      C   sd   d}| }t || |}||j_|| ||}t |}t|t j	|t j	 d S )Nr0   )
r8   r9   r:   r<   r=   rF   lookuptmassert_numpy_array_equalr;   )	r'   r(   r)   r?   r@   r+   rA   resultexpectedr   r   r   test_lookup   s   


 zTestHashTable.test_lookupc                 C   sn   |t jt jfv rd}nd}| }t || |}|| t ||}||}t |dks5J d S )Nd   i   )r8   int8uint8r9   r:   rF   rM   all)r'   r(   r)   r@   r+   rA   Z
wrong_keysrP   r   r   r   test_lookup_wrong   s   

zTestHashTable.test_lookup_wrongc           
      C   s   |t jkr
td d}|dd}t|| |}tg d}||j_	|
|| |||}t|}	t|tj|	tj |td| g|tdg}t|tjtjdgtjd	 d S )
Nr-   r0   Tr.   )FTFr   FrT   r)   )r   r1   r"   r2   r8   r9   r:   rI   r<   r=   rF   rM   rN   rO   r;   )
r'   r(   r)   r?   r@   r+   rA   maskrP   rQ   r   r   r   test_lookup_mask   s   



&zTestHashTable.test_lookup_maskc           	      C   sb   |t jt jfv rd}nd}| }t || |}t |d}||j_||}t	
|| d S )NX     r   )r8   rU   rV   r9   r:   repeatr<   r=   uniquerN   rO   )	r'   r(   r)   r?   r@   r+   rQ   rA   r_   r   r   r   test_unique   s   
zTestHashTable.test_uniquec                 C   s   |t jt jfv rd}nd}t ||}t & | }|| t }| }||ks.J ~t dks6J W d    d S 1 sAw   Y  d S )N   i0u  r   )	r8   rU   rV   r9   r:   r   rF   r   sizeof)r'   r(   r)   r@   rA   r+   usedmy_sizer   r   r   test_tracemalloc_works   s   
"z$TestHashTable.test_tracemalloc_worksc                 C   s\   t  ! | }t }| }||ksJ ~t dksJ W d    d S 1 s'w   Y  d S Nr   )r   r   rb   )r'   r(   r)   r+   rc   rd   r   r   r   test_tracemalloc_for_empty   s   "z(TestHashTable.test_tracemalloc_for_emptyc                 C   sL   |d}|  }|d dksJ |d dksJ d|v sJ d|v s$J d S )Nr]   r   r   Z
n_occupied	n_bucketsupper_bound)	get_state)r'   r(   r)   r+   stater   r   r   test_get_state   s   zTestHashTable.test_get_stater@   r   n   c           	      C   sn   t ||}||}| d }|| | d }||ks#J | }|| || d ks5J d S Nrh   )r8   r9   r:   rj   rF   )	r'   r(   r)   r@   rA   preallocated_tablen_buckets_startn_buckets_endclean_tabler   r   r   test_no_reallocation   s   

z"TestHashTable.test_no_reallocationN)__name__
__module____qualname__r,   r5   rD   rH   rL   rR   rX   r[   r`   re   rg   rl   r"   markparametrizer>   rs   r   r   r   r   r   "   s    $$	
	r   c                   @   s   e Zd Zdd Zdd Zdd Zdd Zd	d
 Zej	
dddgej	
dejejddfejejddfejejddfejejddfejejddfejejddfgdd Zej	
dejejejejejejgdd ZdS )TestHashTableUnsortedc                 C   s   t  }|dd |ddksJ tjtdd |dd W d    n1 s+w   Y  tjtdd |d W d    d S 1 sHw   Y  d S )Nkeyr   z'key' has incorrect typer         z'val' has incorrect type)r   StringHashTabler    r!   r"   r#   	TypeError)r'   tblr   r   r   (test_string_hashtable_set_item_signature	  s   "z>TestHashTableUnsorted.test_string_hashtable_set_item_signaturec              	   C   s^   t ddt jddddg}|j|d t }|| t|	|t j
t|t jd d S )	NgX9v@gQ	@ir   r   r0   writerY   )r8   rI   nansetflagsr   Float64HashTablerF   rN   rO   rM   r9   r   intpr'   r?   Zxsmr   r   r   test_lookup_nan  s
   
&z%TestHashTableUnsorted.test_lookup_nanc                 C   s:   d}t |}|dd |dd t|dksJ d S )Nr{   g        r   g       r   )r   r   r    r   )r'   r@   r   r   r   r   test_add_signed_zeros  s
   
z+TestHashTableUnsorted.test_add_signed_zerosc                 C   s|   t dt ddd }t dt ddd }||ksJ ||ks$J t }||d ||d t|dks<J d S )Ndz=Ql          r   l         r   )structunpackpackr   r   r    r   )r'   ZNAN1ZNAN2r   r   r   r   test_add_different_nans)  s   z-TestHashTableUnsorted.test_add_different_nansc                 C   sX   t jg dt jd}|j|d t }|| t|	|t j
t|t jd d S )N)r   r   l            rY   r   )r8   rI   uint64r   r   UInt64HashTablerF   rN   rO   rM   r9   r   r   r   r   r   r   test_lookup_overflow7  s
   
&z*TestHashTableUnsorted.test_lookup_overflownvalsr   
   z&htable, uniques, dtype, safely_resizesobjectFTfloat64r;   int32r   c           
      C   s   t jtd|d}|j|d | }| }||d | |dd | }|j}	|r3|||dd ntjt	dd |||dd W d    n1 sMw   Y  |  |j|	ks]J d S )Nr]   rY   r   r   rT   zexternal reference.*r   )
r8   rI   r>   r   Z
get_labelsZto_arrayshaper"   r#   
ValueError)
r'   r?   ZhtableZuniquesr)   Zsafely_resizesr   rB   tmpZoldshaper   r   r   test_vector_resize?  s   z(TestHashTableUnsorted.test_vector_resizer   c                 C   s    t t jjd }||d d S )Nr   )	size_hint)r8   Ziinfouint32max)r'   r   r   r   r   r   test_hashtable_large_sizehintn  s   z3TestHashTableUnsorted.test_hashtable_large_sizehintN)rt   ru   rv   r   r   r   r   r   r"   rw   rx   r   r1   ZObjectVectorr}   r   ZFloat64Vectorr7   ZInt64VectorInt32HashTableZInt32Vectorr   ZUInt64Vectorr   r   r   r   r   r   ry     s:    #ry   c                   @   s<   e Zd Zdd Zdd Zdd Zdd Zd	d
 Zdd ZdS )TestPyObjectHashTableWithNansc                 C   sF   t d}t d}||usJ t }||d ||dks!J d S Nr   r   floatr   r1   r    r!   r'   nan1nan2r+   r   r   r   test_nan_float  s   z,TestPyObjectHashTableWithNans.test_nan_floatc                 C   sZ   t tdtd}t tdtd}||usJ t }||d ||dks+J d S r   )complexr   r   r1   r    r!   r   r   r   r   test_nan_complex_both  s   z3TestPyObjectHashTableWithNans.test_nan_complex_bothc                 C   s   t tdd}t tdd}t tdd}||usJ t }||d ||dks.J tjtd d}|| W d    n1 sEw   Y  t	|j
t	|ksUJ d S )Nr   r   r   r   r   r   r   r   r1   r    r!   r"   r#   r$   r%   valuer'   r   r   otherr+   errorr   r   r   test_nan_complex_real     z3TestPyObjectHashTableWithNans.test_nan_complex_realc                 C   s   t dtd}t dtd}t dtd}||usJ t }||d ||dks.J tjtd d}|| W d    n1 sEw   Y  t	|j
t	|ksUJ d S )Nr   r   r   r   r   r   r   r   r   r   test_nan_complex_imag  r   z3TestPyObjectHashTableWithNans.test_nan_complex_imagc                 C   sR   t df}t df}|d |d usJ t }||d ||dks'J d S )Nr   r   r   r   r   r   r   r   test_nan_in_tuple  s   

z/TestPyObjectHashTableWithNans.test_nan_in_tuplec                 C   s   ddt dfff}ddt dfff}d}t }||d ||dks'J tjtd d}|| W d    n1 s>w   Y  t|j	t|ksNJ d S )Nr   r   r   )r   r   r   r   )
r   r   r1   r    r!   r"   r#   r$   r%   r   r   r   r   r   test_nan_in_nested_tuple  s   z6TestPyObjectHashTableWithNans.test_nan_in_nested_tupleN)	rt   ru   rv   r   r   r   r   r   r   r   r   r   r   r     s    r   c                  C   s\   t dt dt dff} t dt dt dff}t| t|ks$J t| |s,J d S )Nr   )r   r   Zobject_hashZobjects_are_equal)abr   r   r   test_hash_equal_tuple_with_nans  s   r   c                 C   sv   t  }tjg dtjd}| |j_||\}}tjg dtjd}tjddgtjd}t	
|| t	
|| d S )N)r   r   rT   r   r   rT   rY   )r   r   rT   r   r   rT   r   r   )r   r7   r8   rI   r;   r<   r=   Zget_labels_groupbyr   rN   rO   )r?   r+   rB   arrr_   Zexpected_arrZexpected_uniquer   r   r   !test_get_labels_groupby_for_Int64  s   r   c                  C   s   d} t | t jt j}t ' t }|| t	 }|
 }||ks)J ~t	 dks1J W d    d S 1 s<w   Y  d S )Nr]   r   )r8   r9   r:   str_object_r   r   r}   rF   r   rb   )r@   rA   r+   rc   rd   r   r   r   *test_tracemalloc_works_for_StringHashTable  s   
"r   c                  C   s^   t  " t } t }|  }||ksJ ~ t dksJ W d    d S 1 s(w   Y  d S rf   )r   r   r}   r   rb   )r+   rc   rd   r   r   r   *test_tracemalloc_for_empty_StringHashTable  s   "r   r@   r   rm   c                 C   s|   t | t jt j}t| }| d }|| | d }||ks)J t }|| || d ks<J d S rn   )	r8   r9   r:   r   r   r   r}   rj   rF   )r@   rA   ro   rp   rq   rr   r   r   r   $test_no_reallocation_StringHashTable  s   


r   c                   @   s$   e Zd Zdd Zdd Zdd ZdS )TestHashTableWithNansc                 C   s   t d}| }||vsJ ||d t|dksJ ||v s!J ||dks*J ||d t|dks8J ||v s>J ||dksGJ d S )Nr   r   r   r   )r   r    r   r!   r&   r   r   r   r,     s   z/TestHashTableWithNans.test_get_set_contains_lenc                 C   sR   d}| }t j|t j|d}|| t|dksJ |t j|d ks'J d S )Nr   rY   r   )r8   fullr   rF   r   r!   )r'   r(   r)   r@   r+   rA   r   r   r   rH     s   
z(TestHashTableWithNans.test_map_locationsc                 C   sJ   d}| }t j|t j|d}||}t t |r!t|dks#J d S )Ni  rY   r   )r8   r   r   r_   rW   isnanr   )r'   r(   r)   r@   r+   rA   r_   r   r   r   r`     s
   
$z!TestHashTableWithNans.test_uniqueN)rt   ru   rv   r,   rH   r`   r   r   r   r   r     s    
r   c                  C   D   t  } tjdd tdD tjd}| |}t|dks J d S )Nc                 S   s   g | ]}t d qS )r   r   r   rC   r   r   r   
<listcomp>  s    z6test_unique_for_nan_objects_floats.<locals>.<listcomp>2   rY   r   r   r1   r8   rI   r>   r   r_   r   r+   rA   r_   r   r   r   "test_unique_for_nan_objects_floats     
r   c                  C   r   )Nc                 S   s   g | ]	}t td dqS )r         ?)r   r   r   r   r   r   r   &  s    z7test_unique_for_nan_objects_complex.<locals>.<listcomp>r   rY   r   r   r   r   r   r   #test_unique_for_nan_objects_complex$  r   r   c                  C   sJ   t  } tjdgdd tdD  tjd}| |}t|dks#J d S )Nr   c                 S   s   g | ]
}d t dd ffqS )r   r   r   r   r   r   r   r   .  s    z5test_unique_for_nan_objects_tuple.<locals>.<listcomp>r   rY   r   r   r   r   r   r   !test_unique_for_nan_objects_tuple+  s   
r   r)   c                   @   sL   e Zd Zdd Zdd Zdd Zdd Zd	d
 Zdd Zdd Z	dd Z
dS )TestHelpFunctionsc           	      C   sf   d}t || |}t |d}||j_t|d\}}}t	t 
|| t |dks1J d S )N+   r   F)r8   r9   r:   r^   r<   r=   r   value_countrN   rO   sortrW   )	r'   r)   r?   r@   rQ   valuesrA   counts_r   r   r   test_value_countH  s   z"TestHelpFunctions.test_value_countc                 C   s~   |t jkr
td t jdgd |d}t jdt jd}d|d< d|d< tj|d|d	\}}}t	|d
ks7J |d
ks=J d S )Nz%mask not implemented for object dtyper   r   rY   )r   Tr{   F)rZ   r   )
r8   r   r"   r2   rI   Zzerosbool_r   r   r   )r'   r)   r   rZ   rA   r   Z
na_counterr   r   r   test_value_count_maskQ  s   

z'TestHelpFunctions.test_value_count_maskc                 C   sP   t g d|}||j_t|d\}}}t|| t 	|dks&J d S )Nr   r   r      r0   rT   rE   Fr   )
r8   rI   r:   r<   r=   r   r   rN   rO   rW   )r'   r)   r?   r   rA   r   r   r   r   r   test_value_count_stable\  s
   z)TestHelpFunctions.test_value_count_stablec                 C   s\   d}t t ||d}||j_t|}t j|t j	d}d|d d d< t
|| d S )NrS   r   rY   F)r8   r^   r9   r:   r<   r=   r   
duplicated	ones_liker   rN   rO   )r'   r)   r?   r@   r   rP   rQ   r   r   r   test_duplicated_firstd  s   
z'TestHelpFunctions.test_duplicated_firstc                 C   s`   d}t ||}t ||}||j_||j_t||}t j|t jd}t	
|| d S )N   rY   )r8   r9   r:   r<   r=   r   ismemberr   r   rN   rO   )r'   r)   r?   r@   r   r   rP   rQ   r   r   r   test_ismember_yesm  s   z#TestHelpFunctions.test_ismember_yesc                 C   sT   d}t ||}t || |}t||}t j|t jd}t|| d S )N   rY   )	r8   r9   r:   r   r   Z
zeros_liker   rN   rO   )r'   r)   r@   r   r   rP   rQ   r   r   r   test_ismember_now  s   z"TestHelpFunctions.test_ismember_noc                 C   sb   |t jt jfv rd}nd}t t ||d}d|d< ||j_t	|dd }|dks/J d S )N5   ig+  r   r   r   F)
r8   rU   rV   r^   r9   r:   r<   r=   r   mode)r'   r)   r?   r@   r   rP   r   r   r   	test_mode  s   zTestHelpFunctions.test_modec                 C   s<   t g d|}||j_t|dd }t|| d S )Nr   Fr   )	r8   rI   r:   r<   r=   r   r   rN   rO   )r'   r)   r?   r   rA   r   r   r   test_mode_stable  s   z"TestHelpFunctions.test_mode_stableN)rt   ru   rv   r   r   r   r   r   r   r   r   r   r   r   r   r   4  s    		
r   c                  C   sR   t jtjt jd g} tjdg| d  tjd}t|dd }|j	t
| ks'J d S )NTr   rY   Fr   )r4   r/   r8   r   ZNaTrI   r   r   r   r   r   )Znullsr   modesr   r   r   test_modes_with_nans  s   r   c                 C   sJ   t jg dt jd}| |j_t|}t jg dt jd}t|| d S )N)r   r   r   r   r   r0   rY   )r   r   r   )	r8   rI   r   r<   r=   r   unique_label_indicesrN   rO   )r?   rA   rP   rQ   r   r   r   test_unique_label_indices_intp  s
   
r   c                  C   s   t jddddt j} t| }t j| ddd }t	j
||dd d	| t jdt| d
< t| }t j| ddd dd  }t	j
||dd d S )Nr   r   i   i   T)Zreturn_indexF)Zcheck_dtyperT   r   )r8   randomZdefault_rngZintegersr:   r   r   r   r_   rN   rO   choicer   )r   leftrightr   r   r   test_unique_label_indices  s   

r   c                   @   s4   e Zd Zdd Zdd Zdd Zdd Zd	d
 ZdS )TestHelpFunctionsWithNansc                 C   s   t jt jt jt jg|d}t|d\}}}t|dksJ t|d\}}}t|dkr5t t |s7J |d dks?J d S )NrY   Tr   Fr   r0   )r8   rI   r   r   r   r   rW   r   )r'   r)   r   rA   r   r   r   r   r   r     s    z*TestHelpFunctionsWithNans.test_value_countc                 C   sB   t jt jt jt jg|d}t|}t g d}t|| d S )NrY   )FTT)r8   rI   r   r   r   rN   rO   )r'   r)   r   rP   rQ   r   r   r   r     s   
z/TestHelpFunctionsWithNans.test_duplicated_firstc                 C   s`   t jt jt jt jg|d}t jt jt jg|d}t||}t jg dt jd}t|| d S )NrY   )TTTr8   rI   r   r   r   r   rN   rO   r'   r)   r   r   rP   rQ   r   r   r   r     s
   z+TestHelpFunctionsWithNans.test_ismember_yesc                 C   sZ   t jt jt jt jg|d}t jdg|d}t||}t jg dt jd}t|| d S )NrY   r   )FFFr   r   r   r   r   r     s
   z*TestHelpFunctionsWithNans.test_ismember_noc                 C   sR   t jdt jt jt jg|d}t|dd dksJ t t|dd s'J d S )Nr   rY   Tr   F)r8   rI   r   r   r   r   )r'   r)   r   r   r   r   r     s   z#TestHelpFunctionsWithNans.test_modeN)rt   ru   rv   r   r   r   r   r   r   r   r   r   r     s    
r   c                  C   s~   dt dfdg} dt dfg}d}tjt|d t| |}W d    n1 s(w   Y  tjddgtjd}t|| d S )	Nr   r   )r   r   z+isin with argument that is not not a Seriesr   TFrY   )	r   rN   Zassert_produces_warningFutureWarningr   r8   rI   r   rO   )r   compsmsgrP   rQ   r   r   r   test_ismember_tuple_with_nans  s   r   c                  C   sR   g d} t td}ttj| tdt|}tjg dtjd}t	|| d S )N)r   r   g      @y      @           rY   )FTTT)
listr>   r   r8   rI   r   Zasarrayr   rN   rO   )r   r   rP   rQ   r   r   r   +test_float_complex_int_are_equal_as_objects  s
   r   )Hcollections.abcr   
contextlibr   rJ   r   r   numpyr8   r"   Zpandas._libsr   r   Zpandasr4   Zpandas._testingZ_testingrN   Zpandas.core.algorithmsr   r   r   rw   rx   r1   r   ZComplex128HashTableZ
complex128r7   r;   r   r   r   r   ZComplex64HashTableZ	complex64r   r   ZUInt32HashTabler   ZFloat32HashTableZfloat32ZInt16HashTableZint16ZUInt16HashTableZuint16ZInt8HashTablerU   ZUInt8HashTablerV   ZIntpHashTabler   r   ry   r   r   r   r   r   r>   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   <module>   s    













 Sx=





	 	J	#