U
    bD                     @   sB  d dl m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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j,ej-fej.ej/fgG d
d dZ0G dd dZ1dd Z2dd Z3dd Z4dd Z5dd Z6e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Z7dd Z8dd Z9d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j-ej/gG d!d" d"Z;d#d$ Z<d%d& Z=ejd ejej%ejejgG d'd( d(Z>d)d* Z?d+d, Z@dS )-    )contextmanagerN)	hashtable)isinc                   c   s"   t   z
d V  W 5 t   X d S N)tracemallocstartstop r	   r	   D/tmp/pip-unpacked-wheel-xg7vt6ln/pandas/tests/libs/test_hashtable.pyactivated_tracemalloc   s    
r   c                  C   s4   t  } | t dt f} ttdd | jS )NTc                 S   s   | j S r   )size)xr	   r	   r
   <lambda>       z,get_allocated_khash_memory.<locals>.<lambda>)	r   Ztake_snapshotZfilter_tracesZDomainFilterhtZget_hashtable_trace_domainsummapZ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S )TestHashTablec              	   C   sp  d}|d}t |dkst||ks(t||d t |dksDt||ksPt||dksbt||d d ||ks~t|d |kstt |dkst||dkst||d dkst||d ||kst|d |kstt |dkst||dkst||d dks*t|d |ks<ttjtt|d d	 ||d  W 5 Q R X d S )
N   7   r   *      )         match)lenAssertionErrorset_itemget_itempytestraisesKeyErrorstrself
table_typedtypeindextabler	   r	   r
   test_get_set_contains_len3   s.    z'TestHashTable.test_get_set_contains_lenc           	      C   s   |t jkr~d}| }t||}t|tj| }||j_||j_||| t	|D ]}|
|| || ks^tq^d S )NM   )r   PyObjectHashTablenparangeastypeint64flags	writeabler   ranger!   r   )	r'   r(   r)   writableNr+   keysvalsir	   r	   r
   test_mapP   s    
zTestHashTable.test_mapc                 C   sX   d}| }t || |}||j_|| t|D ]}||| |ks8tq8d S )N   )	r/   r0   r1   r3   r4   map_locationsr5   r!   r   )r'   r(   r)   r6   r7   r+   r8   r:   r	   r	   r
   test_map_locations]   s    
z TestHashTable.test_map_locationsc           	      C   sd   d}| }t || |}||j_|| ||}t |}t|t j	|t j	 d S )N   )
r/   r0   r1   r3   r4   r=   lookuptmassert_numpy_array_equalr2   )	r'   r(   r)   r6   r7   r+   r8   resultexpectedr	   r	   r
   test_lookupf   s    


zTestHashTable.test_lookupc                 C   sn   |t jt jfkrd}nd}| }t || |}|| t ||}||}t |dksjtd S )Nd   i   )	r/   int8uint8r0   r1   r=   r@   allr   )r'   r(   r)   r7   r+   r8   Z
wrong_keysrC   r	   r	   r
   test_lookup_wrongp   s    

zTestHashTable.test_lookup_wrongc           	      C   sb   |t jt jfkrd}nd}| }t || |}t |d}||j_||}t	
|| d S )NX     r   )r/   rH   rI   r0   r1   repeatr3   r4   uniquerA   rB   )	r'   r(   r)   r6   r7   r+   rD   r8   rO   r	   r	   r
   test_unique|   s    
zTestHashTable.test_uniquec              	   C   sz   |t jt jfkrd}nd}t ||}t @ | }|| t }| }||ks\t	~t dkslt	W 5 Q R X d S )N   i0u  r   )
r/   rH   rI   r0   r1   r   r=   r   sizeofr   )r'   r(   r)   r7   r8   r+   usedmy_sizer	   r	   r
   test_tracemalloc_works   s    
z$TestHashTable.test_tracemalloc_worksc              	   C   sF   t  6 | }t }| }||ks(t~t dks8tW 5 Q R X d S Nr   )r   r   rR   r   )r'   r(   r)   r+   rS   rT   r	   r	   r
   test_tracemalloc_for_empty   s    z(TestHashTable.test_tracemalloc_for_emptyc                 C   sL   |d}|  }|d dks t|d dks0td|ks<td|ksHtd S )NrM   r   r   Z
