U
    f/eS                  8   @   sH  d dl m Z  d dlZd dlZd dlZd dlmZmZmZm	Z	 d dl
mZ dd Zejddgddggejjd	d
d dD edededededgejdddejdddejdddejdddejdddgeded  eded  eded  eded  eded  geddeddeddeddeddgg dd dejdddddddddgfddddddddgfddddddddgfdddddd ddgfd!dddddddgfd!ddd d dd dgfd!ddd"d"dd"dgfd!ddd#d#d d#dgfd$ddd"d"dd"dgfd$ddd#d#dd#dgfd$dddddddgfd$ddddd ddgfd%dddddd"dgfd%ddd ddd#dgfd%ddd"ddddgfd%ddd#dd ddgfd&ddddd"ddgfd&ddd'd'dd'd(gfd&ddd"d"dd"dgfd&ddddd'dd(gfgd)d* Zejddgddggejd	ej ej ejdejejejggejd+ddd,d-d-ejd.ejd/d/gfddd0d1d1d-dd-d2d2gfddd3d-d-d2d"d2d/d/gfddd,d/d/ejd.ejd-d-gfddd0d2d2d-dd-d1d1gfddd3d/d/d2d"d2d-d-gfd!dd,ddejd"ejddgfd!dd0d"d"dddd4d4gfd!dd3ddd4d"d4ddgfd!dd,ddejd"ejddgfd!dd0d4d4dddd"d"gfd!dd3ddd4d"d4ddgfd$dd,ddejd"ejddgfd$dd0dddddd5d5gfd$dd3ddd5d"d5ddgfd$dd,ddejd"ejddgfd$dd0d5d5dddddgfd$dd3ddd5d"d5ddgfd%dd,ddejd"ejddgfd%dd0d"ddddd4d5gfd%dd3ddd4d"d5ddgfd%dd,ddejd"ejddgfd%dd0d4d5dddd"dgfd%dd3ddd4d"d5ddgfd&dd,ddejdejd"d"gfd&dd0dddd"dddgfd&dd3dddddd"d"gfd&dd,d"d"ejdejddgfd&dd0dddd"dddgfd&dd3d"d"dddddgfgd6d7 Zejddgddggejjd	d8d d9D ededejedededejejgejdddejdddejejdddejdddejdddejejgeded  eded  ejeded  eded  eded  ejejgeddeddejeddeddeddejejgg d:d dejd;ddd,dddejdddejejgfddd,dddejdddejejgfddd,dddejdddejejgfddd,dddejd ddejejgfd!dd,dddejdddejejgfd!dd,dd d ejdd dejejgfd!dd,dd"d"ejdd"dejejgfd!dd,dd#d#ejd d#dejejgfd$dd,dd"d"ejdd"dejejgfd$dd,dd#d#ejdd#dejejgfd$dd,dddejdddejejgfd$dd,dddejd ddejejgfd%dd,dddejdd"dejejgfd%dd,dd dejdd#dejejgfd%dd,dd"dejdddejejgfd%dd,dd#dejd ddejejgfd&dd,dddejd"ddejejgfd&dd,dd'd'ejdd'd(ejejgfd&dd,dd"d"ejdd"dejejgfd&dd,dddejd'dd(ejejgfddd3dddd5dddd5d5gfddd3dd<d<d=d>d<d?d=d=gfddd3dddd5dddd5d5gfddd3dd?d?d=d@d?d<d=d=gfd!dd3dddd4dddd4d4gfd!dd3dd@d@dAd>d@d?dAdAgfd!dd3dd"d"d4dd"dd4d4gfd!dd3ddBdBdAd@dBd<dAdAgfd$dd3dd"d"dCdd"ddCdCgfd$dd3ddBdBdd>dBd?ddgfd$dd3ddddCddddCdCgfd$dd3dd>d>dd@d>d<ddgfd%dd3dddd4dd"dd5dCgfd%dd3dd@d<dAd>dBd?d=dgfd%dd3dd"dd4dddd5dCgfd%dd3ddBd?dAd@d>d<d=dgfd&dd3ddddd"ddddgfd&dd3dd<d<ddAd<d?ddgfd&dd3dd"d"ddd"dddgfd&dd3ddAdAdd<dAd?ddgfg(dDdE ZejdFdd"d"d"d"d"gfdd#d#d#d#d#gfgdGdH ZejdIdJdKdLdMdNdOgejdPddgdQdR ZejdSdd!d$d%d&gejdTddgejdUd,d0d3gejdVddgejd	dWdWdXdWdYgdWejdXejdYggdZd[ ZejdUdd\d]gejdSdd!d$d%d&gejdTddgejdVddgejd	dWdWdXdWdYgdWejdXejdYgd]ejd^ejd.ggd_d` Zdadb Zejdcd]d^gd]d]gddgfd]d]d^d^gd]d^d]d^gd?dd?dgfd]d]d^d^gd]d^d]ejgd?ddejgfd]d]d^gd]d^ejgd?dejgfgddde Zdfdg Zejdhddgdidj Zdkdl Zdmdn Z dodp Z!dS )q    )datetimeN)	DataFrameNaTSeriesconcatc                  C   s  t dd} t dd}tjjdddd}tjjdddd}ttjd| |||d}|dd	gj	
 }d
