U
    /eT                    @   sF  d dl Z d dlZd dlZd dlZd dlmZ d dl	m
Z
 d dlmZmZmZ d dlmZ d dlmZ d dlmZmZmZmZmZ d dlmZmZmZmZmZmZ d dl m!Z! d	d
 Z"dd Z#dd Z$dd Z%dd Z&ej'(dddddgdd Z)dd Z*dd Z+dd Z,d d! Z-d"d# Z.ej'(d$d%d&gej'(d'd(d)d*gd+d, Z/ej'(d$d%d&gej'(d'd(d)d*gej'(d-d&d%gd.d/ Z0d0d1 Z1d2d3 Z2d4d5 Z3d6d7 Z4d8d9 Z5d:d; Z6ej'(d<d=d>gd?d@ Z7dAdB Z8ej'(dCddgdDdE Z9ej'(dCddgdFdG Z:ej'(dCddgdHdI Z;dJdK Z<dLdM Z=ej'(dNdOdPdQe>dRfdSdTe>dRfdUe>dRfgdVdW Z?ej'(dddddgej'(dXd%d&gdYdZ Z@ej'(dddddgd[d\ ZAej'(dddgd]d^ ZBd_d` ZCej'(dadbdcddgej'(ddedfgej'(dgejDdhej'jEdidjdkejDdlej'jFdkgdmdn ZGdodp ZHej'jIej'(dddddgdqdr ZJej'(dddddgdsdt ZKdudv ZLej'(dddddgdwdx ZMdydz ZNd{d| ZOej'(d}d&d%gd~d ZPdd ZQej'(dddgej'(d}d&d%gdd ZRdd ZSdd ZTdd ZUdd ZVdd ZWdd ZXdd ZYdd ZZdd Z[dd Z\dd Z]dd Z^ej'(dddddddgdd Z_dd Z`dd Zadd Zbej'jceddjdd Zdej'jceddjdd Zeej'jceddjdd Zfej'jceddjdd Zgdd Zhdd Zidd Zjdd Zkdd Zlej'(dgdhejDdlej'jFdkgdd Zmej'(dd%d&gdd ZnddĄ ZoddƄ ZpddȄ Zqddʄ Zrdd̄ Zsej'(ddddgej'(dddgej'(dddgddԄ Ztej'(dddgddք Zuej'(dddgej'(ddTddgej'(ddTddgdd܄ ZvdS )    N)compute_as_if_collection)PANDAS_GT_140PANDAS_GT_200tm)_Frame)concat)_maybe_align_partitionsalign_partitionsconcat_indexed_dataframes	hash_joinmerge_indexed_dataframes)assert_divisions	assert_eq
check_metaclear_known_categorieshas_known_categories	make_meta)hlg_layer_topologicalc               
      s  t jddddddgtddd	d
ddddgd} t| d	ddg}t jddddgtddddddgd}t|dddg}tjdd	idd}t||\\}}  fdd}|||||  |jdf|jdfg|jdf|jdfg|jdf|jdfg|jdf|jdfg|jdf|jdfggks(t	t|||\\}}} |||||  |jdfd |jdfg|jdfd |jdfg|jdfd |jdfg|jdfd |jdfg|jdfd |jdfggkst	t
|d	 t dddddddgdddddddgd}	t dddddddgdddddddgd}
t|	dt|
dft|	dt|
dft|	dt|
dft|	dt|
dffD ]2\}}t||\\}}}}t
||	 t
||
 q~t jdddddddgdddddddgdtdd}	t jdddddddgdddddddgdtdd}
