U
    NZfLt                     @   s   d dl Zd dlZd dlZd dl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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)	DataFrameIndex
MultiIndex
RangeIndexSeries	Timestamp)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fooZbazZquxZsnapZtwoZoneZthreekey1key2data)nprandomrandnlenr   r    r   I/tmp/pip-unpacked-wheel-nbcvw55c/pandas/tests/reshape/merge/test_multi.pyleft   s    r   c                 C   s    | }ddg|j _dddg|_|S )z:right dataframe (multi-indexed) for multi-index join testsr   r   Zj_oneZj_twoZj_three)indexnamescolumns)Z multiindex_dataframe_random_datadfr   r   r   right   s    r   c                
   C   sb   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$   r%   r&   r   r   	set_indexr   r   r   r   
left_multi'   s    	
r*   c                   C   sv   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$   r6   r7   r'   r(   r   r   r   r   right_multi5   s    	
r8   c                   C   s
   dddgS )Nr"   r#   r$   r   r   r   r   r   on_cols_multiC   s    r9   c                   C   s   dddddgS )Nr"   r#   r$   r%   r6   r   r   r   r   r   idx_cols_multiH   s    r:   c                   @   s  e Z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ejdddgdd  Zejd!d" Zejd#d$ Zejd%d& Zd'd( Zd)d* Zd+d, Zd-d. Zd/d0 ZdS )1TestMergeMultic                 C   s   ddg}|j |||djdd}t|| ||d}t|| |j |||ddjdd}t|| ||dd}t|| d S )Nr   r   onhowT)dropr=   r>   sort)joinreset_indexr	   tmassert_frame_equal)selfr   r   	join_typeon_colsresultexpectedr   r   r   test_merge_on_multikeyN   s        z%TestMergeMulti.test_merge_on_multikeyrA   FTc              
      sn  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|d 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)r-   r   r   r   <lambda>e       zNTestMergeMulti.test_left_join_multi_index.<locals>.bind_cols.<locals>.<lambda>c                    s   |   td S )Nr-   )maprO   )tsZiordr   r   rP   f   rQ   rL   rN   g      Y@rM   r   g     @)fillna)r   fr   rT   r   	bind_colsd   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   r@      4th5thF)Zcheck_names	mergesortkind)r=   rA   r>   )rB   r   AssertionErrorZisnaanyrD   Zassert_series_equalilocnamerE   sort_valuesr	   rC   r   r   )r   r   rA   resrI   outrW   Zicolsr   r   run_assertsi   s    z>TestMergeMulti.test_left_join_multi_index.<locals>.run_assertsr-   zrX   )i     r'   r   i  floatrY   rZ   T)inplace   ri   %      +   )listrR   chrr   arangerO   r   r   choiceinsertrandintr   ZastypeZpermutationra   copyr)   nanloc)rF   rA   rg   lcr   ir   r   rf   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_indexr>   rA   r   )Zright_onZ
left_indexr>   rA   )r	   r   rD   rE   )rF   r   r   rA   rH   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 tdt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|| 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   rX   ri   rn   lev1lev2r   Amount)rb   ZAAABBBr   )r   r   colr=   )r   r   r   r   )
r   ranger   from_productr   rq   r	   rC   rD   rE   )rF   sr   rI   rJ   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 )
N
   i'  ri   rp   i N  )r   r   Zvalue1)r   r   Zvalue2outerr>   )rD   Zrands_arrayr   Ztiler   r   r   r	   )rF   r   r   r   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   rX   ri      r   r
         Zdtype)r   r   v)ri   r
   )rX   r   v2      r   r   r\   T)r^   rk   )r=   rA   )r   r   Zk3r   r]   )r   r   arrayrs   Zint64r   from_tuplesrB   rw   rx   ry   r   r   rD   rE   rc   Zfloat32int32)rF   rH   r   r   r   rI   rJ   r   r   r   #test_left_join_index_preserve_order   sH    


	
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   r-   WeVQr   hRDr{   r.   r/   r   rV   gjdZcolaZcolbZcolctagrn   ri   r   rX   r   r~      r   	   r   r   r   iirp   Zcol1Zcol2Zcol3valr'   r   r<   Tr@   r\   r]   )r   r)   rB   r   rx   rD   rE   rc   )rF   r   r   rI   rJ   r   r   r   +test_left_join_index_multi_match_multiindex  sv    

























z:TestMergeMulti.test_left_join_index_multi_match_multiindexc                 C   sZ  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 )Nr/   r   r.   rX   r-   ri   rn   r   r   r   r   wxr   yrh   rr   qr   charr'   r   r<   Tr@   r\   r]   r>   r=   )r   r)   rB   r   rx   rD   rE   rc   r	   rC   r   r   r   )rF   r   r   rI   rJ   Z	expected2r   r   r    test_left_join_index_multi_matchF  sP    




