B
    `#l                 @   s   d dl Zd dlm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lmZ d dlmZ ejdd Zejdd	 Zejd
d Zejdd Zejdd Zejdd ZG dd dZG dd dZdS )    N)randn)	DataFrameIndex
MultiIndexSeries)concat)mergec           
   C   sP   ddddddddddg
} ddddddddddg
}t jt| }t| ||d	S )
z=left dataframe (not multi-indexed) for multi-index join testsbarfoobazquxsnaptwoonethree)key1key2data)nprandomr   lenr   )r   r   r    r   T/home/dcms/DCMS/lib/python3.7/site-packages/pandas/tests/reshape/merge/test_multi.pyleft   s    r   c              C   sn   t ddddgdddggdddd	d	d
d
dddg
dd	d
dd	d	d
dd	d
g
gddgd} ttjdd| dddgdS )z:right dataframe (multi-indexed) for multi-index join testsr
   r	   r   r   r   r   r   r            r   r   )levelscodesnames
   j_onej_twoj_three)indexcolumns)r   r   r   r   r   )r$   r   r   r   right   s    .r&   c               C   sd   t tdddddgdddddgdddddgddd	ddgd
ddddgddddddgdddddgS )NABCAMIPOPZhbwZnhbZhboi  i?  i	  i  i\  )OriginDestinationPeriodTripPurpTripsr-   r.   r/   r0   r1   )r%   )r   dict	set_indexr   r   r   r   
left_multi%   s    r4   c               C   sx   t tdddddddgdddddddgdddddddgd	d
dd
d	d
d	gdddddddgddddddgdddddgS )Nr'   r(   r)   EFr*   r+   r,   abcd   P   Z   K   #   7   )r-   r.   r/   LinkTypeDistancer-   r.   r/   r@   rA   )r%   )r   r2   r3   r   r   r   r   right_multi3   s    rB   c               C   s
   dddgS )Nr-   r.   r/   r   r   r   r   r   on_cols_multiA   s    rC   c               C   s   dddddgS )Nr-   r.   r/   r0   r@   r   r   r   r   r   idx_cols_multiF   s    rD   c               @   s   e Zd Zdd Zdd Zejdddgdd	 Zejddd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ejddejeegdd Zdd Zd d! ZdS )"TestMergeMultic             C   s   t ddddgdddggdddd	d	d
d
dddg
dd	d
dd	d	d
dd	d
g
gddgd| _ttjdd| jdddgd| _ddddddddddg
}ddddddddddg
}tjt|}t|||d| _d S )Nr
   r	   r   r   r   r   r   r   r   r   r   firstsecond)r   r   r   r    r!   r"   r#   )r$   r%   r   )r   r   r   )	r   r$   r   r   r   r   Zto_joinr   r   )selfr   r   r   r   r   r   setup_methodL   s*    .zTestMergeMulti.setup_methodc             C   s   ddg}|j |||djdd}tj|| ||d}t|| |j |||ddjdd}tj|| ||dd}t|| d S )Nr   r   )onhowT)drop)rJ   rK   sort)joinreset_indexpdr   tmassert_frame_equal)rH   r   r&   	join_typeon_colsresultexpectedr   r   r   test_merge_on_multikeyj   s    z%TestMergeMulti.test_merge_on_multikeyrM   FTc          
      sh  dddgdd   fdd}t ttttdtd	d
 }ttj|dddgd}|	d
dtj
ddt| tjt|}|j|  } ||d<  | |d< |jdd |||| tj|jd
d ddf< tj|jdd ddf< tj|jdd ddf<  ||d< tjt|}|j|d df } | |d< |jdd |||| d S )N1st2nd3rdc                sB   dd   fdd}|| d || d d  | d  dd	  S )
Nc             S   s   | | krdS t | S )Nr   )ord)r7   r   r   r   <lambda>       zNTestMergeMulti.test_left_join_multi_index.<locals>.bind_cols.<locals>.<lambda>c                s   |   td S )Nr7   )mapr[   )ts)iordr   r   r\      r]   rX   rZ   g      Y@rY   r   g     @)fillna)dffr   )r`   r   	bind_cols   s    z<TestMergeMulti.test_left_join_multi_index.<locals>.bind_colsc                s   | j |d|d}t| t|d k s*t|d   r>t|d   rRttj|d |d  dd  |jd d d df }tj|d |dd |jd kst|rt	||j
d	d
 t| | |dd}tt||_t	|| d S )Nr   )rJ   rK   rM   r   4th5thF)Zcheck_names	mergesort)kind)rJ   rM   rK   )rN   r   AssertionErrorZisnaanyrQ   Zassert_series_equalilocnamerR   sort_valuesr   rO   r   aranger$   )r   r&   rM   resrU   out)rd   icolsr   r   run_asserts   s    z>TestMergeMulti.test_left_join_multi_index.<locals>.run_assertsr7   zr   )i  r   )r%   r   i  re   rf   T)inplace   r   %   r   +   )listr^   chrr   ro   r[   r   r   choiceinsertrandintr   Zpermutationrl   copyr3   nanloc)rH   r   r&   rM   rs   lcir   )rd   rr   r   test_left_join_multi_index|   s*    