t|	dt|
dft|	dt|
dft|	dt|
dft|	dt|
dffD ]2\}}t||\\}}}}t
||	 t
||
 qvd S )N                  abdabdxy
         (   2   <   indexabdaF   P   d   )sr   r*   i8c                    s   t | tjstt |tjs tt |tjs0tt |tjs@tt| | t|| dks`tt  tsnttdt  kstd S )N)r   r    r!   r#   r(   r)   r   )
isinstancedd	DataFrameAssertionErrorr   listlenabaabbL	divisions C/tmp/pip-unpacked-wheel-dbjnr7gq/dask/dataframe/tests/test_multi.py_check+   s    

z%test_align_partitions.<locals>._checkr      r3   r4   cdabcdefgZfghijkl)pdr.   r0   r-   repartitioncoreZScalarr	   _namer/   r   from_pandas)Ar3   Br4   r*   r5   r6   r<   ssZldfrdflhsrhsZlresultZrresultdivpartsr:   r7   r;   test_align_partitions   sl     (
,,
$ $ 
rP   c               
   C   s   t dddddddgdddddddgd} tj| dd	}tj| dd
d}|jrTttt t	|| W 5 Q R X tj| d dd
d}tj| dd
d}|jrt|jrttt t	|| W 5 Q R X d S Nr   r   r   r   r   r   r=   r>   npartitionsFrS   sort)
rC   r.   r-   rG   known_divisionsr/   pytestraises
ValueErrorr	   dfddfddf2r:   r:   r;   'test_align_partitions_unknown_divisionsj   s    ,


r^   c               
   C   s  t dddddddgdddddddgd} tj| d dd	}tj| dd	}t||g\}}||ksht||kstttj| d dd
d}tj| dd
d}|jrt|jrtt||g\}}||kst||ksttj| d dd	}tj| dd	}t||g\}}|j|jksttj| d dd
d}tj| dd
d}|jrFt|jrRtt	t
 t||g W 5 Q R X tj| dd	}tj| dd
d}|jrtt	t
 t||g W 5 Q R X d S rQ   )rC   r.   r-   rG   r   r/   rV   r9   rW   rX   rY   )r[   r\   r]   r3   r4   r:   r:   r;   test__maybe_align_partitions~   s8    ,

r_   c                  C   s"  t jdddddddgiddddddgd	} t| dddg}t jd
tdiddddddgd	}t|ddddg}t||dd}|jd |jd kst|jd t|j|j kstt	|| 
| t||dd}|jd |jd kst|jd |jd kstt	|| j
|dd t||dd}|jd dks<t|jd t|j|j ks\tt	| | j
|dd t||dd}|jd dkst|jd dkstt	| | j
|dd tt||ddjtt||ddjksttt||ddjtt||ddjkstd S )Nr   r   r   r   r   r   r   r=   r$   r   abcdef   lefthowr   rightinnerouter)rC   r.   r-   rD   r0   r   r9   r/   maxr   joincomputesorteddask)rH   r3   rI   r4   r@   r:   r:   r;   1test_merge_indexed_dataframe_to_indexed_dataframe   s4    *" rn   c                 C   s   t | tjr| jdd}n| }t |tjr8|jdd}n|}t|j|j t |tjr~|t	|jj
}|t	|jj
}n| j
}| j
}tj|| d S )NsyncZ	scheduler)r,   r-   r.   rk   r   Zassert_index_equalcolumnsrC   sort_valuesr0   values_compatZassert_numpy_array_equal)r5   r6   r3   r4   avZbvr:   r:   r;   list_eq   s    

rv   rd   rg   rb   rf   rh   c           	   	   C   s  t ddddddgddddddgd}t|dddg}t ddddddgddddddgd	}t|dddg}t|d
|d
| }t|jd rt|	 }t 
||| d
}t|| t|d|ddd|d}t|jd rt|jdkst|j	dd}t 
||dd dd}t|| t|d
|d
d|djt|d
|d
d|djksNtt|d
|d
d|djt|d
|d
d|djkstd S )Nr   r   r   r   r   r   r   r   r   zr   re   r   rx   rh   )rS   shufflezsingle-threadedrp   rg   )ry   )rC   r.   r-   rD   r   r   rm   is_materializedr/   rk   mergerv   rS   rF   )	rd   shuffle_methodrH   r3   rI   r4   r@   resultexpectedr:   r:   r;   test_hash_join   s.    ((

r   c            
   	   C   s   t ttdddddddgd} t ttd	d
dddgd}t ttddddddgd}| j|dddd}|j|dddd}tj| dd}tj|dd}tj|dd}|j|dddd}|j|dddd}	t||	 d S )Nr   A0A1A2A3ZA4ZA5)keyrH   r   B0B1B2B3)r   rI   r   r   C0C1C2C3)r   Crg   Z_l_rrd   lsuffixrsuffixr   rR   r   )rC   r.   r0   rangerj   r-   rG   r   )
df1df2df3Zjoin_pdZmulti_join_pdddf1r]   ddf3Zjoin_ddZmulti_join_ddr:   r:   r;   test_sequential_joins  s     "r   c                  C   s   t jdtdiddddddd	d
ddddgd} tj| dd}t jdtdidddddddd	ddddgd}tj|dd}t j| |ddd}tj||ddd}t|| d S )Nleft_valZabcdabcdabcdr   r   r=   	   r            r            r$   r   rR   	right_valZxyzxyzxyzxyzr   r                  T
left_indexright_index)rC   r.   r0   r-   rG   
merge_asofr   rH   r3   rI   r4   r   r@   r:   r:   r;   test_merge_asof_indexed  s    

r   c                  C   s   t dddgdddgd} tj| dd	}t ddd
ddgddd
ddgd}tj|dd	}t j| |dd}tj||dd}t||dd d S )Nr   r   r   r3   r4   r@   r3   r   r   rR   r   r   r=   r3   r   onFcheck_indexrC   r.   r-   rG   r   r   r   r:   r:   r;   test_merge_asof_on_basic-  s    $r   c               	   C   s   t dddgdddgd} tj| dd	}t ddd
ddgddd
ddgd}tj|dd	}tjtdd tj||ddd W 5 Q R X tjtdd tj||ddd W 5 Q R X tjtdd tj||dddd W 5 Q R X d S )Nr   r   r   r3   r4   r@   r   r   rR   r   r   r=   r   combination of bothmatchr   )r   left_onr   )r   right_on)r   r   r   rC   r.   r-   rG   rW   rX   rY   r   rH   r3   rI   r4   r:   r:   r;   'test_merge_asof_on_lefton_righton_error9  s    $r   c               	   C   s   t dddgdddgddd	gd
} tj| dd}t dddddgdddddgdddddgd}tj|dd}tjtdd tj||dddd W 5 Q R X tjtdd tj||dddd W 5 Q R X tjtdd tj||ddddd W 5 Q R X d S )Nr   r   r   r   r   r   r3   r4   r@   )r3   r4   r   r   rR   r=   r      )r3   r4   r   r   r   r   )r   byleft_byr   )r   r   right_by)r   r   r   r   r   r   r:   r:   r;   'test_merge_asof_by_leftby_rightby_errorG  s    $(r   allow_exact_matchesTF	directionZbackwardZforwardZnearestc                 C   s   t dddgdddgd}tj|dd	}t ddd
ddgddd
ddgd}tj|dd	}t j||d| |d}tj||d| |d}t||dd d S )Nr   r   r   r3   r4   r@   r   r   rR   r   r   r=   r   )r   r   r   Fr   r   )r   r   rH   r3   rI   r4   r   r@   r:   r:   r;   test_merge_asof_onW  s&    $        r   unknown_divisionsc                 C   s  t jdddgdddgdddd	gd
}tj|dd}t jdddddgiddddgd
}tj|dd}|rxdt|j |_t j||dd| |d}tj||dd| |d}t|| dD ]}	ddddgddddgfddddgddddgfddddddgdddddddddg	fdddgddddddddddddgffD ]\}
}t jd|
idd |
D d
}tj||	d}t jd|i|d
}tj||	d}|rdt|j |_t j||dd| |d}tj||dd| |d}t|| q0qd S )Nr   r   r   r3   r4   r@   r   r   r    r$   r   rR   r   r   r   r=   NT)r   r   r   r   )r   r   r   r   r   r      r   r   c                 S   s   g | ]}|d  qS )r   r:   ).0r   r:   r:   r;   
<listcomp>  s     z7test_merge_asof_left_on_right_index.<locals>.<listcomp>)rC   r.   r-   rG   r1   r9   r   r   )r   r   r   rH   r3   rI   r4   r   r@   ZnpartsZa1Zidx2r:   r:   r;   #test_merge_asof_left_on_right_indexi  sh    &"

$$r   c                  C   s   t jddddgidddgd} tj| d	d
}t jddd	dddgidd	dddgd}tj|d	d
}t j| |ddd}tj||ddd}t|| d S )Nr   r3   r4   r@   r   r   r   r$   r   rR   r   r   r   r=   Tr   r   r   r:   r:   r;   &test_merge_asof_indexed_two_partitions  s    &r   c                  C   s8  dd dD } 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dg}dddddg}dddddg}t j| |||ddddd gd!}tj|d"d#}	t j||||d$ddd%d&gd!}
tjd'd( |
jd)d* |
jd*d+ gd)d*d"gd,}t j|
|ddd-}tj||	ddd-}t||d.d/ d S )0Nc                 S   s   g | ]}t |qS r:   rC   to_datetimer   rA   r:   r:   r;   r     s   z)test_merge_asof_on_by.<locals>.<listcomp>2016-05-25 13:30:00.023r   z2016-05-25 13:30:00.030z2016-05-25 13:30:00.0412016-05-25 13:30:00.048z2016-05-25 13:30:00.049z2016-05-25 13:30:00.072z2016-05-25 13:30:00.075GOOGMSFTAAPL     @皙I@\(I@QI@(\X@zGJ@=
ףp@{GI@=
ףpI@      J@q=
ףX@ףp=
@p=
J@c                 S   s   g | ]}t |qS r:   r   r   r:   r:   r;   r     s   r   z2016-05-25 13:30:00.038r   r   r   \(@(\@     X@K      r)   timetickerbidaskr   r   r   r   rq   r   rR   r   r   pricequantityr   r   c                 S   s   | S r   r:   r   r:   r:   r;   <lambda>      z'test_merge_asof_on_by.<locals>.<lambda>r   r   r   r9   )r   r   Fr   )rC   r.   r-   rG   from_mapilocr   r   Ztimes_AZ	tickers_AZbids_AZasks_AZtimes_BZ	tickers_BZprices_BZquantities_BrH   r3   rI   r4   r   r@   r:   r:   r;   test_merge_asof_on_by  sB    


r   c                  C   sH  dd dD } 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dg}dddddg}dddddg}t j| |||ddddd gd!}tj|d"d#}	t j||||d$ddd%d&gd!}
tjd'd( |
jd)d* |
jd*d+ gd)d*d"gd,}t j|
|ddt d-d.}tj||	ddt d-d.}t||d/d0 d S )1Nc                 S   s   g | ]}t |qS r:   r   r   r:   r:   r;   r     s   z3test_merge_asof_on_by_tolerance.<locals>.<listcomp>r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   c                 S   s   g | ]}t |qS r:   r   r   r:   r:   r;   r     s   r   r   r   r   r   r   r)   r   r   r   r   r   r   r   rR   r   r   r   c                 S   s   | S r   r:   r   r:   r:   r;   r      r   z1test_merge_asof_on_by_tolerance.<locals>.<lambda>r   r   r   r   Z2ms)r   r   	toleranceFr   	rC   r.   r-   rG   r   r   r   Z	Timedeltar   r   r:   r:   r;   test_merge_asof_on_by_tolerance  sB    


r   c                  C   sL  dd dD } 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dg}dddddg}dddddg}t j| |||ddddd gd!}tj|d"d#}	t j||||d$ddd%d&gd!}
tjd'd( |
jd)d* |
jd*d+ gd)d*d"gd,}t j|
|ddt d-d.d/}tj||	ddt d-d.d/}t||d.d0 d S )1Nc                 S   s   g | ]}t |qS r:   r   r   r:   r:   r;   r   +  s   zDtest_merge_asof_on_by_tolerance_no_exact_matches.<locals>.<listcomp>r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   c                 S   s   g | ]}t |qS r:   r   r   r:   r:   r;   r   ;  s   r   r   r   r   r   r   r)   r   r   r   r   r   r   r   rR   r   r   r   c                 S   s   | S r   r:   r   r:   r:   r;   r   Z  r   zBtest_merge_asof_on_by_tolerance_no_exact_matches.<locals>.<lambda>r   r   r   r   Z10msF)r   r   r   r   r   r   r   r:   r:   r;   0test_merge_asof_on_by_tolerance_no_exact_matches*  s^    


r   c               	   C   s   t dddgdddgd} tj| dd	}t ddd
ddgddd
ddgd}tj|dd	}tj||dd}tjtdd |  W 5 Q R X d S )Nr   r   r   r3   r4   r@   r   r   rR   r   r   r=   r   r   z
right keysr   )	rC   r.   r-   rG   r   rW   rX   rY   rk   )rH   r3   rI   r4   r}   r:   r:   r;   test_merge_asof_unsorted_raisesr  s    $r   c                  C   s<  t ttdttdd} tj| dd}| | jdk   jddid}tj|dd}tj	|
d	|
d	d
d
d}t j	| 
d	|
d	d
d
d}t||dd tj	|
d	|
d	d
d
d}t j	|
d	| 
d	d
d
d}t||dd tj	|
d	|
d	d
d
d}t j	|
d	|
d	d
d
d}t||dd d S )Nr   )	index_colgood_valr   rR   r   r   Z	empty_valr   r   Tr   Fr   )rC   r.   r0   r   r-   rG   r   copyrenamer   	set_indexr   )Zgood_dfZgood_ddempty_dfZempty_dd	result_dd	result_dfr:   r:   r;   test_merge_asof_with_empty~  sX     r  zleft_col, right_col)	endofweek	timestamp)r  r  c                 C   s   t | ddddddgdddddddgi}t |ddddgd	ddddgi}t j||| |d
}tjtj|ddtj|dd| |d
}t||dd d S )Nr   r   r   r   ZGroupColi  i!  r   ZGroupValr   r   rR   Fr   )rC   r.   r   r-   rG   r   )Zleft_colZ	right_colr   r   r  r  r:   r:   r;   test_merge_asof_on_left_right  s       "r
  c                  C   sp   t tt dgd dddgd} t j| | dd}tddD ],}tj| |d	}tj||dd}t|| q>d S )
Nz1-1-2020r   r   r   )tsfoor  )rb   rf   r   r   rR   )	rC   r.   dictr   r   r   r-   rG   r   )r[   r~   rS   r\   r}   r:   r:   r;   (test_merge_asof_with_various_npartitions  s    $r  rj   c              	   C   s6  t jddddddgtddddddddgd	}t|dddg}t jd
ddddddgiddddddgd	}t|ddddg}t j||gd| dd}t  tdt	 t
||g| d}t|| tt
||g| djtt
||g| djksttt
||gddjtt
||gddjks(tW 5 Q R X d S )Nr   r   r   r   r   r=   r`   r   r$   r   r   r   r!   r"   r#   r(   r   ra   r   F)axisrj   rU   ignorerj   rg   rh   )rC   r.   r0   r-   rD   r   warningscatch_warningssimplefilterFutureWarningr
   r   rl   rm   r/   )rj   rH   r3   rI   r4   r~   r}   r:   r:   r;   test_indexed_concat  s$     *