z/TestMergeMulti.test_left_join_index_multi_matchc                 C   s   t tdtjdtjdtdtj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'   r-   r.   gX9v?gʡE@)r   Zsv3r   r<   rX   )axis)r   rq   r   r   r   rx   r	   r?   rB   rD   rE   )rF   r   r   rI   rdfrJ   r   r   r   test_left_merge_na_bugletx  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         ?r   i  i  i  g      @r   i  g      @yearpanelr   r'   r   r   i)r   rx   r   r	   rU   replacerD   rE   )rF   r   frameZ
other_dataotherrI   rJ   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jd	d
dgtjd}t dddg|d}|j|d|gdd}t	
|| t |dddgdddgd}|j||jjgdd}t	
|| d S NrX   ri   rn   z
2016-01-01z
2017-01-01z
2018-01-01r-   r'   i  i  i  r   )r-   Zkey_1innerr<   )Zkey_0Za_xZa_yr   pdZto_datetimer   r   r   r   r   r	   rD   rE   )rF   r   r   	on_vector	exp_yearsrJ   rI   r   r   r   test_merge_datetime_index  s       z(TestMergeMulti.test_merge_datetime_index
merge_typer   r   c                 C   s  t dddgitjtddgtddggddgd	d
}t tjg ddgd	dgd}tjtddgtddggddgd	}|dkrt ddgd d gd|d
}|j|dddgd}|j|dd}n<t d d gddgd|d
}|j|dddgd}|j|dd}t|| t|| d S )Nr   r   z
1950-01-01r   z
1950-01-02r   dater   r   )r   r   state)r   r   r   )r   r   r   r   )r   r   r   )r   r   r   r   r	   rB   rD   rE   )rF   r   r   r   Zexpected_indexrJ   Zresults_mergeZresults_joinr   r   r   (test_merge_datetime_multi_index_empty_df  sH     
 z7TestMergeMulti.test_merge_datetime_multi_index_empty_dfc                 C   s6   t dddgdddgdddgdd	d
dgdd	}|S )NrX   ri   rn   r   ffffAPAi^ )household_idmalewealthr   r   r   r'   r(   )rF   	householdr   r   r   r     s    zTestMergeMulti.householdc                 C   sh   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}|S )NrX   ri   rn   r   nl0000301109nl0000289783gb00b03mlx29lu0197800237nl0000289965ABN AmroRobecoRoyal Dutch ShellAAB Eastern Europe Equity FundPostbank BioTech Fonds      ?皙?333333?333333?      ?)r   asset_idrb   sharer   r   rb   r   r'   )r   r   rx   r)   )rF   	portfolior   r   r   r     s2    
	
