U
    +if                     @   sJ  d dl mZ d dlZd dlZd dlmZ d dl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mZmZ d dlmZ ejj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#d$ Z,d%d& Z-d'd( Z.d)d* Z/d+d, Z0ej1d-d.d/d0g dggd1d2 Z2dS )3    )catch_warningsN)	Timestamp)	DataFrameHDFStoreIndex
MultiIndexSeries_testingbdate_rangeconcat
date_rangeisnaread_hdf)_maybe_removeensure_clean_pathensure_clean_store)Termc                 C   s  t ddddgdddggdddd	d	d
d
dddg
dd	d
dd	d	d
dd	d
g
gddgd}ttjdd|dddgd}t| N}|jd|dd |dg }t|j	ddgd| t|j	ddd| W 5 Q R X t
tjd|dd}t| ,}|jd|dd t|j	ddd| W 5 Q R X d S )NfoobarZbazZquxonetwothreer            Zfoo_nameZbar_name)Zlevelscodesnames
   ABCindexcolumnsdftableformatr#   zcolumns=['A']where)r"   names)r   r   nprandomrandnr   puttmassert_frame_equalselectr   Zassert_series_equal)
setup_pathr"   r$   storeexpectedr,    r7   W/home/mars/bis/venv/lib/python3.8/site-packages/pandas/tests/io/pytables/test_select.pytest_select_columns_in_where!   s    .

 
r9   c              	   C   s^  t tjddddddgd}tdddd|_t| }|d	| |d	}|}t	j
||d
d |jd	|jd}|}t	j
||d
d |jd	dgd}|jd d dgf }t	
|| W 5 Q R X tt tjddddddgdt tjjddddddddgdgdd}tdddd|_t| }|d	| |d	}|}t	j
||d
d |jd	|jd}|}t	j
||d
d |jd d dgf }|jd	dgd}t	j
||d
d |jd d ddgf }|jd	ddgd}t	j
||d
d W 5 Q R X t| b}|d	| |d	| |jd d ddgf }t||g}|jd	ddgd}t	j
||d
d W 5 Q R X d S )Nr      r   r   r(   z20130101 9:30T)periodsfreqr$   T)Z	by_blocksr      )sizer   r    r   Zaxis)r   r-   r.   r/   r   r"   r   appendr3   r1   r2   r#   locr   randintZreshape)r4   r$   r5   resultr6   r7   r7   r8   test_select_with_dups>   sX    

 	


rE   c              
   C   s  t | }tdd t }t|d |d| |jdddgd}|jddgd}t|| |ddg}|jddgd}t|| t|d |jd|dgd |jdd	gddgd}||j	d
k jddgd}t|| t|d |jd|dd |jdd	gddgd}||j	d
k jddgd}t|| t|d |jd|dgd |jdd	gddgd}||j	d
k jddgd}t|| W 5 Q R X W 5 Q R X d S )NT)recordr$   r   r   r(   zcolumns=['A', 'B']data_columnszA > 0r   r    D)
r   r   r1   makeTimeDataFramer   rA   r3   reindexr2   r   )r4   r5   r$   rD   r6   r7   r7   r8   test_select}   s4    



rL   c              	   C   s  t | @}ttdddtjdd}t|d |jd|ddgd |dd	}||j	t
d
k }t|| ttjddddgd}d|d< d|jdddf< |d dk|d< t|d |jd|dd ||jdk jddgd}dD ]*}|jdd| ddgd}t|| q||jdk jddgd}dD ],}|jdd| ddgd}t|| q8ttjdtjdd}t|d |d| |dd}|jt|jdd dgd}t|| ttjdtjdtjdd d!d"}t|d# |d#| |d#d$}|jt|jdd dgd}t|| W 5 Q R X t | .}ttd%td%d&d'd!}|d( d t|d(< |jd)|dd |jd)d*d+}||d, d-k }t|| tj|jd< ||d, d-k }|jd.|ddd/ |jd.d*d+}t|| ttd%td%d&d'd!}|d( d t|d(< tj|jd0< ||d, d-k }|jd1|dd |jd1d*d+}t|| W 5 Q R X t | T}t }||d dk }|jd|dd td}|jdd2gd+}t|| W 5 Q R X d S )3N
2012-01-01,  r<   )tsr   r$   rP   r   rG   ts>=Timestamp('2012-02-01')
2012-02-01   r   r   r(   r   objectr   r:   r   boolvT)Ttruer   z	boolv == F)Ffalser   r>   )r   r   Zdf_intzindex<10 and columns=['A']r   r!   Zf8Zdtype)r   r   r"   Zdf_floatzindex<10.0 and columns=['A']   colsvaluesfloat64r[   df1
values>2.0r)   r\   g       @df2rH   r"   r   df4z	A>np_zero)r   r   r
   r-   r.   r/   r   rA   r3   rP   r   r1   r2   rB   rU   rK   Zrandlistr"   ZarangerangeapplystrnanilocZmakeDataFramer]   )r4   r5   r$   rD   r6   vZnp_zeror7   r7   r8   test_select_dtypes   s    