r  c                 C   s$  t jddddddgtddddddddgd	}t|d}t jddddddgtddd
dddddgd	}t|d}t jddddddgtddd
dddddgd	}t|d}ddi}||||f||||ffD ]J\}}	}
}t j|
|gfd| i|}tj||	gfd| i|}t|| qd S Nr   r   r   r   r   r=   r`   r   r$   ra   r   r   r   r   r   r   rx   rU   Frj   )rC   r.   r0   r-   rG   r   r   rj   pdf1r   pdf2r]   pdf3r   kwargsdd1dd2Zpd1Zpd2r~   r}   r:   r:   r;   test_concat  s,       

r   c                 C   s4  t jddddddgtddddddddgd	}t|d}t jddddddgtddd
dddddgd	}t|d}t jddddddgtddd
dddddgd	}t|d}ddi}|j|j|j|jf|j|j|j|jffD ]J\}}	}
}t j|
|gfd| i|}tj||	gfd| i|}t|| qd S r  )	rC   r.   r0   r-   rG   r   rx   r   r   r  r:   r:   r;   test_concat_series  s,       r!  c                  C   s   t jddddddgtddddddddgd	} t| d}t jd
tdiddddddgd	}t|d}tt|d
 |d< |jddj	 d |d< ddi}t j
| |gf|}tj
||gf|}t|jd rtt|| d S )Nr   r   r   r   r   r=   r`   r   r$   r   ra   r   r   r   r   r   r   )partition_countrU   F)rC   r.   r0   r-   rG   r   r1   Zassignr"  Zcumsumr   r   rm   rz   r/   r   )r  r   r  r]   r  r~   r}   r:   r:   r;   &test_concat_with_operation_remains_hlg)  s     "r#  c            
      C   s   t jddddgidd} t jg dd}t | |g}tj| dd}tj|dd}t||g}t|| t jg d	gdd
}t | |g}tj|dd}t||g}	|	jjd|	j_t||	dd d S )Nr3   r)      ,  int64dtyper   rR   r   )rq   r(  F)Zcheck_dtype)	rC   r.   r   r-   rG   r   _metar%   astype)
r[   r  Z	df_concatr\   Z	empty_ddfZ
ddf_concatZempty_df_with_colZdf_concat_with_colZempty_ddf_with_colZddf_concat_with_colr:   r:   r;   test_concat_dataframe_empty@  s    
r+  zvalue_1, value_2)      ?r   )r,  oner,  z
1970-01-01)r   r-  r   r-  c           
      C   s   t d| gi}t d|gi}t j||gdd}|d j}tj|dd}tj|dd}tj||gdd}t|dd  }	|	||gkst	d S )Nr   r   r  r   rR   c                 S   s   | j S r   )Zdtypesr   r:   r:   r;   r   p  r   z.test_concat_different_dtypes.<locals>.<lambda>)
rC   r.   r   r(  r-   rG   r0   map_partitionsrk   r/   )
Zvalue_1Zvalue_2Zdf_1Zdf_2r[   Zpandas_dtypeZddf_1Zddf_2r\   Zdask_dtypesr:   r:   r;   test_concat_different_dtypesW  s    
r0  on_indexc              	   C   s  t ttdtd ttdd d}t tdtdd}t|d}t|d}dg}| }}|r|	d}|	d}d }t
jdd8}	t
d tj|||| ||d}
td	d
 |	D }W 5 Q R X t|