d |dd	gD }t|dd}||j}t || |dd	gj	j
dd}dd |dd	gD }t|dd}||j}t || d S )N
   d      r   i  )size)valuekey1key2r   r   c                 S   s   g | ]\}}|j  qS  r   rank.0keyZpiecer   r   B/tmp/pip-unpacked-wheel-tiezk1ph/pandas/tests/groupby/test_rank.py
<listcomp>    s     z#test_rank_apply.<locals>.<listcomp>axisTpctc                 S   s   g | ]\}}|j jd dqS )Tr   r   r   r   r   r   r   '   s    )tmZrands_arraynprandomrandintr   ZrandnZtakegroupbyr   r   r   Zreindexindexassert_series_equal)Zlev1Zlev2Zlab1Zlab2dfresultexpectedr   r   r   test_rank_apply   s,    
r$   grpsZquxZquuxvalsc                 C   s$   g | ]}t jd d dd dg|dqS          dtype)r   arrayr   r,   r   r   r   r   2   s   r   )i8i4i2i1u8Zu4u2u1f8f4f2z
2018-01-02z
2018-01-08z
2018-01-06z
US/Pacific)tzDc                 C   s   t | d S Nr   typexr   r   r   <lambda>T       r@   )idszties_method,ascending,pct,expZaverageTF       @g      @g      @g?      ?g?g?ming      @g333333?maxfirstdenseUUUUUU?UUUUUU?c                 C   s   t | t|}|}t|t|  }t|t jr@t j||jd}t||d}|	dj
|||d}	t|t|  dgd}
t|	|
 d S )Nr+   r   valr   )method	ascendingr   rL   columnsr   repeatlenlist
isinstanceZndarrayr-   r,   r   r   r   r   assert_frame_equal)r%   r&   ties_methodrN   r   expr   	orig_valsr!   r"   exp_dfr   r   r   test_rank_args/   s    Ar[   z#ties_method,ascending,na_option,expZkeep      ?   g      @top      @g      @Zbottomg      @g      @c           
      C   sf   t | t|}|t|  }t||d}|dj|||d}t|t|  dgd}	t||	 d S )NrK   r   )rM   rN   	na_optionrL   rO   )r   rR   rS   r   r   r   r   rV   )