"z)TestMergeMulti.test_left_join_multi_indexc             C   sJ   ddg}|j ||dd|d}|j ||dd|d}||j }t|| d S )Nr   r   Tr   )Zleft_onZright_indexrK   rM   r&   )Zright_onZ
left_indexrK   rM   )r   r%   rQ   rR   )rH   r   r&   rM   rT   Zmerged_left_rightZmerge_right_leftr   r   r   test_merge_right_vs_left   s    
z'TestMergeMulti.test_merge_right_vs_leftc          	   C   s   t jtdt jjddgdddggddgd	d
d}t tdddddddgdd}t j|| ddgd}t tdddddddgdgd tdd}t	
|| d S )N   r'   r(   r   r   r   lev1lev2)r   Amount)rm   ZAAABBBr   )r   r   col)rJ   )r   r   r   r   )rP   r   ranger   from_productr   rz   r   rO   rQ   rR   )rH   srb   rU   rV   r   r   r   .test_merge_multiple_cols_with_mixed_cols_index   s     z=TestMergeMulti.test_merge_multiple_cols_with_mixed_cols_indexc             C   s|   t dd}t|d}|d d d }t||tjdd}t|d d d |d d d tjdd}t||dd	 d S )
Nr    i'  r   ry   i N  )r   r   Zvalue1)r   r   Zvalue2outer)rK   )rQ   Zrands_arrayr   Ztiler   r   r   r   )rH   r   r   rb   Zdf2r   r   r    test_compress_group_combinations   s    *z/TestMergeMulti.test_compress_group_combinationsc             C   s  ddg}t dddgd ddgd	 tjtd
tjdd}tddg}t dddgi|d}|j||d}| }tj	|d< d|j
|jdk|jdk@ df< d|j
|jdk|jdk@ df< t|| |j|ddd |j||dd}t|| t dddgd ddgd	 tjdddgd tjdtjtd
tjdd}tddg}t dddgi|d}|j||d}| }tj	|d< d|j
|jdk|jdk@ df< d|j
|jdk|jdk@ df< t|| |j|dd}|j||dd}t|| d S )Nk1k2r   r   r      r
   r	         )Zdtype)r   r   v)r   r	   )r   r
   v2      )r$   )rJ   rh   T)ri   ru   )rJ   rM   )r   r   Zk3r   )ri   )r   r   arrayro   int64r   from_tuplesrN   r   r   r   r   r   rQ   rR   rn   float32int32)rH   rT   r   r$   r&   rU   rV   r   r   r   #test_left_join_index_preserve_order   s@    



z2TestMergeMulti.test_left_join_index_preserve_orderc             C   sx  t ddddgddddgddd	d
gddddgddddgddd	dgddddgddddgddddgddddgg
ddddgdddddddd d!d"g
d#}t ddddgddddgdddd"gddd	dgddd	dgdddd gddddgdddd!gdddd$gdddd%gddddgddddgdddd&gddd	d'ggd(d)d*d+gd,d(d)d*g}|j|dddgd-d.}t dddddgddddd!gddddtjgddd	d
d'gdddddgddddd&gddddtjgddd	ddgddd	ddgdddddgddddd"gdddddgdddddgddddd ggddddd+gdddddddddddd d!d"gd#}t|| |j|dddgd-d/d0}|jdddgd1d2}t|| d S )3NXYr)   r7   WeVQr'   hRDr   r8   r9   r(   rc   gjdZcolaZcolbZcolctagr   r   r   r   r   r      r   	   r   )r%   r$   iiry   Zcol1Zcol2Zcol3val)r%   r   )rJ   rK   T)rJ   rK   rM   rh   )ri   )r   r3   rN   r   r   rQ   rR   rn   )rH   r   r&   rU   rV   r   r   r   +test_left_join_index_multi_match_multiindex   sh    






