t jr|
n|
 }
|
 j }|r|s|rtd S )Nr   r   rH   rI   rH   Trecordalways)r   rd   r   r   c                 s   s   | ]}d t |kV  qdS )z!merge column data type mismatchesN)str)r   rr:   r:   r;   	<genexpr>  s     z,test_merge_columns_dtypes.<locals>.<genexpr>)rC   r.   r0   nparanger*  floatr-   rG   r  r  r  r  r{   anyr,   rk   Zisnars   r/   )rd   r1  r   r   r3   r4   r   r   r   r4  r}   ZwarnedZhas_nansr:   r:   r;   test_merge_columns_dtypesu  s4    *


     r=  c              
   C   s  t ddddddgddddddgd}t|dddg}t ddddddgddddddgd	}t|dddg}ttj||d
d
| |dt j||d
d
| d tj||d| d}t|t j||d| d tdd |jD st	ttj||dd| |dt j||dd| d ttj||dd| d|dt j||dd| dd ttj||| |dt j||| d ttj||| |dt j||| d ttj||| |dt j||| d ttj||| |dt j||| d ttj||d
d
| |dt j||d
d
| d ttj||d
d
| d|dt j||d
d
| dd ttj||dd
| |dt j||dd
| d ttj||dd
| d|dt j||dd
| dd d S ) Nr   r   r   r   r   r   r   r   rw   T)r   r   rd   ry   )r   r   rd   r   r   rd   c                 s   s   | ]}|d kV  qd S r   r:   r   r:   r:   r;   r8    s     ztest_merge.<locals>.<genexpr>r   rx   )r   r   rd   ry   )r   r   rd   )12)r   r   rd   suffixesry   )r   r   rd   rA  rd   ry   rc   )r   r   rd   rA  ry   )r   r   rd   rA  )r   r   rd   ry   r   r   rd   )r   r   rd   rA  ry   )r   r   rd   rA  )
rC   r.   r-   rD   r   r{   rv   allr9   r/   )rd   r|   rH   r3   rI   r4   r}   r:   r:   r;   
test_merge  s    ((     	$$$$     		rE  c                 C   s   t ddddgddddgd}t ddd	d	gd
dddgd}tj|dd}tj|dd}|j|d|d}|j|d|d}t||dd |jdd |jd  d S )Nr   r   r   r   r=   ra   )r3   valr   r   r   r   r   r   r   rR   r3   r>  Fr   c                 S   s   | S r   r:   r   r:   r:   r;   r     r   z*test_merge_empty_left_df.<locals>.<lambda>)meta)	rC   r.   r-   rG   r{   r   r/  r)  rk   )r|   rd   rb   rf   Zdd_leftZdd_rightmergedr~   r:   r:   r;   test_merge_empty_left_df  s      rI  c               	   C   sn   t ddddgddddgd	} t d
dddgddddgd	}tjtdd tj| |dd W 5 Q R X d S )Nr   r   r   r   r   r   r   r   r2  r   r   r   r   ZD0ZD1ZD2ZD3z1dask.dataframe.merge does not support how='cross'r   Zcrossrc   )rC   r.   rW   rX   rY   r-   r{   )rb   rf   r:   r:   r;   test_merge_how_raises   s    



 rJ  rO   )r   r   )r   r   )r   r   ZleftsemiZleftantienginepandasz,Pandas does not support leftsemi or leftanti)reason)Zmarkscudfc              
   C   s0  | dkrt d}t d}ttjddddddd	d
