o
    ҷh0v                     @   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   s6   g d} g d}t jdt| }t| ||dS )z=left dataframe (not multi-indexed) for multi-index join tests)
barr
   r
   foor   bazr   quxr   snap)
twoonethreer   r   r   r   r   r   r      key1key2data)nprandomdefault_rngstandard_normallenr   r    r   X/var/www/html/venv/lib/python3.10/site-packages/pandas/tests/reshape/merge/test_multi.pyleft   s   r   c                 C   s   | }ddg|j _g d|_|S )z:right dataframe (multi-indexed) for multi-index join testsr   r   )j_onej_twoj_three)indexnamescolumns) multiindex_dataframe_random_datadfr   r   r   right   s   
r'   c                   C   :   t g dg dg dg dg ddg ddg dS )	N)Ar)   Br*   C)r)   r*   r)   r+   r)   )AMr,   IPr,   OP)hbwnhbhbor0   r/   )i  i?  i	  i  i\  )OriginDestinationPeriodTripPurpTripsr$   )r2   r3   r4   r5   r   	set_indexr   r   r   r   
left_multi'      
	r:   c                   C   r(   )	N)r)   r)   r*   r*   r+   r+   E)r)   r*   r)   r*   r)   r*   F)r,   r,   r-   r,   r.   r-   r,   )abcr?   r>   r?   r>   )d   P   Z   rB   K   #   7   )r2   r3   r4   LinkTypeDistancer7   )r2   r3   r4   rG   r8   r   r   r   r   right_multi5   r;   rI   c                   C      g dS )N)r2   r3   r4   r   r   r   r   r   on_cols_multiC      rK   c                   C   rJ   )N)r2   r3   r4   r5   rG   r   r   r   r   r   idx_cols_multiH   rL   rM   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rP   rQ   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_multikeyrT   FTc              
      s  g ddd   fdd}t ttttdtdd }ttjd		|d
ddgd}|
ddtjd	ddt|d tjd	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d	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   s    zNTestMergeMulti.test_left_join_multi_index.<locals>.bind_cols.<locals>.<lambda>c                    s   |   td S )Nr>   )maprb   )tsiordr   r   rc   f   s    r_   ra   g      Y@r`   r   g     @)fillna)r&   fr   rf   r   	bind_colsd   s   .z<TestMergeMulti.test_left_join_multi_index.<locals>.bind_colsc                    s   | j |d|d}t| t|d k sJ |d   rJ |d   r)J tj|d |d  dd  |jd d d df }tj|d |dd |jd u sTJ |rat||j	d	d
 t
| | |dd}tt||_t|| d S )Nr   rS      4th5thF)check_names	mergesortkind)rP   rT   rQ   )rU   r   isnaanyrW   assert_series_equalilocnamerX   sort_valuesr	   rV   r   r"   )r   r'   rT   resr\   outrj   icolsr   r   run_assertsi   s   z>TestMergeMulti.test_left_join_multi_index.<locals>.run_assertsr>   zrk   r   )i  r   r_   ra   r7   r`   r   i  floatrl   rm   T)inplace   %      +   )listrd   chrr   arangerb   r   r   r   choiceinsertintegersr   astypepermutationrv   copyr9   nanloc)rY   rT   r}   lcr   ir'   r   r{   r   test_left_join_multi_index`   s6   "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   )left_onright_indexrQ   rT   r'   )right_on
left_indexrQ   rT   )r	   r$   rW   rX   )rY   r   r'   rT   r[   merged_left_right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g dgddgddd	}ttd
g ddd}t|| ddgd}ttd
g ddgd tdd}t	|| d S )N   r)   r*   rk   r   r   lev1lev2r#   Amount)rw   AAABBB)rk   r   r   rk   r   r   r   )r   r   colrP   )r   r   r   r   )
r   ranger   from_productr   r   r	   rV   rW   rX   )rY   sr&   r\   r]   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}t|d}|d d d }t||tjddd}t|d d d |d d d tjddd}t||dd d S )	Ni'  r   r   i N  )r   r   value1)r   r   value2outerrQ   )	rW   makeStringIndexr   tiler   r   r   r   r	   )rY   r   r   r&   df2r   r   r    test_compress_group_combinations   s    
	z/TestMergeMulti.test_compress_group_combinationsc                 C   s  ddg}t g d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 g dd ddgd tjg d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k2)r   rk   r      r   r
         dtype)r   r   v)r   r
   )rk   r   v2      r"   r   r   rk   rp   T)rr   r   )rP   rT   )r   r   k3r   rq   )r   r   arrayr   int64r   from_tuplesrU   r   r   r   r   r   rW   rX   rx   float32int32)rY   r[   r   r"   r'   r\   r]   r   r   r   #test_left_join_index_preserve_order   sH   




	
