U
    /eM                     @   sl  d dl Z d dlZ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 d dlm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mZmZmZ d dlZd dlmZ d dlmZm Z m!Z!m"Z"m#Z#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/ d dl0m1Z1m2Z2m3Z3 d dl4m5Z5m6Z6m7Z7m8Z8 e9dfe9dfe9dfdZ:e;e9dd Z<ee:ddZ=dd Z>dd Z?dd Z@dd ZAdd ZBdd ZCdd ZDd d! ZEd"d# ZFd$d% ZGd&d' ZHd(d) ZIejJKd*eLeMgd+d, ZNd-d. ZOd/d0 ZPd1d2 ZQd3d4 ZRd5d6 ZSd7d8 ZTd9d: ZUd;d< ZVd=d> ZWd?d@ ZXejJKdAdBdCgdDdE ZYdFdG ZZdHdI Z[dJdK Z\dLdM Z]dNdO Z^ejJKdAdBddPgdQdR Z_ejJKdAdBdSgdTdU Z`dVdW ZadXdY ZbejJKdZeejcd[d\ gd]d^ Zdd_d` Zedadb Zfdcdd Zgdedf Zhdgdh Zididj Zjdkdl Zkdmdn Zldodp Zmdqdr Zndsdt Zodud\ Zpdvdw Zqdxdy Zrdzd{ Zsd|d} Ztd~d Zudd Zvdd Zwdd ZxejJjyejJjzejJj{dddd Z|dd Z}dd Z~dd Zdd ZejJjddejJjyejJjzdd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zdd ZdefdefdefgZejJKdedd Zdd Zdd Zdd Zdd Zdd Zdd ZddĄ ZddƄ ZddȄ Zddʄ ZG dd̄ dejZdd΄ ZddЄ Zdd҄ ZddԄ Zddք Zdd؄ Zddڄ Zdd܄ ZejJKdAdBddSdgd<ddZejJKddBdCdddgejJKddBdCdddgdd ZejJKddddddddddg	dd Zdd Zdd Zdd Zdd Zdd Zdd  Zdd ZejJKddddgdd Zd	d
 Zdd Zdd ZejJKdAdBdCdPgdd ZG dd deZdd Zdd Zdd Zdd Zdd Zdd Zdd  Zd!d" Zd#d$ Zd%d& Zd'd( Zd)d* ZÐd+d, ZĐd-d. ZŐd/d0 ZƐd1d2 Zǐd3d4 ZejJKd5d6d7gd8d9 Zɐd:d; ZdS (=      N)BZ2File)Iterator)ProcessPoolExecutor)GzipFile)repeat)groupbyidentityjoinmergepluckuniquevalmap)Bagcollectfrom_delayedinline_singleton_listslazifylazify_taskoptimize	partitionreducebyreifytotal_mem_usage	assert_eq)	Blockwise)Delayed)	filetextstmpdirtmpfile)add	hlg_layerhlg_layer_topologicalinc   )xr   r&      )r&         r&   c                 C   s   | d dkS Nr)   r    r&   r,   r,   ;/tmp/pip-unpacked-wheel-dbjnr7gq/dask/bag/tests/test_bag.pyiseven-   s    r/   c                 C   s   | d dkS )Nr)   r(   r,   r-   r,   r,   r.   isodd1   s    r0   c                   C   s    t jdkstt jdkstd S )Nr&   r*   )bnameAssertionErrornpartitionsr,   r,   r,   r.   test_Bag5   s    r5   c                   C   s   t  tt kstd S N)r1   __dask_keys__sorteddskkeysr3   r,   r,   r,   r.   	test_keys:   s    r;   c                  C   s<   t t } | dddgd fddddgd fgks8td S )NFr(   r*   Tr   r)      )r1   r   r/   computer3   resultr,   r,   r.   test_bag_groupby_pure_hash>   s    r@   c                  C   sX   t dd  } t| dks"tdddgd f| ks:tdddd	gd f| ksTtd S )
Nc                 S   s   t | rdS dS )Nevenodd)r/   r-   r,   r,   r.   <lambda>F       z.test_bag_groupby_normal_hash.<locals>.<lambda>r)   rB   r(   r*   rA   r   r<   )r1   r   r=   lenr3   r>   r,   r,   r.   test_bag_groupby_normal_hashD   s    rF   c               	      s  t jtddd} t jtdddd}|  }| }ddd	 tt  | tt | tt j | d
tt | tt  | |tt || tt  | d fdd|D  tt j d| d fdd|D   fddt||D }tt j | |dd|  fddt||D }tt j | dd| t| fdd|D }tt j |  |dd|  fdd|D }tt j ||  dd|  fdd|D }tt j d|  |d| t	
d}tt  | | fdd|D  tt j | |d fdd|D  t jtddd}tt t  | | W 5 Q R X tt t  |  dd W 5 Q R X t jtddd}tt t j | ||d  W 5 Q R X tt t j | ||d  W 5 Q R X d S )Nd   
   r4      r(   r)   r*   c                 S   s   | | | S r6   r,   ar1   cr,   r,   r.   myaddR   s    ztest_bag_map.<locals>.myaddrL   c                    s   g | ]} |d qS rH   r,   .0irN   r,   r.   
<listcomp>X   s     z test_bag_map.<locals>.<listcomp>r1   c                    s   g | ]} d |dqS rH   rV   r,   rQ   rT   r,   r.   rU   Y   s     c                    s   g | ]\}} ||d dqS rG   r1   rM   r,   rR   rS   jrT   r,   r.   rU   [   s     rY   c                    s   g | ]\}} |d dqS )rG   rM   r,   rZ   rT   r,   r.   rU   ^   s     r\   c                    s   g | ]} |d dqS rX   r,   rQ   rN   Zx_sumr,   r.   rU   b   s     c                    s   g | ]} |d dqS rX   r,   rQ   r]   r,   r.   rU   e   s     c                    s   g | ]} d |dqS )rG   rK   r,   rQ   r]   r,   r.   rU   h   s     rK   c                    s   g | ]} |d qS rP   r,   rQ   rT   r,   r.   rU   l   s     c                    s   g | ]} |d dqS rW   r,   rQ   rT   r,   r.   rU   m   s     r$   n   )r(   r)   r*   )dbfrom_sequenceranger=   r   maplistzipsumdaskdelayedpytestraises
ValueError)r1   b2r&   x2solrL   Zfewer_partsZunequalr,   r]   r.   test_bag_mapL   sD    
 "$
"$ rn   c                     sB  t jtddd} t jtdddd}|  }| }ddd |   tt |ks`t|  | tt ||kst|  d  fd	d
|D kst| j dd  fdd
|D kst| j |dd  fdd
t||D kstt|| j |  dd  fdd
|D ks>td S )NrG   rH   rI   rJ   r)   r*   c                 S   s   | | | S r6   r,   rK   r,   r,   r.   rN      s    ztest_map_method.<locals>.myaddc                    s   g | ]} |d qS rP   r,   rQ   rT   r,   r.   rU      s     z#test_map_method.<locals>.<listcomp>rV   c                    s   g | ]} |d dqS rW   r,   rQ   rT   r,   r.   rU      s     r\   c                    s   g | ]\}} ||d qS rP   r,   rZ   rT   r,   r.   rU      s    c                    s   g | ]} |d qS rP   r,   rQ   r]   r,   r.   rU      s     )r)   r*   )	r_   r`   ra   r=   rb   rc   r3   rd   re   )r1   rk   r&   rl   r,   r]   r.   test_map_method   s    
 $&(ro   c                     s   dddddg} t j| dd}dd	d