rj   c              	   C   s  t | }ttdddtjdtddgd dgd  dgd  d	d
 tdD  d}t|d |jd|ddddgd |	dd}||j
tdk }t|| |	dd}||j
tdk|jdddg@  }t|| dddgdd
 tdD  }|	dd}||j
tdk|j|@  }t|| tdd}|	dd}||j| }t|| t|dks|tt|j
dd j}|	dd}||j
|j }t|| t|dkstW 5 Q R X d S )NrM   rN   rO   a2   bcd   c                 S   s   g | ]}d |dqS rk   Z03dr7   .0ir7   r7   r8   
<listcomp>  s     z0test_select_with_many_inputs.<locals>.<listcomp>)rP   r   r   usersr$   rP   r   r   ru   rG   rQ   rR   z1ts>=Timestamp('2012-02-01') & users=['a','b','c']c                 S   s   g | ]}d |dqS rp   r7   rq   r7   r7   r8   rt   0  s     <   z.ts>=Timestamp('2012-02-01') and users=selector   z
B=selectorr   zts=selector)r   r   r
   r-   r.   r/   rd   r   rA   r3   rP   r   r1   r2   ru   isinr   lenAssertionErrorr   r\   )r4   r5   r$   rD   r6   selectorr7   r7   r8   test_select_with_many_inputs  sL    



r|   c           
   
   C   s^  t | }td}t|d |d| |d}t|jddd}t|}t|| t|jddd}t	|dks|t
t|}t|| t|jddd}t|}t|| W 5 Q R X t| p}td}||d	 d
}tjt|d t|d	dd W 5 Q R X tjt|d t|d	dd W 5 Q R X W 5 Q R X t| j}td}|j|ddd tt|ddd}t|}t	|dkst
t|| t|t|d W 5 Q R X t | }td}|jd|dd tdjdjd}	d|	d< |d|	 t||	gdd}|jddgdd}t|jddgddd}t|}t|| W 5 Q R X d S )Ni  r$   T)iteratorro   	chunksizerS      Zdf_non_tablez0can only use an iterator or chunksize on a tablematchr%   r&   r^   rG   {}_2r(   r   r   r`   r   r@   )r{   )r{   r   )r   r1   rJ   r   rA   r3   rc   r   r2   ry   rz   r   to_hdfpytestraises	TypeErrorr   renamer'   select_as_multiple)
r4   r5   r$   r6   resultsrD   pathmsgr^   r`   r7   r7   r8   test_select_iteratorC  sX    





"



r   c           	   	   C   s  d}t | }tdd}t|d |d| |jd }|jd }|d}t|| d| d}|jd|d	}t|| d
| d}|jd|d	}t|| d| d| d}|jd|d	}t|| W 5 Q R X t | }tdd}t|d |d| |jd }|jd }t|jd|d}t	|}t|| d| d}t|jd||d}t	|}t|| d
| d}t|jd||d}t	|}t|| d| d| d}t|jd||d}t	|}t|| W 5 Q R X d S )N     @ Sr$   r   
index >= ''r)   
index <= '' & index <= 'r~   r*   r   )
r   r1   rJ   r   rA   r"   r3   r2   rc   r   )	r4   r   r5   r6   beg_dtend_dtrD   r*   r   r7   r7   r8   "test_select_iterator_complete_8014  sN    







r   c           
   	   C   s  d}t | }tdd}t|d |d| |jd }|jd }d| d}t|jd||d	}t|}||j|k }	t	|	| d
| d}t|jd||d	}t|}||j|k }	t	|	| d| d| d}t|jd||d	}t|}||j|k|j|k@  }	t	|	| W 5 Q R X t | d}tdd}t|d |d| |jd }d| d}t|jd||d	}dt
|kstW 5 Q R X d S )Nr   r   r   r$   r   r   r   r   r   r   r   z	index > 'r   r   r1   rJ   r   rA   r"   rc   r3   r   r2   ry   rz   
r4   r   r5   r6   r   r   r*   r   rD   Z	rexpectedr7   r7   r8   &test_select_iterator_non_complete_8014  s<    





r   c           
   	   C   s  d}t | p}tdd}t|d |d| |jd }|j|d  }d| d}t|jd||d	}t|}||j|k }	t	|	| d
| d}t|jd||d	}t
|dkstt|}||j|k }	t	|	| d| d| d}t|jd||d	}t
|dkstt|}||j|k|j|k@  }	t	|	| d
| d| d}t|jd||d	}t
|dksztW 5 Q R X d S )Ni'  i r   r$   r   r   r   r   r   r   r   z' & index >= 'r   r   r7   r7   r8   &test_select_iterator_many_empty_frames  s6    

	r   c           
   
   C   s  t  }t| }|jd|dd |jt|d  }td}|jjd |ksPt	d}d}|