z2TestMergeMulti.test_left_join_index_preserve_orderc                 C   s  t g dg dg dg dg dg dg dg dg d	g d
g
g dg dd}t g dg dg dg dg dg dg dg dg dg dg dg dg dg dgg ddg d}|j|g dd d!}t g d"g d#d$d%d&d'tjgg d(g d)g d*d+d%d,d-tjgg d.g d/g d0g d1g d2g d3g d4gg d5g d6d}t|| |j|g dd d7d8}|jg dd9d:}t|| d S );N)XYr+   r>   )Wr   r+   e)VQr)   h)r   RDr   )r   r   r   r?   )r   r   r)   r@   )r   r   r*   ri   )r   r   r+   g)r   r   r+   j)r   r   r*   d)colacolbcolctag)
r   r   r   rk   r   r      r   	   r   r$   r"   )r   r   r+   r   )r   r   r*   r   )r   r   r*   r   )r   r   r)   rk   )r   r   r)   r   )r   r   r*   r   )r   r   r+   r   )r   r   r+   r   )r   r   r+   i)r   r   r+   i)r   r   r+   r   )r   r   r   r   )r   r   r   r   )r   r   r)   )col1col2col3valr7   )r   r   r   )r   r   r   r   rO   )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   r   r   r   r?   )r   r   r)   r@   rk   )r   r   r)   r@   r   )r   r   r*   ri   r   )r   r   r*   ri   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   rk   rk   r   r   r   r   r   r   r   r   TrS   rp   rq   )r   r9   rU   r   r   rW   rX   rx   )rY   r   r'   r\   r]   r   r   r   +test_left_join_index_multi_match_multiindex  sv   
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g d
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 g dg dg dg dddtjgg dg dddtjggg d g d!d}t|| |j|ddd"d#}|jdd$d%}t|| t||	 ddd&}t
t||_t|| d S )'Nr@   r   r?   rk   r>   r   r   r   r   )r   r   rk   r   r   r   wxr   yr~   rr   qr   charr7   r   rO   )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   rk   rk   r   TrS   rp   rq   rQ   rP   )r   r9   rU   r   r   rW   rX   rx   r	   rV   r   r   r"   )rY   r   r'   r\   r]   	expected2r   r   r    test_left_join_index_multi_matchV  sP   



z/TestMergeMulti.test_left_join_index_multi_matchc                 C   s   t tdtjddtjddtdtjdddg d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 )Nabcder   r   )idv1r   dummyv3r7   r>   r?   gX9v?gʡE@)r   sv3r   r   rO   rk   )axis)r   r   r   r   r   r   r   r	   rR   rU   rW   rX   )rY   r   r'   r\   rdfr]   r   r   r   test_left_merge_na_buglet  s$   

z(TestMergeMulti.test_left_merge_na_bugletc              
   C   s   g dg dg dddt jgg dg dddt jgg d	g d
g	}t|g d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|g dd}|j|dd}|dj|ddd}|dt j}t|| d S )N)  r)         ?)r   r*   r   )  r*   r   i  r*   )  r*         @)r   r+   r   r+   )  r+   g      @)r   r+   r   )yearpanelr   r7   r)   r   r   r   r   r   i)r   r   r   r	   rh   replacerW   rX   )rY   r   frame
other_dataotherr\   r]   r   r   r   test_merge_na_keys  s.   