| fdd| D ksJt|jdd fdd| D ksrt|d  |j d  fdd| D ksttd}|j|d fdd| D kstd S )Nr(   r)   )r*   r<   )r$      )      )	   rH   r)   rI   r   c                 S   s   | | | S r6   r,   rK   r,   r,   r.   rN      s    ztest_starmap.<locals>.myaddc                    s   g | ]} | qS r,   r,   rR   rL   rT   r,   r.   rU      s     z test_starmap.<locals>.<listcomp>rH   r\   c                    s   g | ]} |d diqS rM   rH   r,   ru   rT   r,   r.   rU      s     r(   c                    s   g | ]}|d    iqS r\   r=   ru   Z
max_secondrN   r,   r.   rU      s    c                    s   g | ]} |d diqS rv   r,   ru   rT   r,   r.   rU      s     )r   )	r_   r`   starmapr=   r3   r   maxrf   rg   )datar1   rM   r,   rx   r.   test_starmap   s    
$(
r|   c                     sP   t t tt fddtt jD }  j| ks6t j	t tj	ksLtd S )Nc                    s(   i | ] } j |fttttj |fffqS r,   )r2   r   filterr/   r1   rQ   r\   r,   r.   
<dictcomp>   s    ztest_filter.<locals>.<dictcomp>)
r1   r}   r/   r
   r9   ra   r4   rf   r3   r2   )expectedr,   r\   r.   test_filter   s    

r   c                  C   sD   dd } t | }t|ddgd ks*t|jt | jks@td S )Nc                 S   s   | d dkS r+   r,   r-   r,   r,   r.   rC      rD   ztest_remove.<locals>.<lambda>r(   r*   )r1   removerc   r3   r2   )frM   r,   r,   r.   test_remove   s    
r   c                   C   sH   t ttt tkstt tttt ttddd ksDtd S )Nr(   rq   r*   )r8   rc   r1   Lr3   rb   r#   ra   r,   r,   r,   r.   	test_iter   s    r   funcc                 C   sN   t tj| tksttjd d | tks0td| ttdksJtd S )Nr$   r`   )strr1   r4   r3   r2   r_   r`   ra   )r   r,   r,   r.   	test_repr   s    r   c                  C   s   ddgddgd} t | dd}t|dd	dd
dhks<tt|d	ddddhksZtt|d	dgddddhks|t|d	dgj|d	dgjkstd S )Nr(   rH   )r)      )r*      )r<   (   )r%   r'   r&   r)   r   r(   r*   r<   rH   r   r   r   )rH   r(   )r   r)   )r   r*   )r   r<   )r   setr   r3   r2   dr1   r,   r,   r.   
test_pluck   s    "r   c                      s   t dddg tt fdd t dd dd dgksBt dd j dd jksbt dj dd jkstd S )	NZHello ZWorldc                      s   t  dS Nr   )rc   r   r,   rV   r,   r.   rC      rD   z)test_pluck_with_default.<locals>.<lambda>r   HW)	r_   r`   rh   ri   
IndexErrorrc   r   r3   r2   r,   r,   rV   r.   test_pluck_with_default   s
     r   c                  C   s   t tddd } | d\}}}t|ttdks@tt|dd tdD ks^t|j| dd jksxt|j|jkstd S )NrG   c                 S   s   | | d | d fS )Nr(   r)   r,   r-   r,   r,   r.   rC      rD   ztest_unzip.<locals>.<lambda>r*   c                 S   s   g | ]}|d  qS )r)   r,   rQ   r,   r,   r.   rU      s     ztest_unzip.<locals>.<listcomp>r   )r_   r`   ra   rb   Zunziprc   r3   r2   )r1   ZoneZtwoZthreer,   r,   r.   
test_unzip   s    r   c                  C   sf  t t} |  ttkst| jt tjks4tt jtdd}|j| jksRt| ttdt j  kspt|jt jtddjkstt	j
tddd} dd }| j|tjt d}| t| kst|j| j|tjt djkstt	
d}t|jd	d
 dd tdkstt	j
dgdgdggdd}t|jtg djdddddhksbtd S )NrH   )initialr$   r*   rI   c                 S   s   |   } | | | S r6   )copyr    accr&   r,   r,   r.   binop   s    
ztest_fold.<locals>.binopZhelloc                 S   s   d | |gS )Nr   )r	   rL   r1   r,   r,   r.   rC      rD   ztest_fold.<locals>.<lambda>r   r(   r)   sync	scheduler)r1   foldr    r=   re   r   r3   keyr4   r_   r`   ra   r   union)rM   c2r   r   er,   r,   r.   	test_fold   s$    
 
r   c                  C   sF   dd } t j| tjt td}t|ts,tt|ttt	d d S )Nc                 S   s   |  | | S r6   )r    )Ztotr&   r,   r,   r.   r      s    
ztest_fold_bag.<locals>.binop)combiner   Zout_typer$   )
r1   r   r   r   r   
isinstancer3   r   rc   ra   )r   rM   r,   r,   r.   test_fold_bag   s    r   c                  C   s   t t dddddgkstt jt jks6tdt jksHtt   dks`ttjdgd dd	} | 	d   g kstd S )
Nr   r(   r)   r*   r<   distinctr$   2   rI   )
r8   r1   r   r3   r2   countr=   r_   r`   r}   bagr,   r,   r.   test_distinct  s    r   c                  C   s\   dd dD } t j| dd}tt| dd d}t|jd	d| t|jd
d d| d S )Nc                 S   s   g | ]}d |iqS rO   r,   rQ   r,   r,   r.   rU     s     z*test_distinct_with_key.<locals>.<listcomp>)
r   r(   r)   r(   r)   r*   r)   r*   r<   r$   r*   rI   c                 S   s   | d S NrL   r,   r-   r,   r,   r.   rC     rD   z(test_distinct_with_key.<locals>.<lambda>r   rL   c                 S   s   | d S r   r,   r-   r,   r,   r.   rC     rD   )r_   r`   rc   r   r   r   )seqr   r   r,   r,   r.   test_distinct_with_key  s
    r   c                  C   s
  t  } t| ddddddks$tt jdd}t|ddddddksLt| jt  jks`t| j|jkspt|jt jddjksttjtddd}|dd }| }t|tt	td	d
gd	 ksttjddddgdd}|d jdd}t
|g  d S )Nr*   r   r(   r)   r*   r<   r)   split_everyr   partition_sizec                 S   s   | dk S NrH   r,   r-   r,   r,   r.   rC   #  rD   z"test_frequencies.<locals>.<lambda>rH   r(   r   r<   rI   )r1   frequenciesdictr3   r2   r_   r`   ra   r}   rd   r   )rM   r   rk   r   r   Zbag2r,   r,   r.   test_frequencies  s    $r   c                  C   s<   t ddddddg} t| jdd dddgks8td S )	NrL   r1   rM   T)sort)r1   r*   )rM   r)   rL   r(   )r_   r`   rc   r   r=   r3   rV   r,   r,   r.   test_frequencies_sorted+  s    r   c                  C   s   t tdddddgksttjddd d} t | ddddgksHttjddd d	d
}t |ddddgkstt| j|jksttdjtdjkstd S )Nr<   r*   c                 S   s   |  S r6   r,   r-   r,   r,   r.   rC   2  rD   ztest_topk.<locals>.<lambda>r   r   r(   c                 S   s   |  S r6   r,   r-   r,   r,   r.   rC   4  rD   r)   )r   r   )rc   r1   topkr3   r2   )rM   r   r,   r,   r.   	test_topk0  s    r   r4   r(   r)   c                 C   sx   t jdddg| d}t|jdddddgks2tt|jdddddgksPt|jdddj|jdddjksttd S )	Nr   r)   rt   r*   rs   rI   r)   r(   r   r   )r_   r`   rc   r   r3   r2   r4   r1   r,   r,   r.   test_topk_with_non_callable_key:  s    r   c                  C   s:   t jdddgdd} t| jddd dddgks6td S )	Nr   r   r   r)   rI   c                 S   s   |S r6   r,   r   r,   r,   r.   rC   D  rD   z0test_topk_with_multiarg_lambda.<locals>.<lambda>r   )r_   r`   rc   r   r3   rV   r,   r,   r.   test_topk_with_multiarg_lambdaB  s    r   c                   C   s(   t tdd t ttks$td S )Nc                 S   s   | d S Nr(   r,   r-   r,   r,   r.   rC   H  rD   ztest_lambdas.<locals>.<lambda>)rc   r1   rb   r#   r3   r,   r,   r,   r.   test_lambdasG  s    r   c                   C   s   t t dkstt t dks(tt t dks<tt t dksPtt  dksdtt	  dksxtt	 j
t	 j
kstt	 j
t j
kstd S )N   r   r<   r   TF)intr1   r   r3   re   rz   minanyr=   allr   r,   r,   r,   r.   test_reductionsK  s    r   c                   C   sl   t  jdstt ttjds,ttdd t ttjD sLtt jttddjdshtd S )Nre   c                 s   s"   | ]}t |to|d V  qdS )rz   N)r   r   
startswithrR   kr,   r,   r.   	<genexpr>Y  s    z'test_reduction_names.<locals>.<genexpr>foo)r2   )	r1   re   r2   r   r3   	reductionr   rz   rf   r,   r,   r,   r.   test_reduction_namesV  s    r   c                  C   s   t td} | jttdd}| jttdd}| jttdd}| |   kr^| ksdn tt|jt|jks|t| jdd}| jdd}| | kstt|jt|jkst|j	|j	kst|j	| jddj	kst|j	|  j	kstd S )N   r)   r   rq   r$   )
r_   r`   ra   r   re   r=   r3   rE   rf   r   )r1   rM   r   r   r,   r,   r.   test_tree_reductions_  s    &r   r<   c                 C   s\   t td}tjd| d}t| t|t|  t| t| t| t| d S )Nr   rI   )rc   ra   r_   r   meanre   rE   r   )r4   r   r1   r,   r,   r.   test_aggregationt  s
    r   rH   c                 C   sP   t d}ttd}tj|| d}t| || t| || d S )NnumpyrG   rI   )	rh   importorskiprc   ra   r_   r`   r   r   std)r4   npr{   rM   r,   r,   r.   test_non_splittable_reductions}  s
    
