U
    sVcܒ                     @   s   d dl Zd dlZd dlmZmZ d dlZd dlmZm	Z	m
Z
mZ d dlmZ ejddgddggdd	 d
dd ZG dd dZG dd dZdS )    N)np_percentile_argnamenp_version_under1p21)	DataFrameIndexSeries	TimestamplinearZsinglenearesttablec                 C   s
   d | S )N-)joinx r   L/tmp/pip-unpacked-wheel-xj8nt62q/pandas/tests/frame/methods/test_quantile.py<lambda>       r   paramsidsc                 C   s   | j S )z.(interpolation, method) arguments for quantile)param)requestr   r   r   interp_method   s    r   c                
   @   s  e Zd Zejdejdddddddgee	e
ggd	d
 Zejdee	ejddge	ejddgde	ddgddgee	ddddgdde	dgddgg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d0d1 Zd2d3 Zd4d5 Zd6d7 Zd8d9 Zejd:d;d<d=d>gd?d@ ZejdAdBdddBgddggdCdD ZdEdF Z dGdH Z!dIdJ Z"dKdL Z#dMdN Z$dOdP Z%dQdR Z&dSdT Z'dUdV Z(dWdX Z)dYdZ Z*dS )[TestDataFrameQuantilenon_num_colz
2014-01-01   mperiodsfreqabcc           
   	   C   s   |\}}t dddgdddgd}||d< tddgd	d
gdd}|dkrT|tj}|dkrv|rv|jtjj	dd t
jtdd |jd||d}	W 5 Q R X t
|	| d S )N      r      ABC       @      @r'   r(         ?indexnamer	   r
   Axis name incorrectly set.reasonnumeric_onlymatchinterpolationmethod)r   r   astypenpint64node
add_markerpytestmarkxfailtmassert_produces_warningFutureWarningquantileassert_series_equal)
selfr   r   r   using_array_managerr7   r8   dfexpectedresultr   r   r   'test_numeric_only_default_false_warning   s"    z=TestDataFrameQuantile.test_numeric_only_default_false_warningzdf,expectedr#   r$   r%   )r   r#         ?      @r,   r/   g        N      ?r*   zSparse[float]dtypec                 C   s   |  }t|| d S N)rD   rA   rE   )rF   rH   rI   rJ   r   r   r   test_quantile_sparse;   s    z*TestDataFrameQuantile.test_quantile_sparsec           	         s  |\}}|  j ddd||d}t fdd jD  jdd}|dkrVt|| n8t|j|j |jt	j
j|d	d
 |j|jkst j ddd||d}t fdd jD  jdd}|dkrt|| n:t|j|j |jt	j
j|d	d
 |j|jkstd S )N皙?r   Taxisr3   r7   r8   c                    s   g | ]}t  | d qS )
   )r:   
percentile).0colrH   r   r   
<listcomp>]   s     z7TestDataFrameQuantile.test_quantile.<locals>.<listcomp>r-   r   z%Name set incorrectly for arraymanagerr1   ?r#   c                    s   g | ]}t  j| d qS )Z   )r:   rX   loc)rY   dater[   r   r   r\   q   s     )rD   r   columnsrA   rE   Zassert_index_equalr.   r<   r=   r>   r?   r@   r/   AssertionError)	rF   datetime_framer   rG   r   r7   r8   rJ   rI   r   r[   r   test_quantileT   s\              z#TestDataFrameQuantile.test_quantilec                 C   sJ   |\}}t g g djddd||d}t|d rBt|d sFtd S )Nr   yrT   r   TrU   r   rf   )r   rD   r:   isnanrb   )rF   r   r7   r8   qr   r   r   
test_empty   s        z TestDataFrameQuantile.test_emptyc           	   	   C   s   |\}}t ddddgddddgd}|jdd	||d
}tjtdd | d}W 5 Q R X |dkrx|d tj	}|dkr|r|j
tjjdd t|| d S )Nr'   r(   r#   r$   r   r%   )Zcol1Zcol2r,   Tr3   r7   r8   zSelect only validr4   r	   r
   r0   r1   )r   rD   rA   rB   rC   Zmedianrenamer9   r:   r;   r<   r=   r>   r?   r@   rE   )	rF   r   r   rG   r7   r8   rH   rsZxpr   r   r   test_non_numeric_exclusion   s"       z0TestDataFrameQuantile.test_non_numeric_exclusionc           	      C   s2  |\}}t dddgdddgddddgd}|jdd||d}td	d