z!TestMergeMulti.test_merge_na_keysklassNc                 C      t g dg ddgd}t|j|_|jj}|d ur||}tjg dtjd}t g d|d}|j|d|gdd	}t	
|| t |g dg dd
}|j||jjgdd	}t	
|| d S Nr   )z
2016-01-01z
2017-01-01z
2018-01-01r>   r7   )i  i  i  r   )r>   key_1innerrO   )key_0a_xa_yr   pdto_datetimer"   r   r   r   r   r	   rW   rX   )rY   r   r&   	on_vector	exp_yearsr]   r\   r   r   r   test_merge_datetime_index     z(TestMergeMulti.test_merge_datetime_index
merge_typer   r'   c                 C   s2  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rft ddgtjtjtjgtdd|d
}|j|dddgd}|j	|dd}n%t tjtjtjgtd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   r   r   r   objectr	   rU   rW   rX   )rY   r
  r   r'   expected_indexr]   results_mergeresults_joinr   r   r   (test_merge_datetime_multi_index_empty_df  sD   
z7TestMergeMulti.test_merge_datetime_multi_index_empty_dfc                 C   s.   t g dg dg ddg ddd}|S )Nr   )r   rk   r   )ffffAPAi^ )household_idmalewealthr7   r  r8   )rY   	householdr   r   r   r    s   zTestMergeMulti.householdc              
   C   sP   t g dddddddtjgddd	d	d
dtjgg ddg ddddg}|S )Nrk   r   r   r   r   r   r   nl0000301109nl0000289783gb00b03mlx29lu0197800237nl0000289965ABN AmroRobecoRoyal Dutch ShellAAB Eastern Europe Equity FundPostbank BioTech Fonds      ?皙?333333?333333?r&        ?r$  )r  asset_idrw   sharer7   r  r)  )r   r   r   r9   )rY   	portfolior   r   r   r+    s2   
	
zTestMergeMulti.portfolioc                 C   sH   t g dg dg dg dg dg dddd	gjg d
d}|S )N)r   rk   rk   r   r   r   )r  r  r      xAr,  r,  )r  r  r   r   r!  r"  )r$  r%  r&  r'  r&  r(  )rk   r   r   r   r   r   )r  r  r  r  r  r  )r  r  rw   r*  r  r)  r  r)  )r  r  rw   r*  r7   )r   r9   reindex)rY   r]   r   r   r   r]   /  s   
#zTestMergeMulti.expectedc                 C   s.   |  }|  }|j|dd}t|| d S )Nr   r   )r   rU   rW   rX   rY   r+  r  r]   r\   r   r   r   test_join_multi_levelsV  s   z%TestMergeMulti.test_join_multi_levelsc                 C   sD   |  }|  }t| | dgddddg}t|| d S )Nr  r   rO   r)  )r   r	   rV   r9   rW   rX   r.  r   r   r   (test_join_multi_levels_merge_equivalence_  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   rT   r7   )r   rU   r   r   r   r   r   r   r-  r$   rW   rX   r.  r   r   r   test_join_multi_levels_outerl  s$   z+TestMergeMulti.test_join_multi_levels_outerc                 C   s   |  }|  }d|j_tjtdd |j|dd W d    n1 s%w   Y  |  }|jddg tjtdd |j|dd W d    d S 1 sPw   Y  d S )Nr   z+cannot join with no overlapping index names)matchr   r   r  z'columns overlap but no suffix specified)r   r"   rw   pytestraises
ValueErrorrU   	set_names)rY   r+  r  
portfolio2r   r   r   test_join_multi_levels_invalid  s   "z-TestMergeMulti.test_join_multi_levels_invalidc              	   C   s@  t g dddddddtjgg ddg ddd	d
g}t g dg dg ddd
dg}t g dg dg dg dg ddg djddgd}t| | d
gddg d}t|| t g dg dg dg dg ddg djddgd}t| | d
gd dg d}t|| d S )!Nr  r  r  r  r  r#  )r  r)  r*  r7   r  r)  )r  r  r  r  r  )               )S緖?D >?
?HRE?)r)  t
log_returnrC  )r   r   r   r   r   r   r   r   )r  r  r  r  r  r  r  r  )r9  r:  r;  r9  r:  r;  r<  r=  )r&  r&  r&  r'  r'  r'  r&  r&  )r>  r?  r@  r>  r?  r@  rA  rB  )r  r)  rC  r*  rD  )r  r)  rC  r*  rD  r   rO   )rk   r   r   r   r   r   r   r   r   r   r   r   )r  r  r  r  r  r  r  r  r  r  r  N)NNr9  r:  r;  r9  r:  r;  r<  r=  NN)r$  r%  r&  r&  r&  r'  r'  r'  r&  r&  r(  r$  )NNr>  r?  r@  r>  r?  r@  rA  rB  NNr   )	r   r   r   r9   r-  r	   rV   rW   rX   )rY   r  rD  r]   r\   r   r   r   test_join_multi_levels2  s   	