r   c                   C   s2   t t td tt tdks.td S Ng       @)r   r1   r   mathsqrtfloatr3   r,   r,   r,   r.   test_std  s    r   c                   C   s&   t t d tt dks"td S r   )r   r1   varr   r3   r,   r,   r,   r.   test_var  s    r   	transformc                 C   s   t j| ddS )Nr(   rI   )r_   r`   r-   r,   r,   r.   rC     rD   rC   c              	   C   s   | dddg}t j|ttd}t|tttdddgttt  tt |ttttdddgttt  |jt j|ttdjkstd S )Nr(   r)   r*   )Zon_selfZon_other)r1   r	   r0   r/   r   rc   r2   r3   )r   otherrM   r,   r,   r.   	test_join  s
    "*r   c                  C   s   t ttdtd} tttt jdfdft| j ks8t	t
| t
ttdd td ks^t	| jt ttdtdjks|t	t tdd } t
| t
ttdd td kst	d S )Nr   c                 S   s   | | S r6   r,   r   r,   r,   r.   rC     rD   ztest_foldby.<locals>.<lambda>c                 S   s   | | S r6   r,   r   r,   r,   r.   rC     rD   c                 S   s   | | S r6   r,   r   r,   r,   r.   rC     rD   )r1   foldbyr/   r    r   r2   rc   rf   valuesr3   r   r   itemsr\   r,   r,   r.   test_foldby  s    &&r   c            
   
      s   t  } dD ]}tjtd|d}|tt}| |g dD ]b}|jtt|d}|jttd|d}|jttdt|d}|jttdtd|d}| ||||g q<q
t	| }	|	d  t
 fdd|	D std S )	N)r(   rr       rG   rI   )FNr)   r*   r   r   c                    s   g | ]}| kqS r,   r,   )rR   rfirstr,   r.   rU     s     z.test_foldby_tree_reduction.<locals>.<listcomp>)rc   r_   r`   ra   r   r/   r    extendrf   r=   r   r3   )
r9   nr1   rM   mr   r   r   gresultsr,   r   r.   test_foldby_tree_reduction  s    
r   c                   C   s\   t ttdddgkstttjttjks8ttdd jttjksXtd S )Nr$   c                 S   s   t | d S r   )rE   rO   r,   r,   r.   rC     rD   z%test_map_partitions.<locals>.<lambda>)rc   r1   map_partitionsrE   r3   r2   r,   r,   r,   r.   test_map_partitions  s    r   c                  C   sX  dd t dD } dd t dD }tj| dd}tj|dd}ddd	}|| dd
}ttj||dd
| t|j|dd
| t||d| || |}tt|||| t|j||d
| t|||| | t}|| tt|t	| }t|j||d
| t|||| t
|}t|j||d
| t|||| d S )Nc                 S   s   g | ]}t d dqS irG   randomrandintrQ   r,   r,   r.   rU     s     z3test_map_partitions_args_kwargs.<locals>.<listcomp>rG   c                 S   s   g | ]}t d dqS r   r   rQ   r,   r,   r.   rU     s     rH   rI   r   c                 S   s*   t |trt|n|}dd t| |D S )Nc                 S   s   g | ]\}}t ||qS r,   )rz   )rR   rL   r1   r,   r,   r.   rU     s     zDtest_map_partitions_args_kwargs.<locals>.maximum.<locals>.<listcomp>)r   r   r   rd   )r&   yr,   r,   r.   maximum  s    z0test_map_partitions_args_kwargs.<locals>.maximum)r   )r   )ra   r_   r`   r   r   r   applyr   re   rE   rf   rg   )r&   r   ZdxZdyr   rm   Zdy_meanr,   r,   r.   test_map_partitions_args_kwargs  s(    


r   c                  C   s4   t tjdd ddjd} | s"tt| ts0td S )Nc                 S   s   | S r6   r,   r-   r,   r,   r.   rC     rD   z/test_map_partitions_blockwise.<locals>.<lambda>ztest-string)token)r!   r1   r   rf   r3   r   r   Zlayerr,   r,   r.   test_map_partitions_blockwise  s     r  c                  C   s@   t jtddd} dtt| dd  k r6dk s<n tdS )	zH
    Number of randomly sampled elements are in the expected range.
      r$   rI   rH   g?*   i,  N)r_   r`   ra   rE   rc   random_sampler3   rO   r,   r,   r.   test_random_sample_size  s    r  c               	   C   sV   t jtddd} tt | d W 5 Q R X tt | d W 5 Q R X dS )zN
    Specifying probabilities outside the range [0, 1] raises ValueError.
    r   r$   rI   g?N)r_   r`   ra   rh   ri   rj   r  rO   r,   r,   r.   test_random_sample_prob_range  s
    r	  c                  C   s4   t jtddd} | d}t|t|ks0tdS )zf
    Repeated computation of a defined random sampling operation
    generates identical results.
    r   r$   rI   g?Nr_   r`   ra   r  rc   r3   r   r,   r,   r.   'test_random_sample_repeated_computation  s    
r  c                  C   sR   t jtddd} t| dt| dks2t| dj| djksNtdS )z
    Repeatedly defining a random sampling operation yields different results
    upon computation if no random seed is specified.
    r   r$   rI         ?N)r_   r`   ra   rc   r  r3   r2   rO   r,   r,   r.   (test_random_sample_different_definitions  s     r  c                  C   sB   t jtddd} | dd}| dd}t|t|ks>tdS )zF
    Sampling with fixed random seed generates identical results.
    r   r$   rI   r  i  Nr
  rK   r,   r,   r.   test_random_sample_random_state  s    r  c                  C   s   t tttdddgfff} t| t ttdddgffks8ttttdddgff} t| | ks\tttttttdffff}tttttdfff}t||kstd S )Nr(   r)   r*   r   )re   r   rb   r#   r   r3   r}   r/   )ZtaskrL   r1   r,   r,   r.   test_lazify_task  s     r  c                 C   s   | S r6   r,   r-   r,   r,   r.   rC   %  rD   c                  C   s`   t ttt ttdfffftdftdfd} t ttttdffftdftdfd}t| |ks\td S )Nr   r&   )r&   rL   r1   )r   rb   r#   r}   r/   r   r   r3   r   r,   r,   r.   test_lazify(  s    $r  c                  C   s   t dftddfd} dtt dfdfi}t| dg|ks:tdtddfi}t| dgdd|ks`tt| ddg| ksvtt| ddgdd| kstt dftddftddfd	} t| dd