dgdddgdd}|dkrl|tj}|dkr|r|jtj	j
dd t|| |jddgd||d}t d	dgd
dgddgdddgd}|dkr|jdd d f  d8  < |jdd d f  d7  < |tj}tj||dd d S )Nr#   r$   r   r%   r&   r.   r,   rV   r7   r8   rL         @rM   r-   r	   r
   r0   r1         ?g      ?g      @g      @r#   r$   r   r         ?T)Zcheck_index_type)r   rD   r   r9   r:   r;   r<   r=   r>   r?   r@   rA   rE   ilocassert_frame_equal	rF   r   r   rG   r7   r8   rH   rJ   rI   r   r   r   	test_axis   s4    $    
zTestDataFrameQuantile.test_axisc           	      C   s   |\}}t dddgdddgg}|jddd||d	}td
dgddgdd}|dkr^|tj}|dkr|r|jtj	j
dd t|| d S )Nr#   r$   r   r    r!   r%   r,   TrU   r+         @r   r-   r	   r
   r0   r1   )r   rD   r   r9   r:   r;   r<   r=   r>   r?   r@   rA   rE   rv   r   r   r   test_axis_numeric_only_true   s"        z1TestDataFrameQuantile.test_axis_numeric_only_truec                 C   sz   |\}}t jdddd}t|}t|}|jd||d}	tdgdd	d
}
|dkrj|rj|jtjj	dd t
|	|
 d S )N
2016-01-01r   
US/Pacificr   tzFrj   z2016-01-02 00:00:00r,   datetime64[ns, US/Pacific])r/   rQ   r
   r0   r1   )pd
date_ranger   r   rD   r<   r=   r>   r?   r@   rA   rE   )rF   r   r   rG   r7   r8   ZdtiserrH   rJ   rI   r   r   r   test_quantile_date_range   s&        z.TestDataFrameQuantile.test_quantile_date_rangec           
   	   C   s   |\}}t dddgdddgtjdddd	d
dgd}|jddd||d}tdddgdd}|dkrl|d8 }|dkr|r|jtjj	dd t
|| d}	tjt|	d |jdddd W 5 Q R X d S )Nr#   r$   r   r*   r+   rx   Z20130101r   foobarZbaz)r'   r(   r)   Dr,   TrU   rL   rp   rM   rN   r	   r
   r0   r1   z>'<' not supported between instances of 'Timestamp' and 'float'r4   FrV   r3   )r   r   r   rD   r   r<   r=   r>   r?   r@   rA   rE   raises	TypeError
rF   r   r   rG   r7   r8   rH   rJ   rI   msgr   r   r   test_quantile_axis_mixed   s4        z.TestDataFrameQuantile.test_quantile_axis_mixedc           
   	   C   s  |\}}|dkr*|r*|j tjjdd tdddgdddgddddgd	}|jd
d||d}tddgddgd
d}|dkr|t	j
}t|| |jd
d||d}|dkr|t	j
}t|| |jd
d||d}tdddgdddgd
d}|dkr
|t	j
}t|| |jd
d||d}t|| d}	tjt|	d |jdd||d W 5 Q R X d}	tjt|	d |jddd W 5 Q R X d S )Nr
   r0   r1   r#   r$   r   r%   r&   rn   r,   r   ro   r*   r+   r'   r(   r-   r	   r.   rL   rp   rM   ra   z*No axis named -1 for object type DataFramer4   rT   z.No axis named column for object type DataFramecolumnrV   )r<   r=   r>   r?   r@   r   rD   r   r9   r:   r;   rA   rE   r   