zTestMergeMulti.portfolioc                 C   s|   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}|S )Nr   rX   r   r   g    xAr   r   r   r   r   r   r   r   r   r   ri   rn   r   r   r   r   r   )r   r   rb   r   r   r   r   r   r   r   rb   r   r'   )r   r)   reindex)rF   rJ   r   r   r   rJ     sB    	 
#zTestMergeMulti.expectedc                 C   s.   |  }|  }|j|dd}t|| d S )Nr   r   )rw   rB   rD   rE   rF   r   r   rJ   rI   r   r   r   test_join_multi_levelsF  s    z%TestMergeMulti.test_join_multi_levelsc                 C   sD   |  }|  }t| | dgddddg}t|| d S )Nr   r   r<   r   )rw   r	   rC   r)   rD   rE   r   r   r   r   (test_join_multi_levels_merge_equivalenceO  s    z7TestMergeMulti.test_join_multi_levels_merge_equivalencec                 C   sn   |  }|  }|j|dd}t|tddgitjdtjfgddgdd	gd
ddj|j	d}t
|| d S )Nr   r   r   r   r   r   r   r   r   r   T)r   rA   r'   )rw   rB   r   r   r   r   r   rx   r   r   rD   rE   r   r   r   r   test_join_multi_levels_outer\  s&    
 z+TestMergeMulti.test_join_multi_levels_outerc              	   C   s   |  }|  }d|j_tjtdd |j|dd W 5 Q R X |  }|jddg tjtdd |j|dd W 5 Q R X d S )Nr   z+cannot join with no overlapping index names)matchr   r   r   z'columns overlap but no suffix specified)rw   r   rb   pytestZraises
ValueErrorrB   Z	set_names)rF   r   r   Z
portfolio2r   r   r   test_join_multi_levels_invalidr  s     z-TestMergeMulti.test_join_multi_levels_invalidc                 C   s  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 dddddgdddddgdddddgdddg}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 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 )$NrX   ri   rn   r   r   r   r   r   r   r   r   r   r   )r   r   r   r   r   r   r'                  gS緖?gDg >?g
?gHRE?)r   t
log_returnr   )r   r   r   r   r   r   r   r<   r   )	r   r   rx   r)   r   r	   rC   rD   rE   )rF   r   r   rJ   rI   r   r   r   test_join_multi_levels2  s   	
!=>Bz&TestMergeMulti.test_join_multi_levels2)__name__
__module____qualname__rK   r   markparametrizer|   r}   r   r   r   r   r   r   r   r   asarrayr   r   r   r   fixturer   r   rJ   r   r   r   r   r   r   r   r   r   r;   M   s6   
7
6E2 

.


&	r;   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   sD   t | | ||d| }|j||d }t|| d S )Nr   r   )r	   rC   r)   
sort_indexrB   rD   rE   rF   r*   r8   rG   r9   r:   rJ   rI   r   r   r   test_join_multi_multi"  s    z(TestJoinMultiMulti.test_join_multi_multic                 C   s`   |j |jd}|j |jd}t| | ||d| }|j||d }t|| d S )Nr'   r   r   )	r?   r   r	   rC   r)   r  rB   rD   rE   r  r   r   r   test_join_multi_empty_frames4  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jd	d
dgtjd}t dddg|d}|j|d|gdd}t	
|| t |dddgdddgd}|j||jjgdd}t	
|| d S r   r   )rF   r  r   r   r   rJ   rI   r   r   r   r   H  s       z,TestJoinMultiMulti.test_merge_datetime_indexc                 C   s   t jdddgddgd}tddd	gd
ddgd|d}t jddddgddgd}tddddgddddgd|d}||}t| | 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   r<   )	r   r   r   rB   r	   rC   r)   rD   rE   )rF   Z
index_leftr   Zindex_rightr   rI   rJ   r   r   r   test_single_common_level_  s4      
 
   z+TestJoinMultiMulti.test_single_common_levelc                 C   s   t jddgddggddgd}t jdd	d
gddgd}t|dddddgid}t|ddddgid}||}t|ddddgdddtjgdd}t|| d S )NrX   ri   rn   r   r-   r.   r   )r   rX   )rn   ri   )rn   rX   r   r         (   )r   r   r   r   r
   Zfing)r   r   )	r   r   r   r   rB   r   rx   rD   rE   )rF   Zmidx1Zmidx3r   r   rI   rJ   r   r   r   test_join_multi_wrong_orderx  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  )Znumpyr   r   Zpandasr   r   r   r   r   r   r   Zpandas._testingZ_testingrD   Zpandas.core.reshape.concatr   Zpandas.core.reshape.merger	   r   r   r   r*   r8   r9   r:   r;   r  r   r   r   r   <module>   s2    


	



     Y