gdddddgdddddgd}tdddddgdddddgdddd d!gd"}| dkr||}||}|j||d# d$}|j||d% d$}n$tj||d# d$}tj||d% d$}|j	|d&|d'
d&g}	|j	|d&|d'
d&g}
t|
|	d(d) d S )*NrN  	dask_cudfe   j   )stopZJohnZTomZHarryZRahulZSakilZCalZMumZDelZBaniP  i@  i8 i`  i_ )emp_idnameZcityZsalaryi  i  i  i  i  g   i   f   ZDaskZSparkr   PythonR)Zskill_idrS  Z
skill_namer   rR   r   rS  r>  Fr   )rW   importorskiprC   r.   r9  r:  rG   	from_cudfr-   r{   rr   r   )rK  rd   rO   rN  rO  ZempZskillsZdd_empZ	dd_skillsexpectr}   r:   r:   r;   test_merge_tasks_semi_anti_cudf  s4    



r]  c               
   C   s   t dddddddgdddddddgd} t dddddddgdddddddgd	}tj| dd
}tj|dd
}|j|dddd}tdd |jD rtd S )Nr   r   r   r   r   r   r=   r>   r?   rR   r3   rA   tasks)r   r   ry   c                 s   s   | ]}d |d kV  qdS )Zpartdr   Nr:   r   kr:   r:   r;   r8  Q  s     z2test_merge_tasks_passes_through.<locals>.<genexpr>)rC   r.   r-   rG   r{   r<  rm   r/   )r3   r4   r5   r6   ccr:   r:   r;   test_merge_tasks_passes_throughH  s    ,,rb  c                 C   sN  t dddddddgdddddddgd}t dddddddgdddddddgd	}t jdddddddgdddddddgdtd
d}t jdddddddgdddddddgd	td
d}|}t jddddgddddgd	tdd}|}t jddddgddddgd	tdd}	t jdddddddgdddddddgdtdd}
t jddddgddddgd	tdd}t jdddddddgdddddddgdtdd}t jddddgddddgd	tdd}t jdddddddgdddddddgdtd
d}t jddddgddddgd	tdd}dd }||f||f||f||	f|
|f||f||ffD ]\}}dD ]\}}t||}t||}ttj||| dd|d|||| ddd ttj||| dd|d|||| ddd ttj||| dd|dd|||| dddd ttj||| dd|dd|||| dddd t|j|| dd|d|j|| ddd t|j|| dd|d|j|| ddd ttj||| dd|d t j||| ddd! ttj||| d"d#|d t j||| d"d#d! ttj||| dd|dd$t j||| dddd% ttj||| d#d"|dd$t j||| d#d"dd% ttj||| dd|d t j||| ddd! ttj||| d#d"|d t j||| d#d"d! t|j|| dd|d |j|| ddd! t|j|| d"d#|d |j|| d"d#d! t|j|| dd|d |j|| ddd! t|j|| d#d"|d |j|| d#d"d! qpqbd S )&Nr   r   r   r   r   r   r=   r>   r?   rB   r$   ra   r   abdgr   r   r   lmnopqrabcdcdefghifghic                 [   s6   t j| |f|}t|dkr2||j| jjS |S Nr   )rC   r{   r1   r  r%   r*  r(  )rb   rf   r  outr:   r:   r;   pd_mergex  s    z.test_merge_by_index_patterns.<locals>.pd_merge)r   r   )r   r   )r   r   Trd   r   r   ry   rd   r   r   )rd   r   r   ry   	indicator)rd   r   r   rn  r3   r@   rd   r   r   ry   rd   r   r   r4   rA   )rd   r   r   ry   rn  )rd   r   r   rn  )rC   r.   r0   r-   rG   r   r{   rv   )rd   r|   pdf1lpdf1rpdf2lpdf2rpdf3lpdf3rpdf4lpdf4rpdf5lpdf5rpdf6lpdf6rpdf7lpdf7rrj  pdlpdrlpartrpartddlddrr:   r:   r;   test_merge_by_index_patternsT  s   ,,$ $ (($ ($ ($ (	
	     	     
          		                          r  c              
   C   s  t jtddddddddgd	tdd
}t jtddddddddgdtdd
}|}t jtddddddddgdtdd
}|}t jtdddddgdtdd
}t jtddddddddgd	tdd
}t jtdddddgdtdd
}	t jtddddddddgd	tdd
}
t jtdddddgdtdd
}t jtddddddddgd	tdd
}t jtdddddgdtdd
}t jtddddddddgd	tdd
}t jtdddddgdtdd
}||f||f||f||	f|
|f||f||ffD ]\}}dD ]\}}t||}t||}t|j|| |d|j|| d t|j|| |d|j|| d t|j|| dd|d|j|| ddd t|j|| dd|d|j|| ddd qqd S )NrB   r=   r   r   r   r   r   r   r>   r$   r?   Zgfedcbarc  Zabcabcer&   rd  re  rf  ZababZaabbccdZaabbrg  rk  rB  rc   lr7  )rd   r   r   ry   r   )rC   r.   r0   r-   rG   r   rj   )rd   r|   rq  rr  rs  rt  ru  rv  rw  rx  ry  rz  r{  r|  r}  r~  r  r  r  r  r  r  r:   r:   r;   test_join_by_index_patterns+  s       $ $ $ $ $
          r  c                  C   s   t jddddgidddgd} tj| dd}t jdd	d
dgidddgd}| j|ddd}|j|ddd}|j|jksztt|| d S )Nr3   r4   r@   r   r   r   r$   rR   r,  g       @g      @rg   )rd   r   )rC   r.   r-   rG   rj   r9   r/   r   )r[   r\   Zright_dfr~   actualr:   r:   r;    test_join_gives_proper_divisions  s    r  c                 C   s  t jtdtddddddddd	d
dg
dtdd}t jtdtddd
d	dddddddg
dtdd}t jtdtddddddddd	d
dg
dtdd}t jtdtddd
d	dddddddg
dtdd}t jtdtddddddddd	d
dg
dtdd}t jtdtddd
d	dddddddg
dtdd}||f||f||ffD ]\}}	dD ]\}
}t||
}t|	|}t|j|| |d|j|	| d t|j|| |d|	j|| d ttj||| dd|dt j||	| ddd ttj||| dd|dt j|	|| ddd ttj||| dd|d t j||	| ddd! ttj||| d"d#|d t j||	| d"d#d! ttj||| dd|d t j|	|| ddd! ttj||| d#d"|d t j|	|| d#d"d! ttj||| dd"gdd#g|d t j||	| dd"gdd#gd! qhqZd S )$N
abcdefghijr   r   r   r   r   r   r=   ra   r   r   )r3   r4   r@   r$   )rA   efZ
abcdeabcdeZ
abcabcabcaZ
edcbaedcbaZ
aaabbbcccdZ
fghijklmnoZ
aaaaaaaaaaZ
aaabbbccaaZ
abbbbbbbbbZ
ABCDEFGHIJrk  rB  rc   Trl  rm  r3   rA   ro  rp  r4   r  )	rC   r.   r0   r-   rG   r   rj   r{   rv   )rd   r|   rq  rr  rs  rt  ru  rv  r  r  r  r  r  r  r:   r:   r;   test_merge_by_multiple_columns  s   			                       r  c                  C   sz  t tdd tdd tjdd} t| d}tt	|t 	|  ttj	|dd	t j	| dd	 ttj	|dd
t j	| dd
 ttj	|ddgddt j	| ddgdd ttj	|dddgddt j	| dddgdd t|	 | 	  t|j	dd	| j	dd	 t|j	dd
| j	dd
 t|j	ddgdd| j	ddgdd t|j	dddgdd| j	dddgdd d S )Nre  r   ZXYr   r   )rH   rI   r   r   r   )id_vars)
value_varsrH   Zmyvar)r  var_namerI   myval)r  r  Z
value_name)
rC   r.   r0   r9  randomrandnr-   rG   rv   Zmelt)Zpdfr\   r:   r:   r;   	test_melt   s4    "r  c                  C   s   t jddddddgtddd	d
ddddgd} tj| dd}t ddddgtdd}|j|ddd}t|jd r~t	t
dd |jD st	t|j|ddd| j|ddd d S )Nr   r   r   r   r   r   r   r   r   r   r    r!   r"   r#   r$   rR   r&   r  r   rg   r>  re   c                 s   s   | ]}d |d kV  qdS ry   r   Nr:   r_  r:   r:   r;   r8  K  s     z<test_cheap_inner_merge_with_pandas_object.<locals>.<genexpr>)rC   r.   r0   r-   rG   r{   r   rm   rz   r/   rD  rv   )r3   dar4   dcr:   r:   r;   )test_cheap_inner_merge_with_pandas_objectA  s     r  flipc                 C   s  t jddddddgtddd	d
ddddgd}tj|dd}t ddddgtdd}tj|ddd}| rv||fn||f}| r||fn||f}tj|ddd}t|jd rt	t
dd |jD st	t|jt|jd t|j kst	t|t j|ddd d S )Nr   r   r   r   r   r   r   r   r   r   r    r!   r"   r#   r$   rR   r&   r  FrT   r   rg   r>  re   c                 s   s   | ]}d |d kV  qdS r  r:   r_  r:   r:   r;   r8  _  s     z4test_cheap_single_partition_merge.<locals>.<genexpr>)rC   r.   r0   r-   rG   r{   r   rm   rz   r/   rD  r1   rv   )r  r3   r5   r4   r6   	pd_inputsinputsra  r:   r:   r;   !test_cheap_single_partition_mergeP  s     &r  c                  C   s   t jddddddgtddd	d
ddddgd} tj| dd}t ddddgtdd}tj|ddd}|j|ddd}t|jd rt	|j
rt	t| |j|ddd}|j
rt	t| d S )Nr   r   r   r   r   r   r   r   r   r   r    r!   r"   r#   r$   rR   r&   r  FrT   r   rg   r>  re   )rC   r.   r0   r-   rG   r{   r   rm   rz   r/   rV   r   )r3   r5   r4   r6   r  r:   r:   r;   +test_cheap_single_partition_merge_divisionse  s     

r  c           
      C   s  t jtdtddtdd}tj|dd}t jdddd	gtd
dtd	d}tj|dd}|rl||fn||f}|r||fn||f}tj|dd| d}t j|dd| d}	t|jd	 rt
t||	 tj|dd| d}t j|dd| d}	t|jd	 rt
t||	 d S )Nra   Zababbddar  r$   r   rR   r   r   r   r&   Tr   r   r   rd   re   rC  )rC   r.   r   r0   r-   rG   r{   r   rm   rz   r/   r   )
rd   r  r3   r5   r4   r6   r  r  r  r~   r:   r:   r;   +test_cheap_single_parition_merge_left_righty  s     $
r  c                  C   s&  t jddddddgtddd	d
ddddgd} tj| dd}t ddddgtdd}tj|ddd}|j|dddd}| j|dddd}|jd|_t|j	d
 rt|jrtt|| |j|dddd}|j| dddd}|jd|_t|j	d
 rt|jrtt|| d S )Nr   r   r   r   r   r   r   r   r   r   r    r!   r"   r#   r$   rR   r&   r  FrT   Tr   rg   r  r&  re   )r   r   rd   )rC   r.   r0   r-   rG   r{   r%   r*  r   rm   rz   r/   rV   r   )r3   r5   r4   r6   r  r~   r:   r:   r;   *test_cheap_single_partition_merge_on_index  s&     

r  c                  C   s   t jdddgtddddtdd	} t jd
dgddtddtdd	}tj| dd}tj||dd }t|jdks~t	d S )Nr   r   r   abcr  r,  )rH   rI   r   DZDCBAr   r   r   g      @barab)GHIrI   ZGHIBrR   rI   r   )r  r   rI   rH   r  r  r  )
rC   r.   r0   r-   rG   r{   rk   tuplerq   r/   )rL   rM   r\   rH  r:   r:   r;   test_merge_maintains_columns  s      r  c                 C   s   t jddddddgidddddgd}t jddddddgidddddgd}tj|dd	d
}tj|dd}|j|d| d}|j|dd}t|| d S )Nr   r   r   r   r   r   r$   r   FrT   rR   rg   rB  rc   )rC   r.   r-   rG   rj   r   )r|   r3   r4   r5   r6   r}   r~   r:   r:   r;   "test_merge_index_without_divisions  s    &&r  c                  C   s   t dtjjdddi} t jdtjjddditjjdddd}tj| dd}tj|dd}t j| |d	dd
}tj||d	ddd}t|| t	|j
dk std S )Nr   r)   i  sizer   r$   rR   r   Tr   r   r^  )r   r   ry   i  )rC   r.   r9  r  randintr-   rG   r{   rv   r1   rm   r/   )r3   r4   r5   r6   r@   ra  r:   r:   r;   *test_half_indexed_dataframe_avoids_shuffle  s     
r  c               	   C   s   t jddddddgidddddgd} t jddddddgidddddgd}tj| dd	d
}tj|dd}tt tj||d|jd W 5 Q R X tt tj|||j	|jd W 5 Q R X d S )Nr   r   r   r   r   r   r$   r   FrT   rR   r	  )
rC   r.   r-   rG   rW   rX   NotImplementedErrorr{   r   r   r2   r:   r:   r;   &test_errors_for_merge_on_frame_columns  s    &&r  c                  C   sf   t ddddg} tj| ddd}tj|gdd}t|tjsBttj|gdd}t|tjsbtd S )	Nr   r   r   r   FrT   r   r.  )rC   Seriesr-   rG   r   r,   r/   r.   )r3   r5   r@   r:   r:   r;   test_concat_one_series  s    r  c               	   C   s   t ddddg} t ddddg}tj| ddd}tj|ddd}|jrNttt, t	t j
| |gddtj
||gdd W 5 Q R X tj|ddd}tt tj
||gdd W 5 Q R X tjdd	}tj
||gddd
 W 5 Q R X |rtd S )Nr   r   r   r   FrT   r.  Tr3  )r  Zignore_unknown_divisions)rC   r  r-   rG   rV   r/   rW   warnsUserWarningr   r   rX   rY   r  r  )r3   r4   r5   r6   ra  r4  r:   r:   r;   test_concat_unknown_divisions  s    
0r  c               
   C   s   t ddddddg} t ddddg}tj| ddd}tj|ddd}tt2 tt tj	||gdd	
  W 5 Q R X W 5 Q R X d S )
Nr   r   r   r   r   r   FrT   r.  )rC   r  r-   rG   rW   rX   rY   r  r  r   rk   r2   r:   r:   r;   $test_concat_unknown_divisions_errors  s    r  c                  C   s$  t dddgdddgdt dddgdddgdt dd	d
gdddgdd} tdddt  d}t| d|d d g}t dddgdddgdt dddgdddgdt dddgdddgdd} t| d|d d g}t dddgdddgdt dddgdddgdd} tdddt  d}t| d|d d g}t dddgdddgdddgdt jdddgdddgdddgddddgdd} tddddt g dt  d}t| d|dddg}||g||g||gg}t|g|kst|D ]}d d! |D }t j|d"d#}	t|}
|
j|d j|d j ksRt|
jd$|
jd  ksltt	|	|
 t
|
jt
t|jkstt j|d%d"d&}	tj|d%d'}
|
j|d j|d j kst|
jd$|
jd  kstt	|
|
 t
|
jt
tj|d%d'jkstqd S )(Nr   r   r   r   r   r   r>   r=   ra   r   r   )r   r   )r   r   )r   r   r+   parent_metar   r   r   r    r!   r"   r#   r'   r(   Z   r   r   r   r   )r   r   r   r4   r@   r  r  )r4   r@   rA   r$   )r%   r  c                 S   s   g | ]}|  qS r:   rk   )r   Z_cr:   r:   r;   r   9  s     z test_concat2.<locals>.<listcomp>FrU   r   rg   rj   rU   r  )rC   r.   r   r-   ZIndexr   r/   rS   r9   r   setrm   )dskrG  r3   r4   r@   rA   casescasepdcaser~   r}   r:   r:   r;   test_concat2  s\     



 
  
r  c                  C   s  t jtjddtdtdd} t jtjddtdtdd}t jtjddtdtd	d}t| d
}t|d}t|d
}t j| |gdd}t||g}|j	|j	d d |j	 kst
|j|j|j kst
t|| ttj||gddt | |g t j| ||gdd}t|||g}|j	|j	d d |j	d d  |j	 ks\t
|j|j|j |j kszt
t|| ttj|||gddt | ||g d S )Nr   r   ABCDEr`   rq   r%   ZABCFGZghijklZABCHIZmnopqrr   r   Fr  re   Tinterleave_partitions)rC   r.   r9  r  r  r0   r-   rG   r   r9   r/   rS   r   )r  r  r  r   r]   r   r~   r}   r:   r:   r;   test_concat3I  sP          
  