d||g}|j|d d	d
gf }t || |
d|g}|jd d d	gf }t || t  }|d| d}	tjt|	d |
dd W 5 Q R X W 5 Q R X d S )Nframer%   r&   r   zindex>=datedatezcolumns=['A', 'D']z	columns=Ar   rI   Zdf_timez%could not convert string to Timestampr   zindex>0)r1   rJ   r   r0   r"   ry   r   envscoperz   r3   rB   r2   rA   r   r   
ValueError)
r4   r$   r5   r   Zcrit1Zcrit2Zcrit3rD   r6   r   r7   r7   r8   test_frame_select=  s&    
r   c              
   C   s  t  }d|d< d|j|jdd df< t| }|jd|ddgd |dd	}|j|j|jd
 k|jdk@  }t || |dd}|j|j|jd
 k|jdk@  }t || |dd}|j|j|jd
 k|jdkB  }t || |dd}|j|j|jd
 k|j|jd k@ |jdkB  }t || |dd}|j|jdk }t || d}t	j
t|d |dd W 5 Q R X |dd}|jd d |jddgf }t || |dd}|j|j|jd
 k jddgd}t || W 5 Q R X d S )Nr   stringr   r   r:   r$   r%   r'   rH   z index>df.index[3] & string="bar"r   z index>df.index[3] & string="foo"z index>df.index[3] | string="bar"z7(index>df.index[3] & index<=df.index[6]) | string="bar"   zstring!="bar"z6cannot use an invert condition when passing to numexprr   z~(string="bar")z~(columns=['A','B'])r   r   z(index>df.index[3] & columns in ['A','B']r(   )r1   rJ   rB   r"   r   r0   r3   r   r2   r   r   NotImplementedErrorr#   
differencerK   )r4   r$   r5   rD   r6   r   r7   r7   r8   test_frame_select_complex`  sH        "r   c              	   C   s  t ddg}|\}}tddddddgi}|j|ddd	dgd
 t|ddd}ttjdddgtjdd t	dD ddgdd}|j|ddd	d t|ddd}|j
 }t|}	|	jddd}
t|
| |	  t|ddd}
t|
| |j
}t|ddd}
t|
| t|ddd}
t|
| t|ddd}
t|
| t|ddd}
t|
| t|}	|	jddd}
t|
| |	jddd}
t|
| |	jddd}
t|
| |	jddd}
t|
| |	  W 5 Q R X d S )Nz
params.hdfzhist.hdfr   r   r   r   r$   wr%   )moder'   rH   zA=[2,3]r)      datac                 S   s"   g | ]}t d D ]}||fqqS )rS   )rd   )rr   rs   jr7   r7   r8   rt     s     
  z.test_frame_select_complex2.<locals>.<listcomp>rS   l1l2)r   )r#   r"   )r   r'   zl1=[2, 3, 4]zl1=l0zl1=indexzl1=selection.indexzl1=selection.index.tolist()zl1=list(selection.index))r   r   r   r   r-   r.   r/   r   from_tuplesrd   r"   tolistr   r3   r1   r2   close)r4   pathspphhparamsZ	selectionhistr6   Zl0r5   rD   r"   r7   r7   r8   test_frame_select_complex2  sR     
r   c              
   C   s   t  }t| f}|jd|dd d}tjt|d |dd W 5 Q R X tjt|d |dd W 5 Q R X W 5 Q R X d S )Nr$   r%   r&   z unable to collapse Joint Filtersr   zcolumns=['A'] | columns=['B']z!columns=['A','B'] & columns=['C'])r1   rJ   r   r0   r   r   r   r3   )r4   r$   r5   r   r7   r7   r8   test_invalid_filtering  s    
r   c              	   C   st  t | `}t }d|d< d|j|jdd df< |jd|dgd |dd}||jdk }t|| |dd	}||jdk }t|| |	 }t
j|j|jdkdf< |jd
|dgd |d
d	}|t|j }t|| d|d< d|j|jdd df< |jd|dgd |dd}||jdk }t|| |dd}||jdk }t|| W 5 Q R X d S )Nnonex r      r$   rG   zx=nonezx!=noner`   r   intdf3zint=2zint!=2)r   r1   rJ   rB   r"   rA   r3   r   r2   copyr-   rg   r   r   )r4   r5   r$   rD   r6   r`   r7   r7   r8   test_string_select  s4    r   c              
   C   s  t  }t  jdjd}d|d< t| v}d}tjt|d |jd ddgd	d
 W 5 Q R X |j	d	|ddgd |	d| tjt|d |jd ddgd	d
 W 5 Q R X tjt|d |jd gddgd	d
 W 5 Q R X d}tjt