$z:TestMergeMulti.test_left_join_index_multi_match_multiindexc             C   s\  t ddgddgddgddggdd	gddddgd
}t ddgddgddgddgddgddgddgddggddgdd}|j|ddd}t dddgdddgdddgdddgddtjgdddgdddgddtjggdd	dgddddddddgd
}t|| |j|dddd}|jddd}t|| t||	 ddd}t
t||_t|| d S )Nr9   r   r8   r   r7   r   r   r   r   )r%   r$   r   wxr   yrt   rr   qr   char)r%   r   )rJ   rK   T)rJ   rK   rM   rh   )ri   )rK   rJ   )r   r3   rN   r   r   rQ   rR   rn   r   rO   ro   r   r$   )rH   r   r&   rU   rV   Z	expected2r   r   r    test_left_join_index_multi_matche  sD    
z/TestMergeMulti.test_left_join_index_multi_matchc             C   s   t tdtdtdtdtdddddddgd	}t d
dtjtjtjgddtjtjtjgd}t||ddd}|jdgdd}||}t	|| d S )NZabcder   )idv1r   dummyv3r   r   r   r   r   )r%   r7   r8   gX9v?gʡE@)r   Zsv3r   )rJ   rK   r   )axis)
r   rz   r   r   r   r   rL   rN   rQ   rR   )rH   r   r&   rU   rdfrV   r   r   r   test_left_merge_na_buglet  s    

z(TestMergeMulti.test_left_merge_na_bugletc             C   s   dddgdddgdddgddt jgdddgdd	dgdd	t jgd
d	dgdd	dgg	}t|dddgd}ddt jgddt jgddt jgd
dt jgd
dt jgdd	t jgg}t|dddgd}|j|dd}|dj|ddd}|dt j}t|| d S )Ni  r'   g      ?r(   i  i  i  g      @r)   i  g      @yearZpanelr   )r%   r   )rK   i)r   r   r   r   ra   replacerQ   rR   )rH   r   frameZ
other_dataotherrU   rV   r   r   r   test_merge_na_keys  s*    






z!TestMergeMulti.test_merge_na_keysklassNc             C   s   t dddgdddgdgd}t|j|_|jj}|d k	rB||}t dddgd	d
dgd}|j|d|gdd}t|| t d	d
dgdddgdddgd}|j||jjgdd}t|| d S )Nr   r   r   z
2016-01-01z
2017-01-01z
2018-01-01r7   )r%   i  i  i  )r7   key_1inner)rJ   rK   )key_0a_xa_y)r   rP   to_datetimer$   r   r   rQ   rR   )rH   r   rb   	on_vectorrV   rU   r   r   r   test_merge_datetime_index  s     z(TestMergeMulti.test_merge_datetime_indexc             C   s.  t tdddgdddgdddgdd	d
dgdd	}t tdddddddgddddddtjgddddddtjgdddddddgdd	ddd gdd	dg}|j|d!d"}t tddddddgdddd#d#d#gddddddgddddddgddddddgddddddgd$d	dgjd
ddd gd}t|| t	|
 |
 d	gd!d%d	dg}t|| |j|d&d"}t|t tdgd'tjdtjfgd	dgd(d)gdd*d+j|jd}t|| d,|j_tjtd-d. |j|d!d" W d Q R X | }|jd	d,g tjtd/d. |j|d!d" W d Q R X d S )0Nr   r   r   r   gffffAgPAi^ )household_idmalewealthr   r   r   )r%   r   nl0000301109Znl0000289783gb00b03mlx29lu0197800237nl0000289965zABN AmroZRobecozRoyal Dutch ShellzAAB Eastern Europe Equity FundzPostbank BioTech Fondsg      ?g?g333333?g333333?g      ?)r   asset_idrm   sharer   rm   r   r   )rK   g    xA)r   r   rm   r   r   r   )rJ   rK   r   )r   )r   )r$   T)r   rM   r
   z+cannot join with no overlapping index names)matchz'columns overlap but no suffix specified)r   r2   r3   r   r   rN   reindexrQ   rR   r   rO   r   r   r   r%   r$   rm   pytestZraises
ValueErrorr   Z	set_names)rH   	householdZ	portfoliorU   rV   Z
portfolio2r   r   r   test_join_multi_levels  s    