r  c               	   C   sZ  t jtjddtdtdd} t jtjddtdtdd}t jtjddtd	tdd}t| d
}t|d}t|d
}d}||g||g||g||g||g||g||gg}|D ]R}dd |D }	ttj	|ddt j	|	dd ttj	|dddt j	|	ddd qd}t
t}
tj	||gddd W 5 Q R X |t|
jksVtd S )Nr   r   r  r  r  r   Zfghijklmnopqrr   ZCDEXYZr   r   zyAll inputs have known divisions which cannot be concatenated in order. Specify interleave_partitions=True to ignore orderc                 S   s   g | ]}|  qS r:   r  r   r@   r:   r:   r;   r     s     z6test_concat4_interleave_partitions.<locals>.<listcomp>Tr  Fr  rg   rj   r  r  z!'join' must be 'inner' or 'outer'invalid)rC   r.   r9  r  r  r0   r-   rG   r   r   rW   rX   rY   r6  valuer/   )r  r  r  r   r]   r   msgr  r  r  errr:   r:   r;   "test_concat4_interleave_partitionsq  sT          	 r  c                  C   s  t jtjddtdtdd} t jtjddtdtdd}t jtjddtdtdd}t jtjddtd	tdd}t jtjddtd	td
d}t| d}t|d}t|d}t|d}t|d}	||g||g||g||	g||g||	g|	||g|	|g|j|jg|j	|j	g|j|	jg|j|j	g|j	|j
g|j|	jg||jg|j	|g|	|j|gg}
|
D ]}dd |D }ttj|ddt j|dd ttj|dddt j|dd ttj|ddt j|dd ttj|dddt j|ddd qv||g||g| |g| j|jg|j	|j	g||jg|j	|g|| |j	gg}
|
D ]}dd |D }ttj|ddt | ttj|dddt j|dd ttj|ddt j|dd ttj|dddt j|ddd qRd S )Nr=   r   r  rB   r  r   ZFGHIJKrf  ZFGHABZfklmnopr   r   c                 S   s   g | ]}|  qS r:   r  r  r:   r:   r;   r     s     z test_concat5.<locals>.<listcomp>Tr  Fr  rg   r  r  r   r.  )r  rj   c                 S   s"   g | ]}t |tr| n|qS r:   )r,   r   rk   r  r:   r:   r;   r     s     )rC   r.   r9  r  r  r0   r-   rG   rH   Fr  rI   r   r   )r  r  r  Zpdf4Zpdf5r   r]   r   Zddf4Zddf5r  r  r  r:   r:   r;   test_concat5  s              