g| kstdtdfd} t| dg| kstd S )NrL   r1   r(   rY   rM   FZrename_fused_keysr)   r1   rM   r   r   )r<   r$   )rc   r   r   r3   r   )inpoutr,   r,   r.   test_inline_singleton_lists2  s    r  c               	   C   s   t dftddfd} t| dgdd}| dhks6t|d dd  dksNttjd	di t| dg|ksttW 5 Q R X t| dg|kstd S )
NrL   r1   r(   rY   rM   Fr  r   zoptimization.fuse.rename-keys)rc   r   r   r:   r3   rf   configr   )r  Zoutpr,   r,   r.   test_rename_fused_keys_bagF  s    r  c                   C   sH   t tdddgksttddks,tttjdddtsDtd S )Nr)   r   r(   )r   r(   Frw   )rc   r1   taker3   r   r   r,   r,   r,   r.   	test_takeT  s    r  c                	   C   sz   t tjdddddddddgks&ttjdddd	ks<ttjdddd
ksRttt tjddd W 5 Q R X d S )Nrq   r)   rI   r   r(   r*   r<   r  )r   r(   r)   r*   r<   r   )r   r(   r)   r$   )rc   r1   r  r3   rh   ri   rj   r,   r,   r,   r.   test_take_npartitionsZ  s
    &r  c               
   C   s   t jjdd tt td W 5 Q R X tt td W 5 Q R X tj	dd"} tjddd tjdd	d
 W 5 Q R X | rt
W 5 Q R X d S )Nr   r   rG   rr   Trecordr)   rI   F)warn)rf   r  r   rh   ZwarnsUserWarningr1   r  warningscatch_warningsr3   r  r,   r,   r.   test_take_npartitions_warnb  s    r!  c                   C   s"   t tdd dddgtstd S )Nc                 S   s   | S r6   r,   r-   r,   r,   r.   rC   s  rD   z"test_map_is_lazy.<locals>.<lambda>r(   r)   r*   )r   rb   r   r3   r,   r,   r,   r.   test_map_is_lazyr  s    r"  c                   C   s   t tt tstd S r6   )r   r   r1   r   r3   r,   r,   r,   r.   "test_can_use_dict_to_make_concretev  s    r#  ZHangs)reasonc                  C   sJ   t ddg} | jdkstt d}|jdks4td|dksFtd S )Nzhttp://google.comzhttp://github.comr)   z:http://raw.githubusercontent.com/dask/dask/main/README.rstr(   s   Dask
rH   )r_   Zfrom_urlr4   r3   r  r   r,   r,   r.   test_from_urlz  s
    
r%  c               	   C   sn   t dddF} dd t| D tdks0tdd tdD tdksPtW 5 Q R X ttd	d
  d S )NzA
BzC
D)za1.logza2.logc                 S   s   h | ]}|  qS r,   striprR   liner,   r,   r.   	<setcomp>  s     z!test_read_text.<locals>.<setcomp>ABCDc                 S   s   h | ]}|  qS r,   r&  r(  r,   r,   r.   r*    s     za*.logc                   S   s
   t dS )Nznon-existent-*-path)r_   	read_textr,   r,   r,   r.   rC     rD   z test_read_text.<locals>.<lambda>)r   r_   r,  r   r3   rh   ri   rj   )fnsr,   r,   r.   test_read_text  s     *r.  c               
   C   s   t  } t| d}|dtj  d  W 5 Q R X tj| dd}t| }t|j	dks`t
ttt|j ttt|j kst
tj| gdd}t|t|kst
W 5 Q R X d S )NwbzHello, world!rG   	blocksizer$   )r   openwriteoslinesepencoder_   r,  rE   rf   r3   rc   rb   r   r'  )fnr   r1   rM   r   r,   r,   r.   test_read_text_large  s    "
,r8  c               
   C   s   t  } t| d }|dtj dd  W 5 Q R X tj| ddd}tj| dd}t|j	dksht
|j dd	 }|j d
d	 }t|t|kst
tj| gddd}t|t|kst
W 5 Q R X d S )Nr/  u	   你好！gb18030rG   )r1  encodingr:  r$   c                 S   s
   |  dS Nutf-8r6  r-   r,   r,   r.   rC     rD   z)test_read_text_encoding.<locals>.<lambda>c                 S   s
   |  dS r<  r>  r-   r,   r,   r.   rC     rD   )r   r2  r3  r4  r5  r6  r_   r,  rE   rf   r3   r   r'  rb   rc   )r7  r   r1   rM   Zb_encZc_encr   r,   r,   r.   test_read_text_encoding  s    $r?  c               
   C   s   t d} d}t| d}|| |  tt tj| ddd W 5 Q R X tj| d d}|j	dkslt
d	| | kst
W 5 Q R X d S )
Ngzsx  Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
Hello, world!
r/  r   
)r1  Zlinedelimiterr0  r(   r   )r   r   r3  closerh   ri   rj   r_   r,  r4   r3   r	   r=   decode)r7  r{   r   rM   r,   r,   r.   test_read_text_large_gzip  s    


rD  z(https://github.com/dask/dask/issues/6914c                  C   s`   t d d} tjdtddd}|d| ks4ttjdd	d
gtddd}|jdks\td S )NZs3fs)z(total_bill,tip,sex,smoker,day,time,size
z"16.99,1.01,Female,No,Sun,Dinner,2
z 10.34,1.66,Male,No,Sun,Dinner,3
z21.01,3.5,Male,No,Sun,Dinner,3
z 23.68,3.31,Male,No,Sun,Dinner,2
zs3://tip-data/t*.gzT)Zanon)Zstorage_optionsr$   zs3://tip-data/tips.gzzs3://tip-data/tips.jsonzs3://tip-data/tips.csvr*   )rh   r   r_   r,  r   r  r3   r4   )Z	five_tipsr   rM   r,   r,   r.   test_from_s3  s    
	rE  c                  C   sH   t jdddddgdd} t| jdks*tt| dddddhksDtd S )Nr(   r)   r*   r<   r$   rI   )r_   r`   rE   rf   r3   r   rV   r,   r,   r.   test_from_sequence  s    rF  c                  C   s.   t td} t| }t|t| ks*td S )Ni  )rc   ra   r_   r`   r   r3   )r   r1   r,   r,   r.   test_from_long_sequence  s    
rG  c                  C   sH   t d tg } | jdks"t| jddid }|jsDtdd S )Ndask.dataframer(   rL   r   metazDataFrame is not empty)	rh   r   r_   r`   r4   r3   to_dataframer=   empty)r1   dfr,   r,   r.   test_from_empty_sequence  s
    

