U
    c9F                     @   s  d dl mZ d dlmZ d dlmZ d dlZd dlZd dlm	Z	m
Z
 d dlmZmZmZmZmZmZmZ d dlmZ d dlmZ d dlm  mZ d dlmZmZmZm Z m!Z! ej"d	d
 Z#G dd dZ$G dd dZ%ej&'de(d dddd dddgd)ej*e(d dddd dddgd)ej*e(ddd ddd ddgd)ej*gdge(ej+dej*dde(ej+dej*ddgdggdd Z,G dd dZ-dd Z.d d! Z/dS )"    )defaultdict)datetime)productN)is_ci_environmentis_platform_windows)NA	DataFrame
MultiIndexSeriesarrayconcatmerge)	safe_sort)_decons_group_indexget_group_indexis_int64_overflow_possiblelexsort_indexernargsortc                  C   s   d\} }}t tj| ||dftdd}|jdd|d< tjt|}|j| 	 }|j
d d  d	g |_
tt||_|d	  d9  < ||fS )
N)    i      ABCDEFGcolumns   Zaxisleftright)r   nprandomrandintlistsumpermutationleniloccopyr   tolistarangeindex)lowhighnr   ir    r/   =/tmp/pip-unpacked-wheel-g7fro6k3/pandas/tests/test_sorting.py
left_right"   s    
 r1   c                   @   s  e Zd Zejjdd Zdd Zejdddgdd	 Z	ejd
dde
edde
ed e
edd gdde
ede
edd e
edd gdde
eddde
ed e
edd gdde
ede
edd e
eddd ggdd Zejddde
edde
ed e
edd e
gdde
ede
edd e
edd e
gdde
eddde
ed e
edd e
gdde
ede
edd e
eddd e
gdde
edde
ed e
edd dd gdde
ede
edd e
edd dd gdde
eddde
ed e
edd dd gdde
ede
edd e
eddd dd ggdd ZdS ) TestSortingc                 C   sv  t t dt dt df}t d}t||||||||t jdd	}|ddddd	d
ddg}|ddd
d	ddddg}| d }| d }|j	 \}}	t
|j| |j	d\}}	t
|j| ttt|ddddd	d
ddg j}
t|
}
||
 d }| D ]<\}}|| ||d d d  ksFt|| |kstqt|t|ksrtd S )N  i  i	  )	ABCDEFGHvaluesr4   r5   r6   r7   r8   r9   r:   r;   r<   r   r   )r   Zconcatenater)   r   r    randngroupbyr#   r*   Z	sortleveltmassert_index_equalr"   maptupler<   comZasarray_tuplesafeitemsAssertionErrorr%   )selfr5   r4   dfZlgZrgr   r   Z	exp_index_Ztupsexpectedkvr/   r/   r0   test_int64_overflow2   s:    "

$
 zTestSorting.test_int64_overflowc                 C   sF   t d}t||||d}|ddddg}t|t|ksBtd S )NiE  )abcdrM   rN   rO   rP   )ranger   	from_dictr>   r%   rE   )rF   r<   dataZgroupedr/   r/   r0   'test_int64_overflow_groupby_large_rangeZ   s    z3TestSorting.test_int64_overflow_groupby_large_rangeaggZmeanZmedianc                    s~  t jddd}t jt|t|d }t ||| f}t jt|}|| }t|tdd}t j	dt|d \|d	< |d
< |
td}t|jjsttttt }}ttt||d	 |d
 D ]&\}}	}
|| |	 || |
 qt|t|ksttj| tdd} fdd}t || || fj}t|d	d
g|d }tt|  | d S )Ni i   )i         Zabcder      
   jimjoe)namesc                    s   t jttt  | ddS )NZf8Zdtype)r   ZfromiterrA   getattr)rM   rU   r/   r0   <lambda>z       zKTestSorting.test_int64_overflow_groupby_large_df_shuffled.<locals>.<lambda>)r   r*   )r   r    r!   choicer%   Zvstackr$   r   r"   r=   r>   r   ZgroupershaperE   r   ziprA   rB   appendr	   from_tupleskeysr<   TZ
sort_indexr?   assert_frame_equalr^   )rF   rU   arrr.   rG   ZgrrZ   r[   keyrM   rN   mifresr/   r_   r0   -test_int64_overflow_groupby_large_df_shuffleda   s&    "$ z9TestSorting.test_int64_overflow_groupby_large_df_shuffledzorder, na_position, expTlastrV   i   n   firstFh   rW   r   c                 C   sP   t jgd ttd t jgd  g}t|||d}t|t j|t jd d S )NrV   d   )Zordersna_positionr]   )	r   nanr"   rQ   r   r?   assert_numpy_array_equalr   intp)rF   orderrv   exprg   resultr/   r/   r0   test_lexsort_indexer   s    &z TestSorting.test_lexsort_indexerz ascending, na_position, exp, boxc                 C   s   t j| ddS NOr]   r   r   xr/   r/   r0   r`      ra   zTestSorting.<lambda>c                 C   s   t j| ddS r~   r   r   r/   r/   r0   r`      ra   c                 C   s   t j| ddS r~   r   r   r/   r/   r0   r`      ra   c                 C   s   t j| ddS r~   r   r   r/   r/   r0   r`      ra   c                 C   sR   |t jgd ttd t jgd  }t|d||d}tj|t |dd d S )NrV   ru   	mergesort)kind	ascendingrv   F)Zcheck_dtype)r   rw   r"   rQ   r   r?   rx   r   )rF   r   rv   r{   boxrD   r|   r/   r/   r0   test_nargsort   s    7(
   zTestSorting.test_nargsortN)__name__