r  zknown, cat_index, divisions)TTF)TFT)TFF)FTF)FFT)FFFc           	         s  t tdtdtdtjddddt tdtddtd	tjdddddt td
tddtdtjdddddg}|D ] }|jd|_|jd|_q rdd |D }fdd|D }st|d j	dgdd|d _	 fdd}dD ]}||||}t
|js(tt
|jks<t|dd |D dd |D |}t
|ksnt r|dd |D dd |D |}||d ddg g|dd   |d ddg g|dd   |}t|drt
|jstt
|jkstqd S )NZxxxxxr   Zabcbcf8r'  )wr   r   rx   Zyyyyyr   abbbaZzzzzzr   bcbcccategoryc                 S   s   g | ]}| |jqS r:   )r  r   )r   r[   r:   r:   r;   r   (  s     z+test_concat_categorical.<locals>.<listcomp>c                    s   g | ]}t j|d  dqS )r   rT   r-   rG   r   pr   r:   r;   r   *  s     r   r   Tr$   c                    sx   t ||d}tj | |d}t|| r^ttj|j| }dd |D D ]}t|j| qL rtt	|j
kstt|S )Nr  r  c                 S   s   g | ]}|j d d qS rh  )r   r   ir:   r:   r;   r   7  s     zEtest_concat_categorical.<locals>.check_and_return.<locals>.<listcomp>)r   r-   r   r   r.   rm   Z__dask_keys__r   r)  r   r%   r/   )Zddfsdfsrj   ZsolresrO   r  	cat_indexr9   knownr:   r;   check_and_return/  s    
  z1test_concat_categorical.<locals>.check_and_return)rg   rh   c                 S   s   g | ]
}|j qS r:   )r   r  r:   r:   r;   r   C  s     c                 S   s   g | ]
}|j qS r:   r   r  r:   r:   r;   r   H  s     r   r   r  )rC   r.   r0   r9  r:  r  r*  r   r   r)  r   r/   hasattr)	r  r  r9   framesr[   dframesr  rj   r  r:   r  r;   test_concat_categorical  s\    

  r  c                  C   s   t jdddgdd} t jddgdd}t| dj d}t|dd}t| d|dg}t||g}t|| d S )Nr3   r4   r@   r  r'  r   rH   )	rC   r  r-   rG   catZ
as_unknownZto_framer   r   )r3   r4   r  dbr~   r}   r:   r:   r;   $test_concat_categorical_mixed_simpleT  s    r  c                  C   s   t jddgit jdgddd} t jddgit jdgddd}| jd	jddg| d< |jd	jddg|d< t| d
}t|d
}t	| j
d d |j
d d g}|jjdkstt	||g}t 	| |g}t|| d S )Nr   r3   z2015-03-24 00:00:16zdatetime64[ns]r'  r$   r@   z2015-03-29 00:00:44r  r   r   zM8[ns])rC   r.   ZDatetimeIndexr   r*  r  set_categoriesr-   rG   r   r   r%   r(  r/   r   )Zb2Zb3Zdb2Zdb3r}   r~   r:   r:   r;   test_concat_datetimeindex_  s"     r  c              	   C   s`   t r>tjtdd$ ||}| |}t|| W 5 Q R X n||}| |}t|| |S )Nzappend method is deprecatedr   )r   rW   r  r  appendr   )Zdask_objZdask_appendZ
pandas_objZpandas_appendr~   r}   r:   r:   r;   check_append_with_warningx  s    




r  zpandas removed appendc               	   C   sv  t ddddddgddddddgd} t jddddddgddddddgdddd	d
ddgd}t jddddddgddddddgdddd	d
ddgd}t| d}t|d}t|d}t jdd	gdddgd}t||| | t||| | t|j|j| j|j t||| | t|j|j| j|j t||| | t|j|j| j|j t||| | t|j|j| j|j d S )Nr   r   r   r   r   r   r>   r=   ra   r   r   r   r$   r  r3   r4   )rT  r%   )rC   r.   r-   rG   r  r  r3   r4   )r[   r   r   r\   r]   r   r*   r:   r:   r;   test_append  s,    (    r  c                  C   s  t dddgdddgdt dddgdddgdt dd	d
gdddgdd} tdddt  d}t| d|d d g}| }t dddgdddgdt dddgdddgdt dddgdddgdd} t| d|d d g}| }t dddgdddgdt dddgdddgdd} tdddt  d}t| d|d d g}| }t|||| t|||| t|||| t|||| t|||| t|||| t|||| t|||| d S )Nr   r   r   r   r   r   r>   r=   ra   r   r   r  r+   r  r   r   r   r    r!   r"   r#   r'   r(   r  r  r   r  r  )rC   r.   r   r-   rk   r  )r  rG  r   r   r]   r   r   r   r:   r:   r;   test_append2  s6    r  c            
   
   C   s  t tddtdtjdddddt tddtdtjdddddg} g }| D ]H}|jd	|_| }|jj	td
|_|j|_
|||j q\|\}}dD ]}dd | D }|st|d jdgdd|d _|\}}t||||}	t|	j
|kstt|	j|ks"tt|j|j|j|j}	t|	j
|ksLtt|	|ks^tt|j
|j
|j
|j
}	t|	|kstqd S )Nr   r   r  r  r'  )r   r   rx   r   r  r  r  )TFc                 S   s   g | ]}t j|d ddqS )r   FrT   r  r  r:   r:   r;   r     s     z+test_append_categorical.<locals>.<listcomp>r   r   Tr$   )rC   r.   r9  r:  r0   r   r*  r   r  r  r%   r  r  r   r)  r  r   r/   )
r  Zframes2r[   r   r   r  r  r   r]   r  r:   r:   r;   test_append_categorical  sL    

  r  c                  C   sP   t jdddddgiddddgd} tj| dd}t||| | }|jdksLtd S )	Nr   r   r   r   r   r$   rR   F)rC   r.   r-   rG   r  rV   r/   )r[   r\   r  r:   r:   r;   test_append_lose_divisions	  s    "r  c                  C   s^   t jddddgidddgd} tj| dd}|d}|j|dd}|jd	ksRt|  d S )
Nr   r   r   r   r$   rR   r7  )r   )r   r   )	rC   r.   r-   rG   r  rj   r9   r/   rk   )r[   r\   r]   Zjoinedr:   r:   r;   test_singleton_divisions	  s    
r  c                  C   sL   t dddddgi} tj| ddd}|jddd}t|| d d S )Nr   r   r   rR   )r   r   T)r9   force)rC   r.   r-   rG   r  rD   r   rZ   r:   r:   r;   #test_repartition_repeated_divisions	  s    r  c                  C   s   t dddddgi} t ddgi}tj| ddd}tj|ddd}|jdks\tt|t| kspt|j|dddd}| d}|d}|j|dddd}t	|| d S )	Nr   r   r   rR   r   rb   Trm  )
rC   r.   r-   rG   r  rS   r/   r1   r{   r   )r   r   r   r]   Zr1Zsf1Zsf2Zr2r:   r:   r;   test_multi_duplicate_divisions	  s    