r%   r&   rW   rN   r`   rX   r   r!   r"   rZ   r   r   r   test_infs_n_nans~   s    )
  ra   c                 C   s0   g | ](}t jd d t jdd dt jt jg|dqS r'   )r   r-   nanr.   r   r   r   r      s   )r6   r7   r8   c                 C   s   t | d S r;   r<   r>   r   r   r   r@      rA   z'ties_method,ascending,na_option,pct,expg      ?g      ?g      ?      ?g      ?g      ?g      ?g       @c                 C   s   t | t|}|}t|t|  }t|t jr@t j||jd}t||d}	|		dj
||||d}
t|t|  dgd}t|
| d S )Nr+   rK   r   rM   rN   r`   r   rL   rO   rQ   )r%   r&   rW   rN   r`   r   rX   r   rY   r!   r"   rZ   r   r   r   test_rank_args_missing   s     r
   re   zpct,expc                 C   s\   t ddddddddddg
dgd d}|dj| d}t |d d	gd
}t|| d S )Nab   r   rK   r   r   r(   rL   rO   r   r   r   r   rV   )r   rX   r!   r"   rZ   r   r   r   test_rank_resets_each_group  s    "rj   r,   int64Zint32Zuint64Zuint32float64Zfloat32upperc                 C   s   |r(| d   | dd   } | dd} tdgd dgd d}|d | |d< |d j| ksft|d	 }td
d
d
d
gdgd}t	|| d S )Nr   rh   ZUiZUIrf      rK   rL   r   g      @rO   )
rm   replacer   astyper,   AssertionErrorr   r   r   rV   )r,   rm   r!   r"   rZ   r   r   r   test_rank_avg_even_vals  s    rr   rW   rN   r`   r   barfooZbazc                 C   s   t dgd |d}|d  }|d}|j| |||d}| rht dgd dtjdtjd	gd}	nt dgd ddddd	gd}	|	d}
|
j| |||d}t|| d S )
Nrt      rK   rL   r   rd   r   r(   rh   )	r   Zisnar   r   anyr   rb   r   rV   )rW   rN   r`   r   r&   r!   maskgbresZdf2Zgb2altr   r   r   test_rank_object_dtype  s    
$
   r{   badrh   r(   c              	   C   sN   t dgd |d}d}tjt|d |dj| |||d W 5 Q R X d S )Nrt   ru   rK   z3na_option must be one of 'keep', 'top', or 'bottom')matchr   rd   )r   pytestZraises
ValueErrorr   r   )rW   rN   r`   r   r&   r!   msgr   r   r   test_rank_naoption_raises  s    
   r   c                  C   s   d} t dddgdtjdgd}|| jjdd}td	tjdgd
d}t|| || jdd}t d
d	tjdgi}t	|| d S )NAr   rh   rD   rC   r   BTr   rc   r   name)
r   r   rb   r   r   r   r   r   r    rV   )columnr!   r"   r#   r   r   r   test_rank_empty_group  s    r   z"input_key,input_value,output_valuec                 C   s>   t | |d}|djddd}t d|i}t|| d S )Nr   r   rH   TrM   r   r   ri   )Z	input_keyZinput_valueZoutput_valuer!   r"   r#   r   r   r   test_rank_zero_div  s    r   c                  C   s   t dddgttjjttjjttjjgttdddtgd} | d	 }t dddgtj
dtj
gd}t|| d S )	Nrh   r(   i  )grpint_coldatetimeliker   rD   rC   )r   r   )r   r   Ziinfork   rE   rF   r   r   r   r   NaNr   rV   )r!   r"   r#   r   r   r   test_rank_min_int"  s    r   use_nanc                 C   s   | r
t jnd}tddgddgd|gd|ggddgd}|dgd jdd	d
}| rntddt jt jgdd}ntddddgdd}t|| d S )Nr]   rh   r(   grouprL   rO   rH   Tr   rc   r   rI   rJ   )r   rb   r   r   r   r   r   r    )r   Z
fill_valuer!   r"   r#   r   r   r   .test_rank_pct_equal_values_on_group_transition8  s"    	r   c                  C   s   t tddgddgdtddgdd	gd
ddd} | jddd}|jdd}t | d jdd| d jddgdddgd}t|| d S )Nr]   rn   rh   r(   )Zcol1Zcol2ru   r*      r)   )Zcol3Zcol4)rf   rg   r   r   levelr   rf   rg   )r   keys)r   r   r   r   r   rV   )r!   rx   r"   r#   r   r   r   test_rank_multiindexQ  s     r   c                  C   s   t ddddgddddgd	d
ddgdddddgd} | jddd}|jdd}t| jd jdd| jd jddgdd}t|| |jdd}||rtd S Nrh   r]   ru   r   r(   rn   r*   r)   r\   r_   g      @g      @)r   rh   r(   rf   rg   )r   r   r   r   )	r   r   r   r   locr   rV   equalsrq   )r!   rx   ry   r#   rz   r   r   r   test_groupby_axis0_rank_axis1i  s    "
,r   c                  C   s~   t ddddgddddgd	d
ddgdddddgd} | jddd}|jdd}| ddg tj}|d |d< t|| d S r   )r   r   Zcummaxrp   r   rl   r   rV   )r!   rx   Zcmaxr#   r   r   r   test_groupby_axis0_cummax_axis1}  s    "
r   )"r   Znumpyr   r~   Zpandaspdr   r   r   r   Zpandas._testingZ_testingr   r$   markZparametrizeZ	TimestampZ	to_periodr[   infrb   ra   re   rj   rr   r{   r   r   r   r   r   r   r   r   r   r   r   r   <module>   s  "& "#.2""""""  > "
	  
 $	
	