rN  c                  C   s   t t } | jt jd kstt| dd tD ks8ttddddg}tddd	g}||}t|d
d dD ks~t|j| jkst|j||jkstd S )Nr)   c                 S   s   h | ]}t D ]}||fqqS r,   )r   rZ   r,   r,   r.   r*    s       ztest_product.<locals>.<setcomp>r(   r*   r<   rH   r   r   c                 S   s   h | ]}d D ]}||fqqS ))rH   r   r   r,   rZ   r,   r,   r.   r*    s       )r(   r)   r*   r<   )	r1   productr4   r3   r   r   r_   r`   r2   )rk   r&   r   zr,   r,   r.   test_product  s    

rQ  c               	   C   s   t  } tttdd|  t| dddhks6tt| dddhksPtt| dddhksjttt	td| dddgfddgfgkstW 5 Q R X d S )	Nrq   r*   r   r(   r<   r)   r$   r   )
partdZPickler   r   ra   r   getr3   r8   r   )pr,   r,   r.   test_partition_collect  s    
rU  c                  C   s   t t} t| }|dddgdddgdddgdddgdddgdksHt| jt jksXt| jt tjksnt| jt dd jkstd S )	Nr   r(   r)   r*   r<   r   c                 S   s   | d S r   r,   r-   r,   r,   r.   rC     rD   ztest_groupby.<locals>.<lambda>)r1   r   r   r   r3   r4   r2   )rM   r?   r,   r,   r.   test_groupby  s    
rV  c                  C   sd   t dddgdddgdddgg} t| d}tt|dddgdddggdddggdks`td S )Nr(   r)   r*   r<   rt   r   rp   )r_   r`   r   r   r   r8   r3   r1   r?   r,   r,   r.   test_groupby_with_indexer  s    "rX  c                  C   sb   t jdd dd} t| }|dddgdddgdddgdddgdddgd	ksPt| jdks^td S )
Nc                 S   s   | S r6   r,   r-   r,   r,   r.   rC     rD   z7test_groupby_with_npartitions_changed.<locals>.<lambda>r(   rI   r   r)   r*   r<   r   )r1   r   r   r3   r4   )r?   Zresult2r,   r,   r.   %test_groupby_with_npartitions_changed  s    rY  c               	   C   s@   ddl m}  tjj| d tjdd dd  W 5 Q R X d S )Nr   rS  r   c                 S   s   | S r6   r,   r-   r,   r,   r.   rC   )  rD   z2test_groupby_with_scheduler_func.<locals>.<lambda>r(   rI   )Zdask.threadedrS  rf   r  r   r1   r   r=   rZ  r,   r,   r.    test_groupby_with_scheduler_func%  s    r[  c                  C   sh   t dddg} t dddg}t | |g}t|ddddddgksJt|jt | |gjksdtd S )Nr(   r)   r*   r<   r$   rq   )r_   r`   concatrc   r3   r2   rK   r,   r,   r.   test_concat,  s
    r]  c                  C   sJ   t dgddgg} t|  dddgks.t|  j|  jksFtd S )Nr(   r)   r*   )r_   r`   rc   flattenr3   r2   rV   r,   r,   r.   test_flatten4  s    r_  c                  C   sL   t ddg} t ddg}t | t|g}t|ddddgksHtd S )Nr(   r)   r<   r$   r*   )r_   r`   r\  rb   r#   rc   r3   )rL   r1   r?   r,   r,   r.   test_concat_after_map:  s    r`  c                  C   s@   t dd } t| j }t| t|ks,t| j|jks<td S )Nc                 S   s   | d S r   r,   r-   r,   r,   r.   rC   B  rD   ztest_args.<locals>.<lambda>)r1   rb   r   _argsrc   r3   r4   )rM   r   r,   r,   r.   	test_argsA  s    
rb  c            
   	   C   sJ  t d} t d}dd }dgg ddgd}t|d	d
}|j| ddgd}| }| jj||jdddddd |jddgd}| jj||dd ||| |jddgd}| jj||dd ||| |	dd }| }| jj||dd ||| |j
| j
kst|ddgfD ].}|j|d}| jj||dd ||| q$t t |jddg|d W 5 Q R X |dd }t t |  W 5 Q R X |d}|dg }|j|d}| jj||dd ||| |dtdi}tjtddd}ttfD ]4}	||	j|d}| jj||dd ||| qd S )NrH  Zpandasc                    s(   t  fddtj|   D s$td S )Nc                 3   s   | ]}|j  j k V  qd S r6   )Zdtypesr   )rR   rT  rm   r,   r.   r   N  s    z9test_to_dataframe.<locals>.check_parts.<locals>.<genexpr>)r   rf   r=   
to_delayedr3   )rM  rm   r,   rc  r.   check_partsM  s    z&test_to_dataframe.<locals>.check_partsrp   )rH   r   )rG   rJ   ))testr   )rf  r(   )rf  r)   rf  r*   rL   r1   )columnsr   r(   r   F)Zcheck_index)rL   i8)r1   rh  rI  c                 S   s   t tddg| S )NrL   r1   )r   rd   r-   r,   r,   r.   rC   a  rD   z#test_to_dataframe.<locals>.<lambda>)rg  rJ  c                 S   s   | d dkS )NrL   rJ   r,   r-   r,   r,   r.   rC   r  rD   rG   r$   rI   )rh   r   r   Z	DataFramer=   rK  utilsr   renamerb   _namer3   ri   rj   r}   r   ra   r_   r`   itertupler   )
ddpdre  r9   r1   rm   rM  rJ  rk   r   r,   r,   r.   test_to_dataframeI  sN    

 





rp  r@  bz2r   z
ext,myopenc              	   C   s   t jdddgdd}t }|jtj|d|  dd}tj|d	d
i tj	tj|d|  sft
|tj|d|  d}| }t|dr| }d|kst
|  W 5 Q R X d S )Nabc123xyzr)   rI   *.Frw   r   r   1.rbrC  )r_   r`   r   to_textfilesr4  pathr	   rf   r=   existsr3   readhasattrrC  rB  )extmyopenr1   dirrM   r   textr,   r,   r.   test_to_textfiles  s    
r  c                  C   s   ddddddddd	d
ddddddg} t j| dd}t F}|| t jtj|dddt	t	j
 }| |ksztW 5 Q R X d S )NrL   r1   rM   r   r   r   r   hrS   r[   r   lr   r   orT     rI   *asciir;  )r_   r`   r   rx  r,  r4  ry  r	   rb   r   r'  r=   r3   )r   r1   dnr  r,   r,   r.   /test_to_textfiles_name_function_preserves_order  s6    
r  c                  C   sV   ddddddddd	d
ddddddg} t j| dd}t }|j|td W 5 Q R X d S )NrL   r1   rM   r   r   r   r   r  rS   r[   r   r  r   r   r  rT  r  rI   )Zname_function)r_   r`   r   rx  r   )r   rL   r  r,   r,   r.   $test_to_textfiles_name_function_warn  s(    r  c               
   C   s   t jdddgdd} tD ]\}}t }| jtj|d| ddd	}tj	|  tj
tj|d
| snt|tj|d
| d}| }t|dr|d}d|kst|  W 5 Q R X qd S )Nu   汽车u   苹果u   天气r)   rI   ru  r9  F)r:  r=   rv  rw  rC  )r_   r`   ext_openr   rx  r4  ry  r	   rf   r=   rz  r3   r{  r|  rC  rB  )r1   r}  r~  r  rM   r   r  r,   r,   r.   test_to_textfiles_encoding  s       


r  c               
   C   s   t jdddgdd} t F}t 4}| ||g tj|sBttj|sRtW 5 Q R X W 5 Q R X t :}| | tj|sttjtj	|dstW 5 Q R X t