|d |jd	dgddgd	d
 W 5 Q R X tjt
|d |jdgddgd	d
 W 5 Q R X tjt
dd |jd	dgddgdd
 W 5 Q R X |d	ddg}|jd	gddgd	d
}t || |jd	ddgd	d
}t || |jd	dgddgd	d
}t||gdd}||jdk|jdk@  }t j||dd |jd	dgddd
}t||gdd}|dd  }t || |	dt jdd d}tjt|d |jd	dgddgd	d
 W 5 Q R X W 5 Q R X d S )Nr   r(   r   r   zkeys must be a list/tupler   zA>0zB>0r^   )r*   r{   r   r   rG   r`   z!'No object named df3 in the file'r   z!'No object named df4 in the file'rb   r   r@   r   F)Z
check_freqzindex>df2.index[4]rS   rl   )Znperz,all tables must have exactly the same nrows!)r1   rJ   r   r'   r   r   r   r   r   rA   KeyErrorr3   r2   r   r   r   r   )r4   r^   r`   r5   r   rD   r6   r7   r7   r8   test_select_as_multiple  sz                  r   c              	   C   s   t | }ttdtdddd}|d d t|d< tj|jd< tdd	d
gdddgddddgd}|jd|ddgd |j	ddd}t
|| W 5 Q R X d S )Nr   rZ   r]   rX   r[   r   r   z13.0z14.0z15.0g      @g      @g      @r   r:   rS   )r"   r$   Tra   r_   r)   )r   r   rd   re   rf   r-   rg   rh   rA   r3   r1   r2   )r4   r5   r$   r6   rD   r7   r7   r8   test_nan_selection_bug_4858]  s    
r   c              
   C   s~   t ddddddddgdddd	d
dddgd}||jdk }t| $}|jd|ddd |dd}W 5 Q R X t|| d S )Nrk   rn   rm   ztest & tester   r   r   r:   rS   r   r      )rk   rm   testr%   Tr   za = "test & test")r   rk   r   rA   r3   r1   r2   )r4   r$   r6   r5   rD   r7   r7   r8   (test_query_with_nested_special_characterp  s    
r   c              	   C   s   t ddddgi}t| }|jd|ddd d	}|dd
|d}|jsPtd}|dd|d}|jddgd d f }t|| d}|dd|d}|jdgd d f }t|| W 5 Q R X d S )Nr   g~  eAg$  eAg'1  eAr   r%   Tr   g  eAzA < z.4fg   eAzA > r   r   zA == )	r   r   rA   r3   emptyrz   rB   r1   r2   )r4   r$   r5   cutoffrD   r6   exactr7   r7   r8   test_query_long_float_literal~  s    

r   c                 C   s  t ddgtdddddgddgdd	d
ddgd}t| }|jd|ddd td}|jddd}|jdgd d f }t|| dD ]V}ddtdt	
ddfD ]N}d| d}d| dt| d}	tjt|	d |jd|d W 5 Q R X qd}dD ]B}
|
 d| d}d }	tjt|	d |jd|d W 5 Q R X q td!d"dgddd
gD ]\}}
|
 d| d}|jd|d}|d#kr|jd$gd d f }n2|d%kr|jdgd d f }n|jg d d f }t|| qZqW 5 Q R X d S )&Nz
2014-01-01z
2014-01-02r   rO   g?g333333?r   )r   	real_datefloatr   r   r   r   r   r(   r   r%   Tr   zreal_date > tsr)   )<>==g @r,   zdate z vzCannot compare z	 of type z to string columnr   rk   )r   r   r    zcould not convert string to 1z1.1r   r   r   )r   r   r   rA   r   r3   rB   r1   r2   pdZ	Timedeltatyper   r   r   r   zip)r4   r$   r5   rP   rD   r6   opri   queryr   colr7   r7   r8   test_query_compare_column_type  sL    



 

r   r*   r   r7   )Nc              
   C   sb   t dddg}tdB}t|.}|d|d t|d| d}t|| W 5 Q R X W 5 Q R X d S )Nr   r   r   zempty_where.h5r$   tr)   )r   r   r   r0   r   r1   r2   )r*   r$   r   r5   rD   r7   r7   r8   test_select_empty_where  s    

r   )3warningsr   numpyr-   r   Zpandas._libs.tslibsr   Zpandasr   r   r   r   r   r   r	   r1   r
   r   r   r   r   Zpandas.tests.io.pytables.commonr   r   r   Zpandas.io.pytablesr   markZ
single_cpuZ
pytestmarkr9   rE   rL   rj   r|   r   r   r   r   r   r   r   r   r   r   r   r   r   r   Zparametrizer   r7   r7   r7   r8   <module>   s:   4?)k2EH49#6A+H3