n_occupied	n_bucketsZupper_bound)	get_stater   )r'   r(   r)   r+   stater	   r	   r
   test_get_state   s    zTestHashTable.test_get_statec           	      C   s~   t ddD ]n}t||}||}| d }|| | d }||ksTt| }|| || d ks
tq
d S Nr   n   rX   )r5   r/   r0   r1   rY   r=   r   )	r'   r(   r)   r7   r8   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,   r;   r>   rE   rK   rP   rU   rW   r[   rb   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dd ZdS )TestPyObjectHashTableWithNansc                 C   sF   t d}t d}||k	stt }||d ||dksBtd S Nnanr   floatr   r   r.   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}||k	s0tt }||d ||dksVtd S rg   )complexrj   r   r   r.   r    r!   rk   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}||k	s6tt }||d ||dks\ttjt	d d}|| W 5 Q R X t
|jt
|kstd S )Nrh   r   r   r   r   ro   rj   r   r   r.   r    r!   r"   r#   r$   r%   valuer'   rl   rm   otherr+   errorr	   r	   r
   test_nan_complex_real   s    z3TestPyObjectHashTableWithNans.test_nan_complex_realc              	   C   s   t dtd}t dtd}t dtd}||k	s6tt }||d ||dks\ttjt	d d}|| W 5 Q R X t
|jt
|kstd S )Nr   rh   r   r   r   rq   rs   r	   r	   r
   test_nan_complex_imag   s    z3TestPyObjectHashTableWithNans.test_nan_complex_imagc                 C   sR   t df}t df}|d |d k	s(tt }||d ||dksNtd S )Nrh   r   r   ri   rk   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sNttjtd d}|| W 5 Q R X t	|j
t	|kstd S )Nr   r   rh   )r   r   r   r   )rj   r   r.   r    r!   r   r"   r#   r$   r%   rr   rs   r	   r	   r
   test_nan_in_nested_tuple   s    z6TestPyObjectHashTableWithNans.test_nan_in_nested_tupleN)	rc   rd   re   rn   rp   rv   rw   rx   ry   r	   r	   r	   r
   rf      s   rf   c                  C   s\   t dt dt dff} t dt dt dff}t| t|ksHtt| |sXtd S )Nrh   )rj   r   Zobject_hashr   Zobjects_are_equal)abr	   r	   r
   test_hash_equal_tuple_with_nans   s    r|   c                 C   s   t  }tjddddddgtjd}| |j_||\}}tjddddddgtjd}tjddgtjd}t	
|| t	
|| d S )Nr   r   rG   r)   r   )r   Int64HashTabler/   arrayr2   r3   r4   Zget_labels_groupbyintprA   rB   )r6   r+   r9   arrrO   Zexpected_arrZexpected_uniquer	   r	   r
   !test_get_labels_groupby_for_Int64   s    r   c               	   C   sr   d} t | t jjt j}t B t }|	| t
 }| }||ksTt~t
 dksdtW 5 Q R X d S )NrM   r   )r/   r0   r1   compatunicodeobject_r   r   StringHashTabler=   r   rR   r   )r7   r8   r+   rS   rT   r	   r	   r
   *test_tracemalloc_works_for_StringHashTable  s    
r   c               	   C   sH   t  8 t } t }|  }||ks*t~ t dks:tW 5 Q R X d S rV   )r   r   r   r   rR   r   )r+   rS   rT   r	   r	   r
   *test_tracemalloc_for_empty_StringHashTable  s    r   c                  C   s   t ddD ]~} t| tjjtj}t| }|	 d }|