t | d W 5 Q R X d S )Nrr  rs  rt  r)   rI   z0.partr$   )r_   r`   r   rx  r4  ry  rz  r3   r   r	   rh   ri   	TypeError)BrL   r1   dirnamer,   r,   r.   test_to_textfiles_inputs  s    $
$r  c                  C   sz   t jdddgdd} t V}dD ]J}| j|g|d t|}| }W 5 Q R X |dd	|r^d
ndgks tq W 5 Q R X d S )NrL   r1   rM   r(   rI   )FT)last_endlineza
zb
zc
)r_   r`   r   rx  r2  	readlinesr3   )r1   r7  r  r   r?   r,   r,   r.   test_to_textfiles_endlines  s    
r  c                      s   t jdddgdd dt jks&tdt jks8tt j dd	d
gksTtt jdddgddgddggks~tt jdddgkstt	
t fdd  jdj jdjkst jdj jdjkstd S )NAlice Smith	Bob JonesCharlie Smithr)   rI   splitmatchalice smith	bob jonescharlie smith AliceSmithBobJonesCharliez*Smithc                      s    j jS r6   )r   Z	sfohsofhfr,   rV   r,   r.   rC     rD   z'test_string_namespace.<locals>.<lambda>z*John)r_   r`   r  r   r3   rc   lowerr  r  rh   ri   AttributeErrorr2   r,   r,   rV   r.   test_string_namespace  s    
 r  c                  C   s4   t jdddgdd} t| j dddgks0td S )	Nr  r  r  r)   rI   r  r  r  )r_   r`   rc   r   r  r3   rV   r,   r,   r.   "test_string_namespace_with_unicode  s    r  c                  C   s@   t jdddgdd} t| j ddgdd	gd
dggks<td S )Nr  r  r  r)   rI   r  r  r  r  r  )r_   r`   rc   r   r  r3   rV   r,   r,   r.   test_str_empty_split  s    r  c                  C   sT   t jdddgdddggdd} dd	 }| |}t|dddgddd
ggksPtd S )Nr(   r)   r*   r<   r$   rq   rI   c                 s   s   | D ]}|d V  qd S r   r,   )r   r&   r,   r,   r.   r   "  s    z*test_map_with_iterator_function.<locals>.frr   )r_   r`   rb   rc   r3   )r1   r   rM   r,   r,   r.   test_map_with_iterator_function  s    
r  c                  C   s4   t dddg} | dd  }t|tr0td S )Nr(   r)   r*   c                 S   s   | d S r   r,   r-   r,   r,   r.   rC   -  rD   z8test_ensure_compute_output_is_concrete.<locals>.<lambda>)r_   r`   rb   r=   r   r   r3   rW  r,   r,   r.   &test_ensure_compute_output_is_concrete+  s    r  c                   @   s   e Zd ZdddZdd ZdS )
BagOfDictsNc                    s   |   fddS )Nc                    s   |   S r6   rZ  r   defaultr   r,   r.   rC   3  rD   z BagOfDicts.get.<locals>.<lambda>rb   )selfr   r  r,   r  r.   rS  2  s    zBagOfDicts.getc                    s    fdd}|  |S )Nc                    s   |  < | S r6   r,   r  r   valuer,   r.   setter6  s    zBagOfDicts.set.<locals>.setterr  )r  r   r  r  r,   r  r.   r   5  s    zBagOfDicts.set)N)__name__
__module____qualname__rS  r   r,   r,   r,   r.   r  1  s   
r  c                  C   s   t tdddiigj } | dd d dks:t| ddd d dddksbtt| ddt s|td S )NrL   r1   rM   r   r   zEXTENSIBILITY!!!r1   r   )	r  r_   r`   ra  rS  r=   r3   r   r   )Zdictbagr,   r,   r.   test_bag_class_extend=  s     r  c               
   C   sb   d  } t J}t|d}||  W 5 Q R X t|}| d | dksTtW 5 Q R X d S )Nu   €r/  r   r=  )	r6  r   r2  r3  r_   r,  r=   rC  r3   )Zbin_datar7  r   rL   r,   r,   r.   
test_gh715G  s    
r  c                  C   s*   t dddgdd } | jdd d S )Nr(   r)   r*   c                 S   s   | d S r   r,   rO   r,   r,   r.   rC   Q  rD   z1test_bag_compute_forward_kwargs.<locals>.<lambda>rH   )Zbogus_keyword)r_   r`   rb   r=   r-   r,   r,   r.   test_bag_compute_forward_kwargsP  s    r  c                  C   s   t jddddddgdd} | t \}} }tdd	 || |fD sJt|  ddgks^tt jddddddgdd} |   }t	|t
st| d
kstd S )Nr(   r)   r*   r<   r$   rq   rI   c                 s   s   | ]}t |tV  qd S r6   )r   r   rR   r&   r,   r,   r.   r   X  s     z"test_to_delayed.<locals>.<genexpr>   )r_   r`   rb   r#   rd  r   r3   r=   re   r   r   )r1   rL   rM   tr,   r,   r.   test_to_delayedU  s    r  c                 C   s  t jddddddgdd}|ttt}| \}tt|j}|ddksZt	|
 |
 ksnt	|jd	d
\}t|jt|jkst	|
 |
 kst	| | kst	| }| }tt|j}|
 |
 kst	|ddkst	|jd	d
}t|jt|jks(t	|
 |
 ks>t	| | ksTt	|jt| d	d\}tt|j}|ddkst	d S )Nr(   r)   r*   r<   r$   rq   rI   r   FZoptimize_graphr   rw   )r_   r`   rb   r#   rd  r   r   rf   r   r3   __dask_layers__r=   re   rx  )r   r1   rk   r   r  d2r&   r,   r,   r.   test_to_delayed_optimize_grapha  s,    
r  c               
   C   s   ddl m}  | dddg| dddg| d	d
dg  }}}t|||g}|jt|||gjksbtt|tsptt|ddddddd	d
dg	kst| t|}t	j
|}| |   krdksn td S )Nr   rg   r(   r)   r*   r<   r$   rq   rr   rs   rt   )dask.delayedrg   r   r2   r3   r   r   rc   re   r_   Itemr=   )rg   rL   r1   rM   ZbbZ
asum_valueZ	asum_itemr,   r,   r.   test_from_delayed}  s    ."r  c                     sv   ddl m}  dd }| |dd t fddtd	D }tj| |d
 |d
  dddksrt	d S )Nr   r  c                 S   s   dd t | D S )Nc                 s   s   | ]}d ddgiV  qdS )
operationsr(   r)   Nr,   rR   _r,   r,   r.   r     s     zCtest_from_delayed_iterator.<locals>.lazy_records.<locals>.<genexpr>ra   )r   r,   r,   r.   lazy_records  s    z0test_from_delayed_iterator.<locals>.lazy_recordsF)Zpurec                    s   g | ]} d qS )r$   r,   r  Zdelayed_recordsr,   r.   rU     s     z.test_from_delayed_iterator.<locals>.<listcomp>r$   r  r   r   )   r  r   )
r  rg   r_   r   ra   r=   r   r   r^  r3   )rg   r  r   r,   r  r.   test_from_delayed_iterator  s    r  c                  C   sT   dD ]J} t jd| d}t|j| ks(t|j| ks6tt|ttdkstqd S )N)r(   rr   rH      rG   rI   )r_   ra   rE   rf   r3   r4   rc   r   r,   r,   r.   
test_range  s
    r  rr   r  r  c                 C   st   t jtd|d| d}t jtd|d| d}t ||}|j| ksFtt|tttd|dtd|dksptd S )Nr   r)   rI   r(   )r_   r`   ra   rd   r4   r3   rc   )r4   hiZevensZoddspairsr,   r,   r.   test_zip  s
    r  nin      noutr   c                 C   sX   t jtd| d}|j|d}|j|ks,tt|| t|j|	 }t
|sTtd S )NrG   rI   )r_   r`   ra   repartitionr4   r3   r   rf   rS  r7   r   )r  r  r1   rM   r   r,   r,   r.   test_repartition_npartitions  s    
r  z	nin, noutr(   r(   )r)   r(   )r$   r(   rp   )r)   r)   )r$   r)   )r(   r$   )r)   r$   )r$   r$   c                 C   sR   t jtdd| d}t|t }|j|| d}|j|ksDt	t
|| d S )Nr(   rG   rI   r   r_   r`   ra   re   r   r   r=   r  r4   r3   r   )r  r  r1   	total_memrM   r,   r,   r.   test_repartition_partition_size  s
    r  c                  C   sz   t jtdddd} t| t }| j|d d}|jdksDt	t
| | |j|d d}|jdkslt	t
|| d S )Nr(   rG   rI   r)   r   r$   r  )r1   r  rM   r   r,   r,   r.   (test_multiple_repartition_partition_size  s    
r  c                     sl   t d tj fddtdD dd} t| t }|d }| j	|d}|j
dks^tt| | d S )Nr   c                    s   g | ]}  td qS )rG   )arrayra   r  r   r,   r.   rU     s     zBtest_repartition_partition_size_complex_dtypes.<locals>.<listcomp>r<   r(   rI   r   )rh   r   r_   r`   ra   re   r   r   r=   r  r4   r3   r   )r1   r  Znew_partition_sizerM   r,   r  r.   .test_repartition_partition_size_complex_dtypes  s    
 r  c                  C   sp   t jtddd} | d}| j|jks,t| d}| j|jksFt|j|jksVt| d}| |ksltd S )NrG   r$   rI   r)   r   )r_   r`   ra   r  r2   r3   r  r,   r,   r.   test_repartition_names  s    