__module____qualname__pytestmarkslowrL   rT   parametrizero   r"   rQ   r}   r   r/   r/   r/   r0   r2   1   s   
'
&&((
&&((&&((5r2   c                   @   s   e Zd Zdd Zejjdd Zejjejddddd	gd
d Z	ejjdd Z
ejjejddddd	gejdddgdd ZdS )	TestMergec                 C   sb   t tjddtddg d}t tjddtddg d}t||dd}t|d	ks^td S )
Nr3   r   ZABCDEFZG1r   ZG2outerhowi  )r   r   r    r=   r"   r   r%   rE   )rF   Zdf1Zdf2r|   r/   r/   r0   test_int64_overflow_outer_merge   s      z)TestMerge.test_int64_overflow_outer_mergec                 C   s   |\}}t ||dd}t|t|ks*ttj|d |d  dd |jd d d df jdd	}tj|d |dd |jd kstd S )
Nr   r   r   r   F)Zcheck_namesr   r   )r   r%   rE   r?   Zassert_series_equalr&   r#   name)rF   r1   r   r   outr|   r/   r/   r0   !test_int64_overflow_check_sum_col   s    z+TestMerge.test_int64_overflow_check_sum_colr   r   r   r   innerc              	   C   sV   |\}}t ||dd}|j|j dd tt||_t	|t |||dd d S )Nr   r   T)Zinplacer   sort)
r   sort_valuesr   r(   r   r)   r%   r*   r?   ri   )rF   r1   r   r   r   r   r/   r/   r0   test_int64_overflow_how_merge   s
    z'TestMerge.test_int64_overflow_how_mergec                 C   sX   |\}}t ||ddd}t|||j   t ||ddd}t|||j   d S )Nr   Fr   )r   r?   ri   r   r(   )rF   r1   r   r   r   r/   r/   r0   $test_int64_overflow_sort_false_order  s
    z.TestMerge.test_int64_overflow_sort_false_orderr   TFc                 C   s&  d\}}}t tj|||dfdtdd}|tjj	}t
|sJtt||gdd}t tj|||d dfdtdd}tjt||}	t|||j|	 gdd}tjt||d	< tjt||d
< tjt|}	|j|	  }tt||_tjt|}	|j|	  }tt||_tttt }
}|td D ]\}}|
| |d	  qV|td D ]\}}|| |d
  qg }|
 D ]B\}}||tjg}t||D ]\}}||||f  qҐq| D ]4\}}||
kr|D ]}||tj|f  qqdd }t |tdd	d
g d}||}|d	  |d
  |d	  |d
  @ tjt|ddd}|| }|||  }| |  A s|dkstt!||||d}|rtd}t"#||  || j$|dd t"#||| d S )N)r   r   i   r   int64r   r   T)Zignore_indexrX   r   r   c                 S   s$   |  | j } tt| | _| S N)r   r   r(   r   r)   r%   r*   )rG   r/   r/   r0   alignT  s    zCTestMerge.test_int64_overflow_one_to_many_none_match.<locals>.alignboolr]   )r   r   r   r   r   r   r   )r   )%r   r   r    r!   astyper"   applyr
   Znuniquer<   r   rE   r   rb   r%   r&   r=   r$   r'   r)   r*   r   Z	set_indexZiterrowsre   rD   getrw   r   ZnotnaZonesallanyr   r?   ri   r   )rF   r   r   r+   r,   r-   r   rc   r   r.   ZldictZrdictidxrowvalsrJ   Zlvalrvallvrvr   r   Zjmaskmaskframern   Zkcolsr/   r/   r0   *test_int64_overflow_one_to_many_none_match  s    



 
 z4TestMerge.test_int64_overflow_one_to_many_none_matchN)r   r   r   r   r   r   r   r   r   r   r   r   r/   r/   r/   r0   r      s   