ValueErrorr   r   r   r   test_quantile_axis_parameter   sL    $   
   z2TestDataFrameQuantile.test_quantile_axis_parameterc                 C   sV  t dddgdddgddddgd}|jdddd	}tdddgdddgdd
}t|| tjtdddgdddggdfdditdi}t|dddgddd}t|| t dddgdddgddddgd}|jdddd	}tdddgdddgdd
}t|| tjtdddgdddggdfdditdi}t|dddgddd}t|| |jddgddd	}t ddgddgddgdddgd}t	|| t g g d}|jdddd	}t
|d rt
|d stt dddgdddgdddggdddgd}|jd dgd!d"}t d#d#d#gdddggd dgdddgd$}t	|| d S )%Nr#   r$   r   r%   r&   rn   r,   r	   rV   r7   r-   rV   r   r;   r.   r/   rQ   rO   r*   r+   rx   float64rq   lowerrr   re   rT   Zhigherr   rf   r    r!   r"   ra   rs   Zmidpoint)r7   rL   r.   ra   )r   rD   r   rA   rE   r:   rX   arrayr   ru   rg   rb   )rF   rH   rJ   rI   exprh   r   r   r   test_quantile_interpolation  sZ    $$ $*z1TestDataFrameQuantile.test_quantile_interpolationc                 C   s6   |}|j ddddd}|d t|d dks2td S )NrT   r   Tr   )rV   r3   r7   r'   rW   )rD   r:   rX   rb   )rF   rc   rH   rh   r   r   r   $test_quantile_interpolation_datetimeX  s    z:TestDataFrameQuantile.test_quantile_interpolation_datetimec                 C   sf   |}| d}|d t|d dks*t|j dddd}|d t|d dksVtt|| d S )NrT   r'   rW   r   r   r   )rD   r:   rX   rb   rA   rE   )rF   Z	int_framerH   rh   Zq1r   r   r   test_quantile_interpolation_int`  s    
z5TestDataFrameQuantile.test_quantile_interpolation_intc           	      C   s   |\}}t dddgdddgdddggdddgd}|jdd	g||d
}t dddgdddggdd	gdddgd}|dkr|tj}|dkr|r|jtjj	dd t
|| d S )Nr#   r$   r   r    r!   r"   r   rs   r,   r6   rL   r*   r   r	   r
   r0   r1   r   rD   r9   r:   r;   r<   r=   r>   r?   r@   rA   ru   rv   r   r   r   test_quantile_multim  s    *z)TestDataFrameQuantile.test_quantile_multic           	      C   s   |\}}t dddgdddgdddggdddgd}|jdd	gd||d
}t dddggd dd	gdddgd}|dkr|tj}|dkr|r|jtjj	dd t
|| d S )Nr#   r$   r   r    r!   r"   r   rs   r,   ro   rO   r*   r+   r   r   r	   r
   r0   r1   r   rv   r   r   r   test_quantile_multi_axis_1~  s(    *     z0TestDataFrameQuantile.test_quantile_multi_axis_1c                 C   s^   |\}}t g g djddgd||d}t tjtjgtjtjgdddgd}t|| d S )Nre   rT   r]   r   ro   rn   )r   rD   r:   nanrA   ru   )rF   r   r7   r8   rJ   rI   r   r   r   test_quantile_multi_empty  s        z/TestDataFrameQuantile.test_quantile_multi_emptyc                 C   s  t tddgddgd}|jddd}td	gd
gdd}t|| |jddd}ttdd	gdd
gdd}t|| |jdgdd}t tdd	ggdgdd
gd}t|| tddg|d< |ddg jdddd}ttdtdgddgdd}t|| |ddg jdgddd}t tdtdggdgddgd}t|| |ddg jddd}tg g t	j
dd}t|| |ddg jdgdd}t dgd}t|| d S )N20102011r      r    r!   r,   Tr3   rp   r!   r-   Fz2010-07-02 12:00:00r    r   2012r"   r#   r   z2011-07-02 12:00:00r.   rQ   r/   rn   )r   r   to_datetimerD   r   rA   rE   r   ru   r:   r   )rF   rH   rJ   rI   r   r   r   test_quantile_datetime  sN    
    z,TestDataFrameQuantile.test_quantile_datetimerQ   datetime64[ns]r~   timedelta64[ns]z	Period[D]c                 C   sz   |\}}t ddg|d}|jddd||d}tg g d|d}t|| |jdgdd||d}t dgd	}t|| d S )
Nr    r!   ra   rQ   r,   r#   FrU   r   rn   )r   rD   r   rA   rE   ru   )rF   rQ   r   r7   r8   rH   resrI   r   r   r   test_quantile_dt64_empty  s(        z.TestDataFrameQuantile.test_quantile_dt64_emptyinvalidr   c              	   C   s:   d}|\}}t jt|d |j|||d W 5 Q R X d S )Nz2percentiles should all be in the interval \[0, 1\]r4   r6   )r>   r   r   rD   )rF   r   rc   r   r   r7   r8   r   r   r   test_quantile_invalid  s    z+TestDataFrameQuantile.test_quantile_invalidc           	      C   s   |\}}|dkr*|r*|j tjjdd ttdtdtdgtdddtdddtdddgtd	td
tdgd}|j	dd||d}t
tdtdddtd
gddddgd}t|| |j	dgd||d}ttdtdddtd
ggdgdddgd}t|| d S )Nr
   r0   r1   
2011-01-01
2011-01-02
2011-01-03
US/Easternr}   1 days2 days3 days)r'   r(   r)   r,   Frj   r'   r(   r)   r/   r.   r   )r<   r=   r>   r?   r@   r   r   r   	TimedeltarD   r   rA   rE   ru   	rF   r   r   rG   r7   r8   rH   r   r   r   r   r   test_quantile_box  sf    


   
	   
z'TestDataFrameQuantile.test_quantile_boxc                 C   s  t tdtjtdtdgtdtdtjtdgtdddtjtdddtdddgtdddtdddtjtdddgtdtdtdtjgtjtdtdtdgd	td
d}|jddd}ttdtdtdddtdddtdtdgdtd
d}t	|| |jdgdd}t tdtdtdddtdddtdtdggdgtd
d}t
|| d S )Nr   r   r   r   r   r   r   r   )r'   r    r(   r!   r)   r"   ZAaBbCcr   r,   Fr   r   r   )r   r   r   NaTr   listrD   r   rA   rE   ru   )rF   rH   r   r   r   r   r   test_quantile_box_nat)  sz    





&*




z+TestDataFrameQuantile.test_quantile_box_natc           	      C   s  |\}}|dkr*|r*|j tjjdd ttddtddd}tj|j	d< |j
d||d	}td
|dkrtdnd
gddgdd}t|| |j
ddg||d	}td
dg|dkrddgnd
dgdddgd}t|| |j
dd||d}ttdddd}t|| |j
ddgd||d}ttddgd ddgd}|dkr\tj|j	d< t|| tj|d< |j
d||d	}td
tjgddgdd}t|| |j
ddg||d	}td
dgtjtjgdddgd}t|| d S )Nr
   r0   r1   r#   g      @r   )r   r#   r,   r6   r+   r   rp   r    r!   r-   rq   rx   g      
@rn   ro   rO   rN   r$   r	   )r#   r   )r<   r=   r>   r?   r@   r   r:   aranger   rt   rD   r   rA   rE   ru   r   r   r   r   test_quantile_nanu  sT         

"z'TestDataFrameQuantile.test_quantile_nanc           	      C   sT  |\}}|dkr*|r*|j tjjdd tdtjtjtjgi}|jdd||d}t	tjgdgdd}t
|| |jdgd||d}tdtjgidgd	}t
|| ttd
tdtdgtjtjtjgd}|jdd||d}t	tdtjgddgdd}t
|| |jdgd||d}ttdtjggdgddgd}t
|| d S )Nr
   r0   r1   r    r,   Frj   r-   rn   z
2012-01-01z
2012-01-02z
2012-01-03r   r!   r   )r<   r=   r>   r?   r@   r   r   r   rD   r   rA   rE   ru   r   r   r   r   r   test_quantile_nat  sd                  z'TestDataFrameQuantile.test_quantile_natc                 C   s   |\}}t ddgdd}|jd||d}ttjtjgddgdd}t|| |jdg||d}t tjtjggddgdgd}t|| |jdd	||d
}tg g ddd}t|| |jdgd	||d
}t g dgd}t|| d S )Nr    r!   r   r   r,   r6   r-   )ra   r.   r#   ro   r   )r   rD   r   r:   r   rA   rE   ru   rF   r   r7   r8   rH   r   r   r   r   r   "test_quantile_empty_no_rows_floats  s    z8TestDataFrameQuantile.test_quantile_empty_no_rows_floatsc                 C   sR   |\}}t ddgdd}|jd||d}ttjtjgddgdd}t|| d S )Nr    r!   r;   r   r,   r6   r-   )r   rD   r   r:   r   rA   rE   r   r   r   r    test_quantile_empty_no_rows_ints  s
    z6TestDataFrameQuantile.test_quantile_empty_no_rows_intsc                 C   s   |\}}t ddgdd}|jdd||d}ttjtjgddgddd}t|| |d jd	|d< |jdd||d}|	t
}t|| |d jd	|d< |jdd||d}|	|d j}t|| d S )
Nr    r!   r   r   r,   Frj   r   z
US/Central)r   rD   r   r   r   rA   rE   dtZtz_localizer9   objectrQ   r   r   r   r    test_quantile_empty_no_rows_dt64  sB       
      
   z6TestDataFrameQuantile.test_quantile_empty_no_rows_dt64c                 C   s   |\}}t tjddd}d|j_|jdd||d}tg g dtjd}d|j	_t
|| |jdgd||d}t g dgg d	}d|j_t
|| d S )
Nz1/1/18r   r   zcaptain tightpantsr,   Trj   r   r   )r   r   r   ra   r/   rD   r   r:   r   r.   rA   rE   ru   )rF   r   r7   r8   rH   rJ   rI   r   r   r   test_quantile_empty_no_columns  s*          z4TestDataFrameQuantile.test_quantile_empty_no_columnsc                 C   s   |\}}t tjdddddgd}|d d |d< |d }|sTt|jjdksTt|jd	||d
 d|j	d< |j
d |d d kstd S )Nr%   r   r'   r(   r)   r   r$   r   Frj   c   r   )r   r   )r   r:   randomZrandnlenZ_mgrblocksrb   rD   valuesrt   )rF   rG   r   r7   r8   rH   r   r   r   r   test_quantile_item_cache"  s    
z.TestDataFrameQuantile.test_quantile_item_cachec              	   C   s4   t jtdd ttdjddd W 5 Q R X d S )NzInvalid method: foor4   r#   r,   r   )r8   r>   r   r   r   rangerD   rF   r   r   r   test_invalid_method0  s    z)TestDataFrameQuantile.test_invalid_methodc              	   C   s6   t jtdd ttdjdddd W 5 Q R X d S )NzInvalid interpolation: foor4   r#   r,   r
   r   )r8   r7   r   r   r   r   r    test_table_invalid_interpolation4  s    z6TestDataFrameQuantile.test_table_invalid_interpolation)+__name__
__module____qualname__r>   r?   parametrizer   r   r   r   r   rK   ZarraysZSparseArrayrS   rd   ri   rm   rw   ry   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r      sx   


-*9
1	

9L1/r   c                
   @   s  e Zd Zejejeje	dej
jdddejddddejddd	d
ejdddejedddejedddgdd ddd Zejdd Zdd Zdd Zdd Zej
ddd Zd d! Zej
d"d#g g d$gd%g g d$gd#ejejgd&d'gd(gd%ejejgd&d'gd(ggd)d* Zej
d+d,g g d$d,gd,ejejgd&d'gd(d,ggd-d. Zej
d/ejejge	d0d$gg g d(ggd1d2 Zd3S )4TestQuantileExtensionDtyperW   z#raises when trying to add Intervalsr1   )Zmarksrz   	   r   r   r{   r|   z1 Dayr   Int64rP   Float64c                 C   s
   t | jS rR   )strrQ   r   r   r   r   r   I  r   z#TestQuantileExtensionDtype.<lambda>r   c                 C   s   |j }d|_|S Nr'   )r   r/   )rF   r   idxr   r   r   r.   =  s    z TestQuantileExtensionDtype.indexc                 C   s(   ||  }|tkrd|_ndg|_|S r   )copyr   r/   ra   )rF   r.   Zframe_or_seriesobjr   r   r   r   Q  s
    zTestQuantileExtensionDtype.objc                 C   s(   t |tr||}n|j|dd}|S )NFr   )
isinstancer   rD   )rF   r   qsrJ   r   r   r   compute_quantile\  s    
z+TestQuantileExtensionDtype.compute_quantilec                 C   s   t jt|t jd}t j| |j| }dddg}| ||}trn|j	dkrnd}t
jj|td}|j| |j	}	|j	dkrd	}	t|d
 |d |d g|	|dd}
t||
}
t||
 d S )NrP   r,   r   r#   r   Cfailed on Numpy 1.20.3; TypeError: data type 'Int64' not understoodr2   r   r   r   r%   r   r'   rQ   r.   r/   )r:   r   r   intpr   shufflert   r   r   rQ   r>   r?   r@   r   r<   r=   r   typerA   assert_equalrF   r   r   r.   indexerr   rJ   r   r?   Z	exp_dtyperI   r   r   r   test_quantile_eac  s(    


   z+TestQuantileExtensionDtype.test_quantile_eac                 C   s   |j |jd< |j |jd< tjt|tjd}tj| |j| }dddg}| ||}t	|d |d |d g|j
|dd	}t||}t|| d S )
Nr   r   rP   r,   r#   r%   r'   r   )	_na_valuert   r:   r   r   r   r   r   r   r   rQ   r   rA   r   )rF   r   r.   r   r   rJ   rI   r   r   r   test_quantile_ea_with_na  s    

   z3TestQuantileExtensionDtype.test_quantile_ea_with_naz-ignore:Using .astype to convert:FutureWarningc                 C   s   |j |jd d < ||j}t|j|jks2ttjt	|tj
d}tj| |j| }dddg}| ||}|jdddgd|j d}t||dd	}t||}t|| d S )
NrP   r,   r   r#   r   T)Z
allow_fillZ
fill_valuer'   r-   )r   rt   r9   rQ   r:   allZdtypesrb   r   r   r   r   r   r   Ztaker   r   rA   r   )rF   r   r   r.   r   r   rJ   rI   r   r   r   test_quantile_ea_all_na  s    

z2TestQuantileExtensionDtype.test_quantile_ea_all_nac                 C   s   t jt|t jd}t j| |j| }d}| ||}trh|j	dkrhd}t
jj|td}|j| |j	}	|j	dkr|d}	td|d	 i|	dd
}
t|tr|
d }
||
kstnt||
 d S )NrP   r,   r   r   r   r   r   r'   r%   )rQ   r/   )r:   r   r   r   r   r   rt   r   r   rQ   r>   r?   r@   r   r<   r=   r   r   rb   rA   rE   r   r   r   r   test_quantile_ea_scalar  s"    


z2TestQuantileExtensionDtype.test_quantile_ea_scalarz*dtype, expected_data, expected_index, axisr   r#   r;   r    r!   r   c                 C   sB   t ddg|d}|jd|d}t|dt|dd}t|| d S )Nr    r!   r   r,   r   r   r/   r.   rQ   r   rD   r   r   rA   rE   )rF   rQ   expected_dataexpected_indexrV   rH   rJ   rI   r   r   r   test_empty_numeric  s       z-TestQuantileExtensionDtype.test_empty_numericz:dtype, expected_data, expected_index, axis, expected_dtyper   c           	      C   sD   t ddg|d}|jd|dd}t|dt||d}t|| d S )Nr    r!   r   r,   Fr   r   r   )	rF   rQ   r   r   rV   Zexpected_dtyperH   rJ   rI   r   r   r   test_empty_datelike  s       z.TestQuantileExtensionDtype.test_empty_datelikez#expected_data, expected_index, axisr$   c                 C   sf   t tddgddgtddgd}|ddg jd	|d
d}t|d	t|tjd}t	|| d S )Nr   r   r   r   r   )r    r!   r"   r    r"   r,   Tr   r   )
r   r   r   rD   r   r   r:   r   rA   rE   )rF   r   r   rV   rH   rJ   rI   r   r   r   test_datelike_numeric_only  s    	   z5TestQuantileExtensionDtype.test_datelike_numeric_onlyN) r   r   r   r>   fixturer   r   ZIntervalIndexZfrom_breaksr   r?   r@   Zperiod_ranger   Ztimedelta_ranger   r:   r   r.   r   r   r   r   filterwarningsr   r   r   r   r   r   r   r   r   r   r   r   r   9  sZ   






	
	
r   )Znumpyr:   r>   Zpandas.compat.numpyr   r   Zpandasr   r   r   r   r   Zpandas._testingZ_testingrA   r   r   r   r   r   r   r   r   <module>   s$    
      $