r  c               	   C   s6   t t" ttd} | jddd W 5 Q R X d S )NrH   r$   Z5MiB)r4   r   )rh   ri   rj   r_   r`   ra   r  r   r,   r,   r.   test_repartition_input_errors  s    r  c               	   C   s  dddgddgg ddgg} dd	 t | D }t|d
t| }|t}|j|tjks^t|j|tdjksvt| dddddddgkst|td ddddddddgkst|t	t
 dddddddgksttjdddgdd}|t dddgkstd S )Nr(   r)   r*   r<   r$   rq   rr   c                 S   s   i | ]\}}d |f|qS )rf  r,   )rR   rS   rT  r,   r,   r.   r~     s      z#test_accumulate.<locals>.<dictcomp>rf  r  rH   r   r  r  r   rt      r      r  r        rI   )	enumerater_   r   rE   
accumulater    r2   r3   r=   rb   r#   r`   )partsr9   r1   r   r,   r,   r.   test_accumulate  s    
(*r  c                  C   sx  t jtddd} | jdd ddd}t|j| }|D ]6}|D ],} || k	rDttd|ttd| @ rDt	qDq<t jtd	d
d} | jdd dd}t
|jdk st	t|j| }|D ]6}|D ],} || k	rttd|ttd| @ rt	qqt jtddd} | jdd ddd}t|j| }|D ]>}|D ]2} || k	r<ttd|ttd| @ r<t	q<q4d S )N   r<   rI   c                 S   s   | d S r   r,   r-   r,   r,   r.   rC     rD   z$test_groupby_tasks.<locals>.<lambda>tasks
max_branchshuffler   r  rG   c                 S   s   | d S )N{   r,   r-   r,   r,   r.   rC     rD   r  i'  iY  c                 S   s   | d S )Ni  r,   r-   r,   r,   r.   rC     rD      )r_   r`   ra   r   rf   rS  r7   r   r   r3   rE   )r1   r  Z
partitionsrL   r,   r,   r.   test_groupby_tasks	  s,    $$
r  c                  C   s   t jtddd} dd }dd }t| j|dddjt| j|dddjksRtt| j|dddjt| j|d	ddjkstt| j|dddjt| j|dddjkstd S )
Nr  r<   rI   c                 S   s   | d S r   r,   r-   r,   r,   r.   rC   )  rD   z*test_groupby_tasks_names.<locals>.<lambda>c                 S   s   | d S )Nr   r,   r-   r,   r,   r.   rC   *  rD   r  r  r)   )r_   r`   ra   r   r   rf   r3   )r1   r   Zfunc2r,   r,   r.   test_groupby_tasks_names'  s    

r  zsize,npartitions,groups)r  r   rG   )i90     i  )rG   r(   r   c                    sN    fdd}t j| |dj|dd}|jdd}t|t|t| ksJtd S )Nc                    s   |   S r6   r,   r-   groupsr,   r.   rC   :  rD   z&test_groupby_tasks_2.<locals>.<lambda>rI   r  r  r   r   r_   ra   r   r=   r   r3   )sizer4   r  r   r1   r?   r,   r  r.   test_groupby_tasks_26  s    r  c                  C   sL   dd } t jdddj| ddd}|jd	d
}t|t| tdksHtd S )Nc                 S   s   | d S r   r,   r-   r,   r,   r.   rC   A  rD   z&test_groupby_tasks_3.<locals>.<lambda>r   r$   rI   r  r)   )r  r  r   r   r  )r   r1   r?   r,   r,   r.   test_groupby_tasks_3@  s    r  c               	   C   s^   t  N} tjddddd t}|tj	| d t
t| dksPtW 5 Q R X d S )Nr$   rI   c                 S   s   | dkS r   r,   r-   r,   r,   r.   rC   J  rD   z4test_to_textfiles_empty_partitions.<locals>.<lambda>z*.txt)r   r_   ra   r}   rb   r   rx  r4  ry  r	   rE   listdirr3   r   r,   r,   r.   "test_to_textfiles_empty_partitionsH  s    r  c                  C   sF   t jtddd} t| dd  d t| dd  d d S )	NrH   rG   rI   c                 S   s   | d dkS r+   r,   r-   r,   r,   r.   rC   Q  rD   z&test_reduction_empty.<locals>.<lambda>rs   c                 S   s   | d dkS r+   r,   r-   r,   r,   r.   rC   R  rD   r   )r_   r`   ra   r   r}   rz   r   rV   r,   r,   r.   test_reduction_emptyO  s    r  c              	   C   s   t jddddg| dd }t|jddd t j|jdd|jdddd}|dksZtt	t
6 t jddddg| d}|d jddjdd W 5 Q R X d S )	Nr   r(   rI   r)   r   r   r   r  )r_   r`   r}   r   r   r=   rz   r3   rh   ri   rj   )r4   r1   valsr,   r,   r.   test_reduction_empty_aggregateU  s     r  c                   @   s   e Zd Zdd Zdd ZdS )
StrictRealc                 C   s   t |tst| j|jkS r6   r   r   r3   realr  r   r,   r,   r.   __eq__a  s    zStrictReal.__eq__c                 C   s   t |tst| j|jkS r6   r  r  r,   r,   r.   __ne__e  s    zStrictReal.__ne__N)r  r  r  r  r  r,   r,   r,   r.   r   `  s   r   c                  C   s6   t jdd tdD dd} t| tttd d S )Nc                 S   s   g | ]}t |qS r,   )r   r  r,   r,   r.   rU   k  s     z>test_reduction_with_non_comparable_objects.<locals>.<listcomp>rH   r)   r   rt   )r_   r`   ra   r   r   rz   r   rV   r,   r,   r.   *test_reduction_with_non_comparable_objectsj  s    r  c                     sX   t d tj fddtdD dd}  fdd}| ||jd	d
jdksTtd S )Nzscipy.sparsec                    s   g | ]}  d gqS r   )Z
csr_matrixr  spr,   r.   rU   q  s     z7test_reduction_with_sparse_matrices.<locals>.<listcomp>r<   r)   r   c                    s     | |gS r6   )Zvstackr   r  r,   r.   	sp_reduces  s    z6test_reduction_with_sparse_matrices.<locals>.sp_reducer   r   )r<   r(   )	rh   r   r_   r`   ra   r   r=   shaper3   )r1   r
  r,   r  r.   #test_reduction_with_sparse_matriceso  s    
 r  c                   C   s   t tg g kstd S r6   )rc   r_   r`   r3   r,   r,   r,   r.   
test_emptyy  s    r  c                  C   sj   ddl m} m} ttd}|| |}| | ks>t| }|| |}| | ksftd S )Nr   )dumpsloadsrG   )	pickler  r  r_   r`   ra   r=   r3   re   )r  r  r1   rk   ss2r,   r,   r.   test_bag_picklable}  s    r  c                  C   s@   t ddigd} | jdd}t|dddigiks<td S )NrL   r(   r   r   )r_   r`   r   r=   r   r3   rW  r,   r,   r.   test_msgpack_unicode  s    r  c                  C   s$   dd } t | g}t|| g d S )Nc                   S   s   d S r6   r,   r,   r,   r,   r.   rC     rD   z/test_bag_with_single_callable.<locals>.<lambda>)r_   r`   r   )r   r1   r,   r,   r.   test_bag_with_single_callable  s    r  c                     s   t jddd} | t}|t}||j|  |j   @ rJt|j|j| | d t	 fdd| D std S )NrH   r)   rI   )Z	fuse_keysc                 3   s   | ]}| kV  qd S r6   r,   r   r9   r,   r.   r     s     z*test_optimize_fuse_keys.<locals>.<genexpr>)
r_   ra   rb   r#   Z__dask_optimize__rf   r7   r:   r3   r   )r&   r   rP  r,   r  r.   test_optimize_fuse_keys  s    