r  c               	   C   s   d} t ddddddgddd	d	ddgd
}tj|dd}tjt  dd}td| d D ]4}ttj|||j|k dd||j|k ddd qZd S )Nr   rH   rI   r   r  Er  r   r   )userclusterr   rR   r   rh   rc   F)r   Zcheck_divisions)rC   r.   r-   rG   r   r   r{   r  )Z
k_clustersr[   r  r   r:   r:   r;   test_merge_outer_empty/	  s     r  c               	   C   s~   t dtjddgttjdi} t dtjddgttjdi}tjdd}t	j
| |dd W 5 Q R X |rztd S )Nr3   r   r   r'  Tr3  )rC   r.   r9  arrayr(  r&  Zlonglongr  r  r-   ZmultiZwarn_dtype_mismatchr/   )r   r   r4  r:   r:   r;   test_dtype_equality_warningA	  s
    $$r   c              	   C   sZ  d}d}t tjt|tjjd|dd}t tjt|tjjd|dd}| dkrtd}td}|	|}|	|}|
||}|
||}nt	||}t	||}|d	g }	|d
g }
t|	|
gdd}|d	g }|d
g }t $ tdt tj||gdd}W 5 Q R X t|  |  d S )Nr   r   r)   r  r>   r?   rN  rO  r3   r@   r   r.  r  )rC   r.   r9  r  Zpermutationr:  r  rW   rZ  rG   r[  r-   groupbysumr   r  r  r  r  r   rk   Z
sort_index)rK  r  rS   d1Zd2rN  rO  r  r  Z
grouped_d1Z
grouped_d2r  Zgrouped_dd1Zgrouped_dd2Zres_ddr:   r:   r;   test_groupby_concat_cudfK	  s<    




r  orderedc                 C   s   t dt jddddgdddg| di}t|d}t dt jdddgdddg| di}t|d}t ||g}|d d|d< tj||gdd	}t|| d S )
Nr   r3   r4   r@   )
categoriesr  r   r  T)Zignore_order)rC   r.   Categoricalr-   rG   r   r*  r   )r  r  r   r  r]   r~   r}   r:   r:   r;   test_concat_ignore_order}	  s0     
     r  c                  C   s   t ddddgddddgd} t jdddddgiddddgd}tj| dd	}tj|dd	}|d
 d|d
< |jd|_| j| d
dd}|j|d
dd}|jj	dkst
| }|jj	dkst
t||st
d S )Nr3   r4   r   r   )join_colr3   r   r   r$   rR   r	  r  rb   r>  )rC   r.   r-   rG   r*  r%   rk   rj   r	  r(  r/   r   )r[   r   r\   r]   r~   actual_daskr  r:   r:   r;   test_categorical_join	  s     

"r  c                  C   s   t ddgt ddgd} t dt ddgi}t j|| ddd	}tj| d
d}tj|d
d}tj||ddd	 }|jjdkst	|j
jdkst	|jjdkst	|jjdkst	t||st	d S Nr   r   r3   r4   r2  r   TrH   r  r   rR   r  r&  rC   r.   r  r{   r-   rG   rk   r   r(  r/   rI   rH   r%   r   r   r   r~   r   r]   r  r:   r:   r;   5test_categorical_merge_with_columns_missing_from_left	  s    r  c                  C   s   t t ddgt ddgd} t dt ddgi}t j|| ddd	}tj| d
d}tj|d
d}tj||ddd	 }|jjdkst	|j
jdkst	|jjdkst	|jjdkst	t||st	d S r  r  r  r:   r:   r;   Itest_categorical_merge_with_merge_column_cat_in_one_and_not_other_upcasts	  s    $r  c                  C   s   t ddddgddddgd	} t ddddgddddgd
}t| d}|jd|d< t|d}|jd|d< |j|dd}|jjdkst|	 }|jjdkstd S )Nr   r   r   r   r   r   r   r=   r2  )rH   r   r  rH   r   )
rC   r.   r-   rG   rH   r*  r{   r(  r/   rk   )r3   r4   r   r   r
  r  r:   r:   r;   -test_categorical_merge_retains_category_dtype	  s      r  c                  C   sr   t jddddgidddgd} t jd	d
ddgidddgd}tj| dd}|d	di}t| ||| d S )NrH   r3   r4   r@   r*   vr  )datar%   rI   trA   r  r7  r   rR   r  )rC   r.   r-   rG   r*  r   rj   )r   r   r   r:   r:   r;   ?test_categorical_merge_does_not_raise_setting_with_copy_warning	  s
    r  rS   r       baselgsmc              	   C   s0  d}d}d}d}d}t tjtd|t|d}tj||d}	t tjtd|t|d}
tj|
|d}|d	kr|}|	}|
}|}n|
}|}|}|	}tj||d
| ||dd}t j||d
| d}|d
 	tj
|d
< |d
 	tj
|d
< |r|j|kstt| d
|d
dd d S )Ni  r%  r    r   r,  r)   r   rR   r  r   r^  )r   rd   rS   	broadcastry   r>  Fr   )rC   r.   r9  r  choicer:  r-   rG   r{   r*  Zint32rS   r/   r   rk   rr   )rd   rS   r  Zsize_lgZsize_smZnpartitions_lgZnpartitions_smZbroadcast_biasr  Zddf_lgr  Zddf_smrb   Zddf_leftrf   Z	ddf_rightZ	dd_resultZ	pd_resultr:   r:   r;   test_merge_tasks_large_to_small	  sZ    	r  c              	   C   s   t tjdddgdddgddddgd	d}t tjdddgdddgdddd
gd	dt tjdd
dgdddgddddgd	dg}tt}|j|| d W 5 Q R X |d d S )Nr   r   r   r   r   r   r>   r   r$   ra   r=   r   rc   z-merge_multi only supports left or outer joins)	r-   rG   rC   r.   rW   rX   rY   rj   r   )rd   Zbase_dfr  r  r:   r:   r;   ,test_pairwise_rejects_unsupported_join_types,
  s     $ $ $ 	r  npartitions_baser   r   npartitions_otherc           	   	   C   s   g }t dD ]T}tj| dddddg| ddd	d
dgiddd
d	gd}t||}|| qttjddd	gd|}|D ]}|j|| d}qttjddd	gd|}|j|| d}t|| d S )Nr   rH   r   r   r=   ra   rI   r   r   r   r   r   r$   rc   )r   rC   r.   r-   rG   r  rj   r   )	rd   r  r  Zdfs_to_merger  r[   r\   Zddf_loopZddf_pairwiser:   r:   r;   2test_pairwise_merge_results_in_identical_output_df?
  s$     
 

r   )wr  Znumpyr9  rL  rC   rW   Zdask.dataframeZ	dataframer-   Z	dask.baser   Zdask.dataframe._compatr   r   r   Zdask.dataframe.corer   Zdask.dataframe.methodsr   Zdask.dataframe.multir   r	   r
   r   r   Zdask.dataframe.utilsr   r   r   r   r   r   Zdask.utils_testr   rP   r^   r_   rn   rv   markZparametrizer   r   r   r   r   r   r   r   r   r   r   r   r   r  r
  r  r  r   r!  r#  r+  r   r0  r=  rE  rI  rJ  paramZxfailZgpur]  rb  Zslowr  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  Zskipifr  r  r  r  r  r  r  r  r   r  r  r  r  r  r  r  r  r  r   r:   r:   r:   r;   <module>   sF   L-#
"?::H4 




#
Q
& V
c
 !
8(/]
L

*
0


 
/
<