!

=B
z&TestMergeMulti.test_join_multi_levels2)__name__
__module____qualname__r^   r3  markparametrizer   r   r   r   r   r   r   r   r   r   asarrayr   r   r  r  fixturer  r+  r]   r/  r0  r1  r8  rE  r   r   r   r   rN   M   s8    
=
6E2 

.


&	rN   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	   rV   r9   
sort_indexrU   rW   rX   rY   r:   rI   rZ   rK   rM   r]   r\   r   r   r   test_join_multi_multi2  s   z(TestJoinMultiMulti.test_join_multi_multic                 C   s`   |j |jd}|j |jd}t| | ||d| }|j||d }t|| d S )Nr7   r   r   )	rR   r$   r	   rV   r9   rN  rU   rW   rX   rO  r   r   r   test_join_multi_empty_framesD  s   z/TestJoinMultiMulti.test_join_multi_empty_framesboxNc                 C   r   r   r  )rY   rR  r&   r  r  r]   r\   r   r   r   r  X  r	  z,TestJoinMultiMulti.test_merge_datetime_indexc                 C   s   t jg dddgd}tg dg dd|d}t jg d	dd
gd}tg dg dd|d}||}t| | dgddg d}t|| d S )N))K0X0)rS  X1)K1X2keyr   r   )A0A1A2)B0B1B2)r)   r*   r   ))rS  Y0)rV  Y1)K2Y2)ra  Y3r   )C0C1C2C3)D0D1D2D3)r+   r   r   rO   )rX  r   r   )	r   r   r   rU   r	   rV   r9   rW   rX   )rY   
index_leftr   index_rightr'   r\   r]   r   r   r   test_single_common_levelo  s(   

z+TestJoinMultiMulti.test_single_common_levelc                 C   s   t jddgddggddgd}t jg dddgd}t|d	g d
id}t|dg did}||}t|g d
dddtjgdd}t|| d S )Nrk   r   r   r   r>   r?   r   ))r   rk   )r   r   )r   rk   r   )
         (   )r"   r   r   )r   r
   fingrs  r   r
   )r   r   )	r   r   r   r   rU   r   r   rW   rX   )rY   midx1midx3r   r'   r\   r]   r   r   r   test_join_multi_wrong_order  s   
z.TestJoinMultiMulti.test_join_multi_wrong_order)rF  rG  rH  rP  rQ  r3  rI  rJ  r   rK  r   r   r  rn  rv  r   r   r   r   rM  1  s    
rM  )numpyr   r3  pandasr  r   r   r   r   r   r   pandas._testing_testingrW   pandas.core.reshape.concatr   pandas.core.reshape.merger	   rL  r   r'   r:   rI   rK   rM   rN   rM  r   r   r   r   <module>   s4     


	



     i