| |	 d }||ksbtt }|
| ||	 d ks
tq
d S r\   )r5   r/   r0   r1   r   r   r   r   r   rY   r=   r   )r7   r8   r^   r_   r`   ra   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d Zd	S )
TestHashTableWithNansc                 C   s   t d}| }||kst||d t|dks6t||ksBt||dksTt||d t|dkspt||ks|t||dkstd S )Nrh   r   r   r   )rj   r   r    r   r!   r&   r	   r	   r
   r,   6  s    z/TestHashTableWithNans.test_get_set_contains_lenc                 C   sn   d}| }t j|t j|d}t || t j}||| t|dksNt|	t jd| d ksjtd S )NiL  r}   r   r   )
r/   fullrh   r0   r1   r2   r   r   r   r!   )r'   r(   r)   r7   r+   r8   r9   r	   r	   r
   r;   E  s    zTestHashTableWithNans.test_mapc                 C   sR   d}| }t j|t j|d}|| t|dks6t|t j|d ksNtd S )N
   r}   r   )r/   r   rh   r=   r   r   r!   )r'   r(   r)   r7   r+   r8   r	   r	   r
   r>   N  s    
z(TestHashTableWithNans.test_map_locationsc                 C   sJ   d}| }t j|t j|d}||}t t |rBt|dksFtd S )Ni  r}   r   )r/   r   rh   rO   rJ   isnanr   r   )r'   r(   r)   r7   r+   r8   rO   r	   r	   r
   rP   V  s
    
z!TestHashTableWithNans.test_uniqueN)rc   rd   re   r,   r;   r>   rP   r	   r	   r	   r
   r   ,  s   
	r   c                  C   sD   t  } tjdd tdD tjd}| |}t|dks@td S )Nc                 S   s   g | ]}t d qS )rh   rj   .0r:   r	   r	   r
   
<listcomp>`  s     z6test_unique_for_nan_objects_floats.<locals>.<listcomp>2   r}   r   	r   r.   r/   r   r5   r   rO   r   r   r+   r8   rO   r	   r	   r
   "test_unique_for_nan_objects_floats^  s    
r   c                  C   sD   t  } tjdd tdD tjd}| |}t|dks@td S )Nc                 S   s   g | ]}t td dqS )rh         ?)ro   rj   r   r	   r	   r
   r   g  s     z7test_unique_for_nan_objects_complex.<locals>.<listcomp>r   r}   r   r   r   r	   r	   r
   #test_unique_for_nan_objects_complexe  s    
r   c                  C   sJ   t  } tjdgdd tdD  tjd}| |}t|dksFtd S )Nr   c                 S   s   g | ]}d t dd ffqS )r   rh   r   r   r	   r	   r
   r   o  s     z5test_unique_for_nan_objects_tuple.<locals>.<listcomp>r   r}   r   r   r   r	   r	   r
   !test_unique_for_nan_objects_tuplel  s     
r   r)   c                   @   sD   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S )TestHelpFunctionsc                 C   sd   d}t || |}t |d}||j_t|d\}}t	t 
|| t |dks`td S )N+   r   F)r/   r0   r1   rN   r3   r4   r   value_countrA   rB   sortrJ   r   )r'   r)   r6   r7   rD   valuesr8   countsr	   r	   r
   test_value_count  s    z"TestHelpFunctions.test_value_countc              	   C   sX   t dddddddg|}||j_t|d\}}t|| t 	|dksTt
d S 	Nr   r   r      r?   rG   r<   F)r/   r   r1   r3   r4   r   r   rA   rB   rJ   r   )r'   r)   r6   r   r8   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 )NrF   r   r}   F)r/   rN   r0   r1   r3   r4   r   
duplicated	ones_likebool_rA   rB   )r'   r)   r6   r7   r   rC   rD   r	   r	   r
   test_duplicated_first  s    
z'TestHelpFunctions.test_duplicated_firstc                 C   s`   d}t ||}t ||}||j_||j_t||}t j|t jd}t	
|| d S )N   r}   )r/   r0   r1   r3   r4   r   ismemberr   r   rA   rB   )r'   r)   r6   r7   r   r   rC   rD   r	   r	   r
   test_ismember_yes  s    z#TestHelpFunctions.test_ismember_yesc                 C   sT   d}t ||}t || |}t||}t j|t jd}t|| d S )N   r}   )	r/   r0   r1   r   r   Z
zeros_liker   rA   rB   )r'   r)   r7   r   r   rC   rD   r	   r	   r
   test_ismember_no  s    z"TestHelpFunctions.test_ismember_noc                 C   s^   |t jt jfkrd}nd}t t ||d}d|d< ||j_t	|d}|dksZt