r   zcodes_list, shaper   rX      ru   rW   rV   )rW   rV      '  r]   )r   r   c                 C   s>   t | |ddd}t||}t| |D ]\}}t|| q$d S )NT)r   Zxnull)r   r   rd   r?   rx   )Z
codes_listrc   Zgroup_indexZcodes_list2rM   rN   r/   r/   r0   test_deconss  s    
r   c                   @   s  e Zd Zejddddddgdddddggedejede	d	gg g ggd
d Z
ejdddgejdddddddddgddddddddgdgddddddddgddddddddgdgg g dggdd Zejje oe ddejdddgdd Zejddd egdd Zdd Zd d! Zejd"dd#ed$gdddgded%gddddgddged&ggd'd( Zejddddgdddggddejdgdddejgggd)d* Zejdddgejdddgd+d, Zd#S )-TestSafeSortzarg, expr   r   rX   r   rW   ZbaaacbZaaabbcr]   c                 C   s"   t |}t|}t|| d S r   )r   r   r   r?   rx   )rF   argr{   r|   rI   r/   r/   r0   test_basic_sort  s    	
zTestSafeSort.test_basic_sortverifyTFzcodes, exp_codes, na_sentinelr   c   c           
      C   sb   dddddg}t dddddg}t||||d\}}t j|t jd}	t|| t||	 d S )Nr   r   rX   r   rW   na_sentinelr   r]   r   r   r   ry   r?   rx   )
rF   r   codesZ	exp_codesr   r<   rI   r|   result_codesexpected_codesr/   r/   r0   
test_codes  s    
   
zTestSafeSort.test_codeszRIn CI environment can crash thread with: Windows fatal exception: access violation)reasonr   c              	   C   s   dddddg}t dddddg}ddddddddg}t|||d	\}}t jd||ddd|dgt jd
}t|| t|| d S )Nr   r   rX   r   rW   e   f   r   )r   r]   r   )rF   r   r<   rI   r   r|   r   r   r/   r/   r0   test_codes_out_of_bound  s     z$TestSafeSort.test_codes_out_of_boundr   c                 C   s   t j| tdS )Nr]   )r   r   objectr   r/   r/   r0   r`     ra   zTestSafeSort.<lambda>c                 C   sF   |ddddddg}t |}tjddddddgtd}t|| d S )NrN   r   r   rM   r]   )r   r   r   r   r?   rx   )rF   r   r<   r|   rI   r/   r/   r0   test_mixed_integer  s    zTestSafeSort.test_mixed_integerc                 C   s   t jddddgtd}dddddddg}t||\}}t jddddgtd}t jdddddddgt jd}t|| t|| d S )	NrN   r   r   rM   r]   rX   r   r   )r   r   r   r   ry   r?   rx   )rF   r<   r   r|   r   rI   r   r/   r/   r0   test_mixed_integer_with_codes  s    z*TestSafeSort.test_mixed_integer_with_codesc              	   C   sF   t jddt ddgtd}d}tjt|d t| W 5 Q R X d S )Nr   rX   r   r   r]   z,'[<>]' not supported between instances of .*match)	r   r   r   nowr   r   raises	TypeErrorr   )rF   rj   msgr/   r/   r0   test_unsortable  s    zTestSafeSort.test_unsortablezarg, codes, err, msgNz"Only list-like objects are allowedzOnly list-like objects or Nonezvalues should be uniquec              	   C   s*   t j||d t||d W 5 Q R X d S )Nr   )r<   r   )r   r   r   )rF   r   r   errr   r/   r/   r0   test_exceptions  s    	zTestSafeSort.test_exceptionsc                 C   s0   t |dd}t|}t |dd}t|| d S )NInt64r]   )r   r   r?   assert_extension_array_equal)rF   r   r{   rM   r|   rI   r/   r/   r0   test_extension_array  s    z!TestSafeSort.test_extension_arrayc                 C   st   t dddgdd}t|dd|dg||d\}}t dddgdd}tj dd|dgtjd}t|| t|| d S )Nr   r   rX   r   r]   r   r   )r   r   r   ry   r?   r   rx   )rF   r   r   rM   r|   r   Zexpected_valuesr   r/   r/   r0   test_extension_array_codes  s     
  
z'TestSafeSort.test_extension_array_codes)r   r   r   r   r   r   r"   r   r   r   r   r   Zskipifr   r   r   r   r   r   r   
ValueErrorr   rw   r   r   r/   r/   r/   r0   r     sV   
((

	

 .
r   c                  C   sH   t jdt jddgtd} t| }t jt jdddgtd}t|| d S )NrN   rM   r]   )r   r   rw   r   r   r?   rx   )r<   r|   rI   r/   r/   r0   test_mixed_str_nan  s    r   c                  C   sl   t ddttgdd} ddddg}t| |g}t|}tt ddttgddddddgg}t|| d S )NrX   r   r   r]   r   )r
   r   r	   Zfrom_arraysr   r?   r@   )Zarr1Zarr2Zmidxr|   rI   r/   r/   r0   test_safe_sort_multiindex  s    r   )0collectionsr   r   	itertoolsr   Znumpyr   r   Zpandas.compatr   r   Zpandasr   r   r	   r
   r   r   r   Zpandas._testingZ_testingr?   Zpandas.core.algorithmsr   Zpandas.core.commoncorecommonrC   Zpandas.core.sortingr   r   r   r   r   Zfixturer1   r2   r   r   r   Ztiler   r   r)   r   r   r   r   r/   r/   r/   r0   <module>   sF   $		
 5 """

n