r  c                     sR   d g  fdd}  fdd}t d|  idd}||t}t|ttd d S )	Nc                  3   s    t dD ]} |  d< | V  qd S )NrH   r   r  rS   currentr,   r.   part  s    z&test_reductions_are_lazy.<locals>.partc                    s    d dkst t| S r   )r3   re   )r  r  r,   r.   r     s    z&test_reductions_are_lazy.<locals>.func)r   r   r   r(   rH   )r   r   re   r   ra   )r  r   r1   resr,   r  r.   test_reductions_are_lazy  s    r  c                  C   s@   t jddd} | dd }ttt|ttt|ks<td S )NrH   r<   rI   c                 S   s   | d S )Nr*   r,   r-   r,   r,   r.   rC     rD   z'test_repeated_groupby.<locals>.<lambda>)r_   ra   r   r   rE   r   r3   rY   r,   r,   r.   test_repeated_groupby  s    r  c              
   C   s|   t jddd}td\}tjjt| |d< |dd }|  t	dd t
t| D sdtW 5 Q R X W 5 Q R X d S )	NrH   r<   rI   )Ztemporary_directorypoolc                 S   s   | d S )Nr)   r,   r-   r,   r,   r.   rC     rD   z*test_temporary_directory.<locals>.<lambda>c                 s   s   | ]}| d V  qdS )z.partdN)endswith)rR   r7  r,   r,   r.   r     s     z+test_temporary_directory.<locals>.<genexpr>)r_   ra   r   rf   r  r   r   r   r=   r   r4  r  r3   )r   r1   r  rk   r,   r,   r.   test_temporary_directory  s    
r!  c                  C   s^   t g } t| t d t| t d t| t d t| t d d S )NTF)	r_   r`   r   rb   r#   r   r   re   r   rV   r,   r,   r.   test_empty_bag  s
    
r"  c                  C   sZ   t jdddgdd} | d}|d ds0t|d	 d
sBttd td
 d S )Nrr  rs  rt  r)   rI   zfoo*r   Zfoo0r(   Zfoo1)r_   r`   rx  r   r3   r4  r   )r1   pathsr,   r,   r.   test_bag_paths  s    

r$  c                  C   sR   dd } t dddg}t|| dddd	g t|| tdddd	g d S )
Nc                    s    fdd| D S )Nc                    s   g | ]}|  qS r,   r,   r  r  r,   r.   rU     s     z?test_map_partitions_arg.<locals>.append_str.<locals>.<listcomp>r,   )r   r  r,   r%  r.   
append_str  s    z+test_map_partitions_arg.<locals>.append_strrL   r1   rM   r   ZafooZbfooZcfoo)r_   r`   r   r   rf   rg   )r&  Zmybagr,   r,   r.   test_map_partitions_arg  s     r'  c                  C   sd   t dddg} t| t }dttjj|ks8t	t
| t t
| t ks`t	d S )Nr(   r)   r*   r#   )r_   r`   r   rb   r#   Z__dask_graph__rf   ri  Z	key_splitr3   r   r   r  r,   r,   r.   test_map_keynames  s    r(  c                     s   G dd d t   fdd} dd }tjtddd	| || | }zt	  |j
d
d W 5 t  X d S )Nc                   @   s   e Zd Zdd ZdS )zCtest_map_releases_element_references_as_soon_as_possible.<locals>.Cc                 S   s
   || _ d S r6   r  )r  rS   r,   r,   r.   __init__  s    zLtest_map_releases_element_references_as_soon_as_possible.<locals>.C.__init__N)r  r  r  r)  r,   r,   r,   r.   C  s   r*  c                    s&   t dkst | }| |S r   )rE   r3   r    )rS   r  r*  Z	in_memoryr,   r.   f_create  s    
zJtest_map_releases_element_references_as_soon_as_possible.<locals>.f_createc                 S   s
   | j d S )NrG   r  )r  r,   r,   r.   f_drop  s    zHtest_map_releases_element_references_as_soon_as_possible.<locals>.f_dropr)   r(   rI   r   r   )weakrefWeakSetr_   r`   ra   rb   re   gcenabledisabler=   )r,  r-  r1   r,   r+  r.   8test_map_releases_element_references_as_soon_as_possible  s$    
r3  c                  C   sJ   t d} | jddd d }t|}| }t|dddddg d S )N
dask.arrayrH   r$   )chunksr   g      ?)rh   r   onesrd  r_   r   r=   r   )daobjr   r1   r,   r,   r.   test_bagged_array_delayed"  s
    

r9  c                  C   s   t jddgdd} |  | jfks&t| jj | jhks>t| jj| jt	 iksVt| 
 }| |jfksrt|jj | j|jhkst|jj| jt	 |j| jhikstd S )Nr(   r)   rI   )r_   r`   r  r2   r3   rf   layersr:   dependenciesr   r   r   )rL   rS   r,   r,   r.   test_dask_layers+  s    r<  r   FTc              	   C   s  t d}tj|d d }|jd }|jdd  dksDt|j	j
 d| hks^t|j	jd| t iksxt| d| fkst|dd }|j| dd }tj|}|j|jkst|j	|j	kst| | kst|j| d}| | kst| sP|j	|j	ks&tt jtdd t|j	|j W 5 Q R X t jtdd tj|j	|jfd	d
 W 5 Q R X d S )Nr4  r(   r   r  zdelayed-r  znot in)r  r   r  )rh   r   r_   r  r   r6  rd  r   r3   rf   r:  r:   r;  r   r  ri   rj   r2   )r   r7  rS   r2   Zarrrg   backr,   r,   r.   test_dask_layers_to_delayed6  s*    

r>  c               	   C   s   t d ddlm}  tjddddddgdd	}tjd
d0 |dd }|dd }|dd }W 5 Q R X t	|j|j
d kst| }t	|jt	|jk st|jdd}t	|jt	|j|j
 kstt|jdjdd
ikst| || d S )NrH  r   r   Ztest1r(   )r2   v1Ztest2r)   rI   T)r   c                 S   s   t f | d| d d iS )Nv2r?  r(   r   rO   r,   r,   r.   rC   `  rD   z2test_to_dataframe_optimize_graph.<locals>.<lambda>c                 S   s   t f | d| d d iS )Nv3r@  r(   rA  rO   r,   r,   r.   rC   a  rD   c                 S   s   t f | d| d d iS )NZv4rB  r(   rA  rO   r,   r,   r.   rC   b  rD   r<   Fr  r   )rh   r   Zdask.dataframe.utilsr   r_   r`   rf   Zannotaterb   rE   r4   r3   rK  r"   annotations)Zassert_eq_dfr&   r   r   r  r,   r,   r.    test_to_dataframe_optimize_graphV  s"    
 rD  )r  )r0  r   r4  r   r  r.  rq  r   collections.abcr   concurrent.futuresr   gzipr   	itertoolsr   rR  rh   Ztlzr   r   r	   r
   r   r   r   rf   Zdask.bagr   r_   Zdask.bag.corer   r   r   r   r   r   r   r   r   r   r   Zdask.bag.utilsr   Zdask.blockwiser   r  r   Z
dask.utilsr   r   r   Zdask.utils_testr    r!   r"   r#   ra   r9   rc   r   r1   r/   r0   r5   r;   r@   rF   rn   ro   r|   r   r   r   markZparametrizer   reprr   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   rg   r   r   r   r   r   r  r  r	  r  r  r  r  r   r  r  r  r  r  r!  r"  r#  Zslownetworkskipr%  r.  r8  r?  rD  ZxfailrE  rF  rG  rN  rQ  rU  rV  rX  rY  r[  r]  r_  r`  rb  rp  r2  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  r   r   r  r  r  r  r  r  r  r  r  r!  r"  r$  r'  r(  r3  r9  r<  r>  rD  r,   r,   r,   r.   <module>   s~  $44
		
	

	

	 







	
=
 
	
		

 




4	