d S )N5   ig+  r   r   r   F)r/   rH   rI   rN   r0   r1   r3   r4   r   moder   )r'   r)   r6   r7   r   rC   r	   r	   r
   	test_mode  s    zTestHelpFunctions.test_modec              	   C   sB   t dddddddg|}||j_t|d}t|| d S r   )	r/   r   r1   r3   r4   r   r   rA   rB   )r'   r)   r6   r   r8   r	   r	   r
   test_mode_stable  s    z"TestHelpFunctions.test_mode_stableN)
rc   rd   re   r   r   r   r   r   r   r   r	   r	   r	   r
   r   u  s   		
r   c                  C   sN   t jtjt jd g} tjdg| d  tjd}t|d}|j	t
| ksJtd S )NTr   r}   F)pdZNAr/   rh   ZNaTr   r   r   r   r   r   r   )Znullsr   modesr	   r	   r
   test_modes_with_nans  s    r   c                 C   sT   t jddddddgt jd}| |j_t|}t jdddgt jd}t|| d S )Nr   r   r?   r}   r   r   )	r/   r   r   r3   r4   r   Zunique_label_indicesrA   rB   )r6   r8   rC   rD   r	   r	   r
   test_unique_label_indices_intp  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:tt|d\}}t|dkrft t |sjt|d dksztd S )Nr}   Tr   Fr   r?   )	r/   r   rh   r   r   r   r   rJ   r   )r'   r)   r   r8   r   r	   r	   r
   r     s     z*TestHelpFunctionsWithNans.test_value_countc                 C   sD   t jt jt jt jg|d}t|}t dddg}t|| d S )Nr}   FT)r/   r   rh   r   r   rA   rB   )r'   r)   r   rC   rD   r	   r	   r
   r     s    
z/TestHelpFunctionsWithNans.test_duplicated_firstc                 C   sb   t jt jt jt jg|d}t jt jt jg|d}t||}t jdddgt jd}t|| d S )Nr}   Tr/   r   rh   r   r   r   rA   rB   r'   r)   r   r   rC   rD   r	   r	   r
   r     s
    z+TestHelpFunctionsWithNans.test_ismember_yesc                 C   s\   t jt jt jt jg|d}t jdg|d}t||}t jdddgt jd}t|| d S )Nr}   r   Fr   r   r	   r	   r
   r     s
    z*TestHelpFunctionsWithNans.test_ismember_noc                 C   sJ   t jdt jt jt jg|d}t|ddks0tt t|dsFtd S )Nr   r}   TF)r/   r   rh   r   r   r   r   )r'   r)   r   r	   r	   r
   r     s    z#TestHelpFunctionsWithNans.test_modeN)rc   rd   re   r   r   r   r   r   r	   r	   r	   r
   r     s
   
r   c                  C   sL   dt dfdg} dt dfg}t| |}tjddgtjd}t|| d S )Nrz   rh   )r{   r   TFr}   )rj   r   r/   r   r   rA   rB   r   compsrC   rD   r	   r	   r
   test_ismember_tuple_with_nans  s
    
r   c                  C   sJ   ddddg} t td}t| |}tjddddgtjd}t|| d S )	Nrz   r   g      @y      @           FTr}   )listr5   r   r/   r   r   rA   rB   r   r	   r	   r
   +test_float_complex_int_are_equal_as_objects  s
    
r   )A
contextlibr   r   Znumpyr/   r"   Zpandas._libsr   r   Zpandasr   Zpandas._testingZ_testingrA   Zpandas.core.algorithmsr   r   r   markZparametrizer.   r   ZComplex128HashTableZ
complex128r~   r2   ZUInt64HashTableZuint64ZFloat64HashTableZfloat64ZComplex64HashTableZ	complex64ZInt32HashTableZint32ZUInt32HashTableZuint32ZFloat32HashTableZfloat32ZInt16HashTableZint16ZUInt16HashTableZuint16ZInt8HashTablerH   ZUInt8HashTablerI   ZIntpHashTabler   r   rf   r|   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r	   r	   r	   r
   <module>   s   














 =




	)	?	#	