z%TestMergeMulti.test_join_multi_levelsc             C   s  t tdddddddgddddddtjgd	d
ddddd	gddddgdddg}t tdddddgdddddgdddddgdddg}t tddddddddgddddddddgddddddddgddddddddgddddddddgddddgjdd gd}t| | dgd!d"dddg}t	|| t tddddddddddddgdddddddddddd gd d ddddddddd d gd	d
dddddddddd	gd d ddddddddd d gddddgjdd gd}t| | dgd#d"dddg}t	|| d S )$Nr   r   r   r   r   r   r   r   g      ?g?g333333?g333333?g      ?)r   r   r   r   r   r   )r%                  gS緖?gDg >?g
?gHRE?)r   t
log_returnr   )r   r   r   r   r   r   r   )rJ   rK   r   )
r   r2   r   r   r3   r   r   rO   rQ   rR   )rH   r   r   rV   rU   r   r   r   test_join_multi_levels2T  s    

z&TestMergeMulti.test_join_multi_levels2)__name__
__module____qualname__rI   rW   r   markparametrizer   r   r   r   r   r   r   r   r   r   asarrayr   r   r   r   r   r   r   r   r   rE   K   s   77E2  mrE   c               @   sL   e Zd Zdd Zdd Zejddej	e
egdd Zd	d
 Zdd ZdS )TestJoinMultiMultic             C   sF   t j| | ||d| }|j||d }t|| d S )N)rK   rJ   )rK   )rP   r   rO   r3   
sort_indexrN   rQ   rR   )rH   r4   rB   rS   rC   rD   rV   rU   r   r   r   test_join_multi_multi  s    
z(TestJoinMultiMulti.test_join_multi_multic             C   sb   |j |jd}|j |jd}tj| | ||d| }|j||d }t	|| d S )N)r%   )rK   rJ   )rK   )
rL   r%   rP   r   rO   r3   r   rN   rQ   rR   )rH   r4   rB   rS   rC   rD   rV   rU   r   r   r   test_join_multi_empty_frames  s    
z/TestJoinMultiMulti.test_join_multi_empty_framesboxNc             C   s   t dddgdddgdgd}t|j|_|jj}|d k	rB||}t dddgd	d
dgd}|j|d|gdd}t|| t d	d
dgdddgdddgd}|j||jjgdd}t|| d S )Nr   r   r   z
2016-01-01z
2017-01-01z
2018-01-01r7   )r%   i  i  i  )r7   r   r   )rJ   rK   )r   r   r   )r   rP   r   r$   r   r   rQ   rR   )rH   r   rb   r   rV   rU   r   r   r   r     s     z,TestJoinMultiMulti.test_merge_datetime_indexc             C   s   t jjdddgddgd}t jddd	gd
ddgd|d}t jjddddgddgd}t jddddgddddgd|d}||}t j| | dgdddddg}t	|| d S )N)K0ZX0)r   ZX1)K1ZX2keyr   )r   ZA0A1A2B0ZB1ZB2)r'   r(   )r$   )r   ZY0)r   ZY1)K2ZY2)r   ZY3r   ZC0ZC1ZC2ZC3ZD0ZD1ZD2ZD3)r)   r   r   )rJ   rK   )
rP   r   r   r   rN   r   rO   r3   rQ   rR   )rH   Z
index_leftr   Zindex_rightr&   rU   rV   r   r   r   test_single_common_level  s    
z+TestJoinMultiMulti.test_single_common_levelc             C   s   t jjddgddggddgd}t jjdd	d
gddgd}t j|dddddgid}t j|ddddgid}||}t j|ddddgdddtjgdd}t	|| d S )Nr   r   r   r   r7   r8   )r   )r   r   )r   r   )r   r   r   r          (   )r$   r   r   r
   r	   Zfing)r   r   )
rP   r   r   r   r   rN   r   r   rQ   rR   )rH   Zmidx1Zmidx3r   r&   rU   rV   r   r   r   test_join_multi_wrong_order8  s     
 z.TestJoinMultiMulti.test_join_multi_wrong_order)r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r     s
    r   )numpyr   Znumpy.randomr   r   ZpandasrP   r   r   r   r   Zpandas._testingZ_testingrQ   Zpandas.core.reshape.concatr   Zpandas.core.reshape.merger   Zfixturer   r&   r4   rB   rC   rD   rE   r   r   r   r   r   <module>   s(        