U
    -e                    @   s>  d dl mZ d dlmZ 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 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e	jdddddge fdddd ge fd!d"dge fd#dd$ge fg e e! fddgd%gge e fd#gdd$d&gge e fd'd(dge"e#d#e e#d$e gfgd)d* Z$e	jj%d+d, Z&e	jj%d-d. Z'e	jj%d/d0 Z(e	jj%d1d2 Z)e	jj%d3d4 Z*e	jj%e	jj+d5d6 Z,d7d8 Z-d9d: Z.d;d< Z/d=d> Z0d?d@ Z1dAdB Z2dCdD Z3dEdF Z4dGdH Z5dIdJ Z6dKdL Z7dMdN Z8dOdP Z9dQdR Z:dSdT Z;dUdV Z<dWdX Z=dYdZ Z>d[d\ Z?d]d^ Z@d_d` ZAdadb ZBe	jj%dcdd ZCdedf ZDdgdh ZEdidj ZFdkdl ZGe	jj%dmdn ZHdodp ZIdqdr ZJdsdt ZKdudv ZLdwdx ZMdydz ZNd{d| ZOd}d~ ZPdd ZQdd ZRdd ZSdd ZTdd ZUdd ZVdd ZWdd ZXdd ZYdd ZZdd Z[dd Z\dd Z]dd Z^dd Z_e	jj%dd Z`dd Zadd Zbdd Zcdd Zddd Zedd Zfdd Zge	jddddgdddgfehfdddgdddgfejifdgddggdgddggfejjfgdd Zke	jdejlejmgdd Zne	jddddgdddgfejifdgddggdgddggfejjfgdd Zoe	jddddgdddgfehfdddgdddgfejifdgddggdgddggfejjfgdd Zpe	jdejlejmgdd Zqe	jj%dd Zre	jj%dd ZsddĄ Zte	jj%ddƄ ZuddȄ Zvddʄ Zwdd̄ Zxdd΄ ZyddЄ Zzdd҄ Z{ddԄ Z|ddք Z}e	jj~dd؄ Zddڄ Ze	jj~dd܄ Ze	jj~ddބ Ze	jj~dd Ze	jj~dd Ze	jj~dd Zdd Zdd Zdd Ze	jdejejgdd ZdS )    )OrderedDict)IterableNc               	   C   s  t jg t  d} | jt  ks$t|  g ks4t|   t jg dd}|jt  ks\tt	t
 t g  W 5 Q R X t dddgdddgd	d
dgg} t| jtsttdd | jD sttdd |  D stt| jdkst|  tdd |  D kstt| t| |   ks8t| jdksHt|   t| }| d k	sjt~ | d ks|td S )Ntypebinary                        	   c                 s   s   | ]}t |tjjV  qd S N
isinstancepalibZ
Int64Array.0c r   Y/var/www/html/Darija-Ai-Train/env/lib/python3.8/site-packages/pyarrow/tests/test_table.py	<genexpr>0   s     z,test_chunked_array_basics.<locals>.<genexpr>c                 s   s   | ]}t |tjjV  qd S r   r   r   r   r   r   r   1   s     c                 s   s   | ]}|  V  qd S r   )get_total_buffer_sizer   r   r   r   r   3   s   H   )r   chunked_arraystringr   AssertionError	to_pylistvalidater   pytestraises
ValueErrorr   chunkslistall
iterchunkslenr   sumsys	getsizeofobject
__sizeof__nbytesweakrefref)datadata2wrr   r   r   test_chunked_array_basics   s>    
r5   c               	   C   s  t dddgdddgddd	gg} | jt  ks4tt| d	ksDtt| jdksVtt dddgd
ddgddd	gg} | jt  kstt| d	kstt| jdkstt jdddgd
ddgddd	ggt  d} | jt  kstt| d	kstt| jdkstt dddgg g} | jt  ks4tt| dksFtt| jdksZtd}tj	t
|d t g s~tW 5 Q R X t jg t  djt  kstt g gjt  kstt jg gt  djt  kstd S )Nr   r   r	   r
   r   r   r   r   r         @      @      @r   z@cannot construct ChunkedArray from empty vector and omitted typematch)r   r   r   int64r   r)   r%   int8r"   r#   r$   r   null)arrmsgr   r   r   test_chunked_array_construction@   sN    "r@   c                  C   sH   t ddg} t | | g}| }t ddddg}||sDtd S )Nr   r   )r   arrayr   combine_chunksequalsr   )r>   chunked_arrresexpectedr   r   r   test_combine_chunksm   s
    rG   c                   C   sl   t jg t  d t jg t  dks,tt jt jg t  dgt  d t jg t  dkshtd S )Nr   )r   r   bool_rB   rA   r   r   r   r   r   4test_chunked_array_can_combine_chunks_with_no_chunksv   s      rI   c                  C   s`   t dddgdddgg g} t| }|  }t|tjs>t|jdksLtt	||s\td S )Nr   r   r	   r
   r   r   )r   )
r   r   npasarrayZto_numpyr   Zndarrayr   shapeZarray_equal)r2   arr1Zarr2r   r   r   test_chunked_array_to_numpy   s    
rN   c               	   C   s   d} t jt| d, ttdddgtdddgg W 5 Q R X t jt| d& tjtdddggt d	 W 5 Q R X d S )
Nzchunks must all be same typer9   r   r   r	         ?       @      @r   )r"   r#   	TypeErrorr   r   rA   float64)r?   r   r   r   !test_chunked_array_mismatch_types   s    rT   c                  C   s>   t dddgt dddgg} t | } t| dks:td S )Nr   r   r	   r
   r   r   z<[
  [
    1,
    2,
    3
  ],
  [
    4,
    5,
    6
  ]
])r   rA   r   strr   r2   r   r   r   test_chunked_array_str   s
    
rW   c               	   C   s.  t dddgt dddgg} t | } | d  dks>t| d  dksRt| d  dksfttt | d  W 5 Q R X tt | d	  W 5 Q R X | t	d  dkst| dd }|
 ddgkst| dd }|
 dgkst| d
d
 }|j| jkst|
 g ks*td S )Nr   r   r	   r
   r   r   iic   )r   rA   r   as_pyr   r"   r#   
IndexErrorrJ   int32r    r   )r2   
data_slicer   r   r   test_chunked_array_getitem   s&    
r^   c                  C   s   t dddgt dddgg} t | } | t| }|j| jksHt| g ksXt| t| d }|j| jkszt| g kstt jj	| gdgd	}|t|}t|d
kstt jj	| gdgd	}|t|d }t|d
kstd S )Nr   r   r	   r
   r   r   
   anamesr   )
r   rA   r   slicer)   r   r   r    Tablefrom_arrays)r2   r]   tabletable_slicer   r   r   test_chunked_array_slice   s     
rh   c                  C   s|   t dgt dddgt dddgt dd	d
gg} t | }ttd|D ]\}}|| ksPtqPt|tsxtd S )Nr   r   r   r	   r
   r   r   r   r   r   r_   )	r   rA   r   ziprangerZ   r   r   r   )r2   r>   ijr   r   r   test_chunked_array_iter   s    

rm   c                  C   s*  dd } dd }| t jg t  dt jg t  d |t jg t  dt jg t  d t jddgt  d}t jddgt  d}t jddgt  d}t jddddgt  d}| |g|g ||g|g | ||g||g | ||g|g |||g||g t jg t  dd r&td S )	Nc                 S   st   t | tjr| }n
t| }t |tjr.|}n
t|}||sFt||sTt||ks`t|t|ksptd S r   )r   r   ChunkedArrayr   rC   r   rU   ZxarrsZyarrsxyr   r   r   eq   s    

z%test_chunked_array_equals.<locals>.eqc                 S   sd   t | tjr| }n
t| }t |tjr.|}n
t|}||rFt||rTt||ks`td S r   )r   r   rn   r   rC   r   ro   r   r   r   ne  s    

z%test_chunked_array_equals.<locals>.ner   r   r   r	   )r   r   r\   r;   rA   rC   r   )rr   rs   r`   br   dr   r   r   test_chunked_array_equals   s$    rv   )r2   typTFr   r   r
   r   rO         @r`   rt   r	   r   )r   r`   )r   r   c                 C   sr   g }| r2| tj| d d |d | dd  } qtj||d}|  tt|}|  ||snt	d S )Nr   r   )
appendr   rA   r   r!   pickleloadsdumpsrC   r   )r2   rw   arraysrA   resultr   r   r   test_chunked_array_pickle"  s    r   c                  C   s   dd l } tdddddgg}tj|dgd}|d}t|tjsHt| }t|| j	s`t|j
dksnt|d dks~t|jdkstd S )	Nr   r   r_   r`   ra   )r   )pandasr   rA   rf   columnr   rn   r   	to_pandasSeriesrL   name)pdr2   rf   colZseriesr   r   r   test_chunked_array_to_pandas<  s    
r   c                  C   s   dd l } dd lm} tdddgt| dddgt| jdddt| jddd	d
tdddgtdtdddgtdd	fD ]L}td|i}|	d
 }|jdkst| j|
 dd}||| qd S )Nr   r   r   r	   r`   rt   Z2012)periodszEurope/Brussels)r   tzmsr   )r   )r   pandas.testingtestingr   rA   ZCategoricalZ
date_range	timestamprf   r   r   r   r   r   Zassert_series_equal)r   tmr2   rf   r~   rF   r   r   r   *test_chunked_array_to_pandas_preserve_nameM  s    r   c                  C   s   dd l } | j| dddd}t|}| }|jdks>t|jdksLt|jdksZt|j	
|j	slt| j| dddd}t|}| }|jdkst|jdkst|jdkst|j	
|j	std S 	Nr   r_   r   )index)r_   r   r	   r
   )r
   r   )r   	DataFrame
RangeIndexr   rf   r   num_rowsr   rL   r   rC   )r   r2   rf   r~   r   r   r   .test_table_roundtrip_to_pandas_empty_dataframea  s    

r   c                  C   s   dd l } | j| dddd}tj|}| }|jdks@t|j	dksNt|j	dks\t|j
|j
snt| j| dddd}tj|}| }|jdkst|j	dkst|j	dkst|j
|j
std S r   )r   r   r   r   RecordBatchfrom_pandasr   r   r   rL   r   rC   )r   r2   batchr~   r   r   r   4test_recordbatch_roundtrip_to_pandas_empty_dataframe{  s    r   c                  C   sl   dd l } dd lm} | ddgddgd}t|}|j  }|j	dksRt
|||jd d  d S )Nr   r   r   g?g?r`   rt   r   r   )r   r   r   r   r   rf   schemaZempty_tabler   rL   r   assert_frame_equalZiloc)r   r   dfrf   r~   r   r   r   test_to_pandas_empty_table  s    
r   c               	   C   s  t dgt dddgg} t | }t|}| ddddgksHt|jtdks\ttj|dd}| dd	d
dgkstt dd gt dddgg} t | }t|}| }|d dkstt|d st|dd  dddgkst|jtdkstt j	
t dddddgt dddg}t ||g}t|}|jtdksft| dddddgd kstd S )Nr   r   r   r	   r;   rU   dtype0123rO   rP   rQ   rS   r`   rt   r   r-   )r   rA   r   rJ   rK   tolistr   r   isnanZDictionaryArrayre   )r2   rD   Znp_arrelementsr>   r   r   r   test_chunked_array_asarray  s6    




 
r   c                  C   s  t t dt  t dt  g} t jdddg| d}t |}| \}}|t t jddd	gt  dsxt	|t t jd
ddgt  dst	t jg | d}t |}| \}}|t t jg t  dst	|t t jg t  ds
t	d S )Nrp   rq   r   rx   r	         @)r         @r   r   r	   r   rx   r   r   )
r   structfieldint16float32rA   r   flattenrC   r   )tyr`   Zcarrrp   rq   r   r   r   test_chunked_array_flatten  s    
(
"r   c                  C   s   t t ddd dg t dd dg g} | djt ddgsPt| djt ddgspt|  } t dddg}| dj|st| dj|st| 	 ddd ddd dgkstd S )Nfoobarquuxr   r   )
r   r   rA   dictionary_encodechunk
dictionaryrC   r   unify_dictionariesr    )r>   expected_dictr   r   r   %test_chunked_array_unify_dictionaries  s      r   c               	   C   s  t jtdddt jdddd dgddg} t | d	d
g}|jjrFtt|dksVt|jdksdt|j	t| ksvt|
 dkst|jdk t|t||
  kst| }|td	dddddgfd
dddd dgfgkstt|tksttt |d  W 5 Q R X t jt jd	t  ddidt d
t  gddid}t j| |d}|j|ksjtt | |}|j|kstt|dkst|jdddkstt|}| d k	st~| d kstd S )Nr   r   r   r   r   r   r_   r\   c0c1   r   r   r	   r
   keyvaluemetadata   foo   barr   zSpyarrow.RecordBatch
c0: int16
c1: int32
----
c0: [0,1,2,3,4]
c1: [-10,-5,0,null,10]TZshow_metadataznpyarrow.RecordBatch
c0: int16
  -- field metadata --
  key: 'value'
c1: int32
-- schema metadata --
foo: 'bar')r   rA   rj   record_batchr   r   r   r)   r   num_columnsr   r/   r+   r,   r-   r.   	to_pydictr   r   dictr"   r#   r[   r   r   r\   rU   	to_stringr0   r1   )r2   r   pydictr   r4   r   r   r   test_recordbatch_basics  sP    

	
r   c                	   C   s&   t jtdd t  W 5 Q R X d S )Nr   r9   )r"   r#   rR   r   r   r   r   r   r   test_recordbatch_dunder_init(  s    r   c                  C   s   t jtdddt jdddd dgddg} t | d	d
g}g }| D ]}|| qF|j|ksdt|t j||jdks|t|t j|dd  |jdd  dkst||kstd S )Nr   r   r   r   r   r   r_   r\   r   r   ra   r   )	r   rA   rj   r   itercolumnsry   columnsr   column_names)r2   r   r   r   r   r   r   test_recordbatch_itercolumns-  s    (r   c                  C   s  t jtdddt jdddd dgddg} t d	d
dgt dgdgdggg}ddg}t | |}|t | |ksxt|t | |st|t ||kst|t ||rtt j| |ddid}||st|j|ddrt|d rt|dkstd S )Nr   r   r   r   r   r   r_   r\   r`   rt   r   ru   efr   r   r   r   rb   r   TZcheck_metadatar   )r   rA   rj   r   r   rC   )Zdata1r3   r   r   Z
batch_metar   r   r   test_recordbatch_equals>  s&    r   c               
   C   s   t t dddd dgt ddddd	ggd
dg} | t ddg| ddsZt| t dd gt t dd gt dd ggd
dgstd S Nr   r   r	   r   r`   rt   r   ru   r   f1f2)r   r   rA   takerC   rc   r   r   r   r   r   test_recordbatch_takeZ  s    &r   c                  C   s6   t jt ddddggdgd} | d jdks2td S Nr   r   r	   r
   Za0ra   r   )r   r   rA   _namer   )rbr   r   r   )test_recordbatch_column_sets_private_namee  s     r   c               	   C   sV   t ddg} t t dt t  g}tt t j	| g|d W 5 Q R X d S )Nr   r   Zf0r   )
r   rA   r   r   list_utf8r"   r#   NotImplementedErrorr   )r>   r   r   r   r   ,test_recordbatch_from_arrays_validate_schemak  s    r   c               	   C   sR   t dgt dddgt dgg} tt t | dddg W 5 Q R X d S )	Nr   ZtokyolikeZhappyZderekidtagsr   )r   rA   r"   r#   r$   r   rV   r   r   r   -test_recordbatch_from_arrays_validate_lengthss  s
    
r   c                  C   s<   t g g } t| dkst| jdks*t| jdks8td S )Nr   )r   r   r)   r   r   r   r   r   r   r   test_recordbatch_no_fields|  s    r   c               	   C   sv   t tdt dddddgg} tt t j| dddgd	 W 5 Q R X tt t j| dgd	 W 5 Q R X d S 
Nr   r   r   r   r_   r`   rt   r   ra   )r   rA   rj   r"   r#   r$   r   rV   r   r   r   *test_recordbatch_from_arrays_invalid_names  s    r   c                  C   sF   t tdt dddddgg} t | ddg}|jjd ksBtd S )Nr   r   r   r   r_   r   r   )r   rA   rj   r   r   r   r   r2   r   r   r   r   test_recordbatch_empty_metadata  s
    r   c                  C   s   t jtdddt jdddddgddg} t d	t  t d
t  g}t j|ddid}t j| |d}t	t
|}||st|j|kstd S )Nr   r<   r   r   r   r   r_   r   intsfloatsr   r   r   r   )r   rA   rj   r   r<   r   r   r   rz   r{   r|   rC   r   )r2   fieldsr   r   r~   r   r   r   test_recordbatch_pickle  s    r   c               	   C   s   t tdt dddddgt tddg} t jj| dd}|d|jds\t|d|jdsxtt	
t |d	 W 5 Q R X t	
t |d  W 5 Q R X t	
t |d
 W 5 Q R X d S Nr   r   r   r   r_   r`   rt   r   ra   r`   ru   r
   )r   rA   rj   r   re   r   rC   r   r   r"   r#   KeyErrorrR   r[   r   r   r   r   test_recordbatch_get_field  s    r   c               	   C   s   t tdt dddddgt tddg} t jj| dd}|d|dsZttj	t
d	d
 |d W 5 Q R X t	t |d  W 5 Q R X t	t |d W 5 Q R X d S Nr   r   r   r   r_   r   ra   r`   z"Field "d" does not exist in schemar9   ru   r
   )r   rA   rj   r   re   r   rC   r   r"   r#   r   rR   r[   r   r   r   r   test_recordbatch_select_column  s     r   c               	   C   s4  t dddd dg} t ddddd	g}t ddgdd
gddgd ddgg}t | ||gdddg}|dg}t | gdg}||st|ddg}t ||gddg}||st|dg}t | gdg}||st|ddg}t ||gddg}||st|ddi}|ddg}d|jjksHtt	j
tdd |dg W 5 Q R X t	j
tdd |dg W 5 Q R X |ddg}t ||gddg}||stt | ||gdddg}t	j
tdd |dg W 5 Q R X |dg}t |gdg}||s0td S Nr   r   r	   r   r`   rt   r   ru   r   r
   r   r   r_   r   r   Zf3r   test   azField "f5" does not existr9   Zf5zindex out of boundszField "f1" exists 2 times)r   rA   r   selectrC   r   replace_schema_metadatar   r   r"   r#   r   r[   )a1a2a3r   r~   rF   batch2r   r   r   test_recordbatch_select  s@    $r  c                	   C   s0   t t tjttd W 5 Q R X d S )Nr   )r"   r#   rR   r   r   from_struct_arrayrA   rj   r   r   r   r   *test_recordbatch_from_struct_array_invalid  s    r  c               	   C   s   t jddiddigt dt  fdt  fgd} t j| }|t jt jdd gt  dt jd dgt  dgddgst	d S )Nr   r   r   rO   r   )
r   rA   r   r\   r   r   r  rC   re   r   )Zstruct_arrayr~   r   r   r   "test_recordbatch_from_struct_array  s    
r	  c              	   C   s  t tdt dddddgg}ddg}| ||d}|d	}|jd
ksNt| dd |D |d}||srt|d	d	}| dd |D |d}||st|d|stt|t|dkstt	t
 |d W 5 Q R X |dd|d d st|dd	|d d	 s0t|d	d	|d	d sPt|d	t|d	 |d	d  sxt|t|d	 d	|dd  st|t|d d	|dd std S )Nr   r   r   r   r_   r   r   ra   r   r	   c                 S   s   g | ]}| d qS r   rc   r   rp   r   r   r   
<listcomp>'  s     z+_table_like_slice_tests.<locals>.<listcomp>c                 S   s   g | ]}| d d qS r
  r  r  r   r   r   r  +  s     rX   r
   )r   rA   rj   rc   r   r   rC   r)   r"   r#   r[   )factoryr2   rb   objZslicedrF   Zsliced2Z	expected2r   r   r   _table_like_slice_tests  s,    
   ((r  c                   C   s   t tjjS r   )r  r   r   re   r   r   r   r   test_recordbatch_slice_getitem@  s    r  c                   C   s
   t tjS r   )r  r   rf   r   r   r   r   test_table_slice_getitemD  s    r  c                  C   sZ   t dt jg t ddi} | dd}|  t dt jg t  di} |   d S )Nr`   usr   r   )r   rf   rA   r   rc   r   r   r   )rf   rg   r   r   r   test_slice_zero_length_tableH  s
    r  c               	   C   sz   t jdgdd} t jddgdd}tt| gdg}tt|gdg}ttj tj||g W 5 Q R X d S )Nr   uint32r   r6   r7   rS   r   )	rJ   rA   r   r   r"   r#   ArrowInvalidrd   from_batches)r  r  batch1r  r   r   r   "test_recordbatchlist_schema_equalsS  s    r  c                  C   s6   t jt ddddggdgd} | d jdks2td S r   )r   rf   rA   r   r   )tr   r   r   #test_table_column_sets_private_name]  s     r  c                  C   sf   t jjg g d} | | st| d r,tt jjg g ddid}| j|ddrTt| |sbtd S )Nra   r   r   r   Tr   )r   rd   re   rC   r   )rf   otherr   r   r   test_table_equalsc  s    r  c               	   C   s(  t t dt  t dt  g} t jt dgt dggddgd}t j|g| }|j	| slt
|dt dggkst
|dt dggkst
t t dt  g}tt j t j|g| W 5 Q R X t t dggdg}tt j t j|g|  W 5 Q R X d S )Nr`   rt   r   gQ	@ra   r   )r   r   r   r;   rS   r   rA   rd   r  rC   r   r   r   r"   r#   r  )r   r   rf   Zincompatible_schemaZincompatible_batchr   r   r   "test_table_from_batches_and_schemao  s"    r   c            
      C   s.  ddl m}  dd l}|dttdi}|dttddi}tjj|dd}tjj|dd}tj	
|||g}|j|||gdd	}| }t|d
kst| tj	
| | |jdd}ttt|ddddgkst| | | | tj	
| | tj	
t|||g}	||	s*td S )Nr   )r   r`   r_      F)preserve_indexT)Zignore_indexr	      Zmax_chunksizer   )r   r   r   r   r&   rj   r   r   r   rd   r  concat
to_batchesr)   r   r   mapiterrC   )
r   r   df1df2r  r  rf   Zexpected_dfZbatchesZtable_from_iterr   r   r   test_table_to_batches  s*    r+  c               
   C   s  t jtdddt jdddddgddg} t j| dd	}|  t|dksPt|jdks^t|jd
kslt|j	dkszt|
 dkst|jdkstt|t||
  kst| }|tdddd
ddgfddddddgfgkstt|tkstg }| D ]p}|| | D ]}|d k	s$tq$tt |d W 5 Q R X tt ||j W 5 Q R X q|j|kst|t j||jd	kst|t j|dd  |jdd  d	kst||kstt|}| d k	st~| d kstd S )Nr   r;   r   r   r   r   r_   r   ra   r   )r   r   P   r`   r   r	   r
   rt   rX   ) r   rA   rj   rf   r!   r)   r   r   r   rL   r   r/   r+   r,   r-   r.   r   r   r   r   r   ry   r(   r"   r#   r[   r   
num_chunksr   r   r0   r1   )r2   rf   r   r   r   r   r4   r   r   r   test_table_basics  sN    

*
r.  c                	   C   s&   t jtdd t  W 5 Q R X d S )Nrd   r9   )r"   r#   rR   r   rd   r   r   r   r   test_table_dunder_init  s    r/  c                  C   s   t ddg} t ddg}t jdt  tdddd	}t jd
t  dd}t jj| |gt ||gd}d|djks~t	|dj
dkst	d S )Nr   r   r	   r
   field1ABr   r   Zfield2Fnullabler   r   r   )r   rA   r   r;   r   rd   re   r   r   r   r4  )Zarr0rM   Zfield0r0  rf   r   r   r   0test_table_from_arrays_preserves_column_metadata  s    r5  c               	   C   sz   t tdt dddddgg} tt t jj| dddgd	 W 5 Q R X tt t jj| dgd	 W 5 Q R X d S r   )r   rA   rj   r"   r#   r$   rd   re   rV   r   r   r   $test_table_from_arrays_invalid_names  s     r6  c                  C   s   t tddddddgg} tj| ddgd}tjj| ddgd}||sNttt	dt
 t	dt g}tj| |d	}tjj| |d	}||std S )
Nr   r   r   r   r_   r`   rt   ra   r   )r&   rj   r   rf   rd   re   rC   r   r   r   Zuint16r;   )r2   r~   rF   r   r   r   r   test_table_from_lists  s    
r7  c                  C   s   t jddgddggt  dt jddd dggt  dg} t jt d	t  t d
t  gddid}t jj| |d}t	t
|}|  ||std S )Nr   r   r	   r
   r   Zsomestrings r   strsr   r   r   r   )r   r   r  r   r   r   rd   re   rz   r{   r|   r!   rC   r   )r2   r   rf   r~   r   r   r   test_table_pickle  s    r;  c               	   C   s   t tdt dddddgt tddg} t jj| dd}|d|jds\t|d|jdsxtt	
t |d	 W 5 Q R X t	
t |d  W 5 Q R X t	
t |d
 W 5 Q R X d S r   )r   rA   rj   rd   re   r   rC   r   r   r"   r#   r   rR   r[   r2   rf   r   r   r   test_table_get_field  s    r=  c               	   C   s   t tdt dddddgt tddg} t jj| dd}|d|dsZttj	t
d	d
 |d W 5 Q R X t	t |d  W 5 Q R X t	t |d W 5 Q R X d S r   )r   rA   rj   rd   re   r   rC   r   r"   r#   r   rR   r[   r<  r   r   r   test_table_select_column%  s    r>  c               	   C   sf   t jt dddgt dddgt ddd	ggd
dd
gd} tjtdd | d
 W 5 Q R X d S )Nr   r   r	   r
   r   r   r   r   r   r`   rt   ra   z"Field "a" exists 2 times in schemar9   )r   rf   rA   r"   r#   r   r   rf   r   r   r   !test_table_column_with_duplicates:  s    r@  c                  C   s   t tdt dddddgt tddg} t jj| dd}t d| d	 j}|d
|| d	 }||| d	 }t jj| | d	 g dd}|	|st
|	|st
|d|| d	 }t jj| d	 g|  dd}|	|st
d S )Nr   r   r   r   r_   r   ra   ru   r   r	   )r`   rt   r   ru   )ru   r`   rt   r   )r   rA   rj   rd   re   r   r   
add_columnappend_columnrC   r   )r2   rf   	new_fieldt2t3rF   Zt4r   r   r   test_table_add_columnE  s$    rF  c                  C   s   t tdt dddddgt tddg} t jj| dd}t d| d	 j}|d|| d	 }t| }| d	 |d< t jj|d
d}|	|st
d S )Nr   r   r   r   r_   r   ra   ru   r   )ru   rt   r   )r   rA   rj   rd   re   r   r   Z
set_columnr&   rC   r   )r2   rf   rC  rD  expected_datarF   r   r   r   test_table_set_column\  s    rH  c               	   C   s   t td} t dddddg}t tdd}t jj| ||gdd}|dd	g}|d}t jj|gd
d}||stt jj||gdd}||sttj	t
dd |dg W 5 Q R X dS )z& drop one or more columns given labelsr   r   r   r   r_   r   ra   r`   rt   r   rt   r   Column 'd' not foundr9   ru   N)r   rA   rj   rd   re   Zdrop_columnsrC   r   r"   r#   r   r`   rt   r   rf   rD  rE  Zexp_t2Zexp_t3r   r   r   test_table_drop_columnsn  s    
rM  c               	   C   s   t td} t dddddg}t tdd}t jj| ||gdd}|dd	g}|d}t jj|gd
d}||stt jj||gdd}||sttj	t
dd |dg W 5 Q R X dS )z, verify the alias of drop_columns is workingr   r   r   r   r_   r   ra   r`   rt   rI  rJ  rK  r9   ru   N)r   rA   rj   rd   re   ZdroprC   r   r"   r#   r   rL  r   r   r   test_table_drop  s    
rN  c                  C   s|   t tdt dddddgt tddg} t jj| dd}|d}|  t jj| dd  d	d}||sxtd S )
Nr   r   r   r   r_   r   ra   r   rJ  )	r   rA   rj   rd   re   remove_columnr!   rC   r   r2   rf   rD  rF   r   r   r   test_table_remove_column  s    
rQ  c                  C   sz   t tdg} t jj| dgd}|d}|  t|t|ksHt|	d|
d|d }|  ||svtd S )Nr   r`   ra   r   )r   rA   rj   rd   re   rO  r!   r)   r   rA  r   rC   )r2   rf   rD  rE  r   r   r   test_table_remove_column_empty  s    
rR  c               	   C   s>   g } ddg}d}t jt|d tjj| |d W 5 Q R X d S )Nr`   rt   z?Length of names [(]2[)] does not match length of arrays [(]0[)]r9   ra   )r"   r#   r$   r   rd   re   )r2   rb   messager   r   r   test_empty_table_with_names  s    rT  c                  C   s4   t g } | jg kst| t jg g s0td S r   )r   rf   r   r   rC   rd   re   r?  r   r   r   test_empty_table  s    
rU  c                  C   s   t tdt dddddgt tddg} t jj| dddgd	}|jdddgksZt|d
ddg}|  |jd
ddgkstt jj| d
ddgd	}|	|std S )Nr   r   r   r   r_   r`   rt   r   ra   ehbeeZsea)
r   rA   rj   rd   re   r   r   Zrename_columnsr!   rC   rP  r   r   r   test_table_rename_columns  s    rX  c                  C   s  t t dt  t dt  g} t t d| g}t jddg| d}t jddg|d}t jd	d
gt  d}t jj|||gdddgd}|	 }|
  t jjt jddgt  dt jddgt  dt jddg| d|gddddgd}||std S )Nrp   rq   nestr   r   r   )   g      )@)   g      -@FTr`   rt   r   ra   r   r	   rx   r   rZ  r\  za.xza.yzb.nest)r   r   r   r   r   rA   rH   rd   re   r   r!   rC   r   )Zty1Zty2r`   rt   r   rf   rD  rF   r   r   r   test_table_flatten  s&    
r^  c                  C   s   t jt dgt dggddgd} t jt dgt dggddgd}t j| |g}| }|  ||szt|j	D ]}|j
dkstqd S )Nr   r`   r   r   ra   r   rt   )r   r   rA   rd   r  rB   r!   rC   r   r   r-  )r  r  rf   combinedr   r   r   r   test_table_combine_chunks  s    
r`  c                  C   s`  t jt ddd dg t ddddg t ddd d ggdd	d
gd} t jt ddd dg t dddd g t dd d dggdd	d
gd}t j| |g}|ddi}|ddj	
t ddgst|ddj	
t ddgst|ddj	
t dddgs.t|ddj	
t ddgsVt|t  }t dddg}t dddg}|ddj	
|st|ddj	
|st|ddj	
|st|ddj	
|st| 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ksFt|jjddiks\td S )Nr   r   {   i  i  TFr`   rt   r   ra   r   s   key1s   value1r   r   r   )r   r   rA   r   rd   r  r  r   r   r   rC   r   r   Zdefault_memory_poolr   r   r   )r  r  rf   Zexpected_dict_0Zexpected_dict_1r   r   r   test_table_unify_dictionaries  sR    


rb  c                  C   s   t tddddddgg} t tdddd	d
ddgg}tjjdd | D dd}tjjdd |D dd}t||g}|  t|dksttjjdd t	| |D dd}|
|std S )Nr   g      $g      g        r7   g      $@r_   rO   rP   rQ   r6   c                 S   s   g | ]}t |qS r   r   rA   r  r   r   r   r  &  s     z&test_concat_tables.<locals>.<listcomp>r   ra   c                 S   s   g | ]}t |qS r   rc  r  r   r   r   r  (  s     c                 S   s   g | ]\}}t || qS r   rc  )r   rp   rq   r   r   r   r  /  s   )r&   rj   r   rd   re   concat_tablesr!   r)   r   ri   rC   )r2   r3   t1rD  r~   rF   r   r   r   test_concat_tables  s*    
rf  c                	   C   s&   t t td g W 5 Q R X d S r   )r"   r#   AttributeErrorr   rd  r   r   r   r   test_concat_tables_none_table6  s    rh  c                  C   s   dd l } ttdt tdt g}td}| ||d}| tjgd |d}tj	j
||dd}tj	j
||dd}|j|jst|jj|jd	d
rtt||g}|jj|jd	d
st|j|jstd S )Nr   r`   rt   Zabcdefghr   r   F)r   r"  Tr   )r   r   r   r   r   r&   r   rJ   nanrd   r   rC   r   rd  )r   r   valuesr)  r*  table1table2table3r   r   r   1test_concat_tables_with_different_schema_metadata<  s    rn  c               
   C   s   t jt jddgt  dgdg} t jt jddgt  dgdg}t j| |gdd	}|t jt jddd d gt  dt jd d ddgt  dgddgstd S )
Nr   r   r   Zint64_fieldrO   rP   Zfloat_fieldTZpromote)	r   rd   re   rA   r;   r   rd  rC   r   re  rD  r~   r   r   r   !test_concat_tables_with_promotionS  s      
rq  c               	   C   sv   t jt jddgt  dgdg} t jt jddgt  dgdg}tt j t j	| |gdd W 5 Q R X d S )Nr   r   r   r   Tro  )
r   rd   re   rA   r;   r   r"   r#   r  rd  )re  rD  r   r   r   'test_concat_tables_with_promotion_errora  s      rr  c               
   C   s   t tdt dddddgt dddd	d
gt dddddgg} t jj| tdd}|d |d spt|d |d st|d |d st|d |d stt	t
 |d  W 5 Q R X t	t
 |d  W 5 Q R X d S )Nr   r   r   r   r_   rO   rP   rQ   r6   r7   abbccddeefabcdra   rX   r	   r  r   r   r  r
   )r   rA   rj   rd   re   tuplerC   r   r"   r#   r[   r<  r   r   r   test_table_negative_indexingk  s    r{  c               	   C   s   t tdt dddddgg} t jj| tdd}t t dt  t d	t 	 g}t t d
t  g}d}t
jt|d || W 5 Q R X t
jt|d || W 5 Q R X d S )Nr   r   r   r   r_   rs  ra   r1  rt   r`   zGTarget schema's field names are not matching the table's field names:.*r9   )r   rA   rj   rd   re   rz  r   r   r\   r   r"   r#   r$   cast)r2   rf   Ztarget_schema1Ztarget_schema2rS  r   r   r   &test_table_cast_to_incompatible_schema  s     r}  c               
   C   sT  t jtdt  dt jdddddgt  dt jddd	d
dgt  dt jdddddgt  dg} t jj| t	dd}t jtdt  dt jdddddgt 
 dt jdddddgt  dt jdddddgt  dg}t jj|t	dd}t t dt  t dt 
 t dt  t dt  g}||}||sPtd S )Nr   r   r   r   r   r_   rO   rP   rQ   r6   r7   rs  rt  ru  rv  rw  rx  ra   r   r   r	   r
   r`   rt   r   ru   )r   rA   rj   r;   r\   rS   r   rd   re   rz  r   r   r   r|  rC   r   r2   rf   rG  Zexpected_tableZtarget_schemaZcasted_tabler   r   r   test_table_safe_casting  s(    
r  c               
   C   s~  t jtdt  dt jdddddgt  dt jddd	d
dgt  dt jdddddgt  dg} t jj| t	dd}t jtdt  dt jdddddgt 
 dt jdddddgt  dt jdddddgt  dg}t jj|t	dd}t t dt  t dt 
 t dt  t dt  g}tjt jdd || W 5 Q R X |j|dd}||sztd S )Nr   r   r   r   r   r_   g?g@gffffff
@g@g      @rs  rt  ru  rv  rw  rx  ra   r   r   r	   r
   r`   rt   r   ru   	truncatedr9   F)safe)r   rA   rj   r;   r\   rS   r   rd   re   rz  r   r   r   r"   r#   r  r|  rC   r   r~  r   r   r   test_table_unsafe_casting  s,    r  c               	   C   sr   t jddgt jd} t }tj| |dtj| dd  |dg}ttjj tjj	|ddgd W 5 Q R X d S )Nr   r   r   r   r  r  ra   )
rJ   rA   uint8r   r"   r#   r   r  rd   re   )rA   u8r}   r   r   r   test_invalid_table_construct  s
    $r  zdata, klassr9  r   r   r   r   r7   c              	   C   s  || d || d g} t dt  fdt  fg}t jj| |d}|jdksTt|jdksbt|j|ksptt dt  fg}t	
t t jj| |d W 5 Q R X t dt  fdt  fg}t jj| |d}t j|djst|jdkst|jdkst|j|kstt dt  fdt dfg}t	
ttf t jj| |d W 5 Q R X t	
t t jj| |ddgd	 W 5 Q R X t	
t t jj| |d
did W 5 Q R X d S )Nr   r   r:  r   r   r   r	   s)r   rb   r   r   r   r   )r   r   r   r   rd   re   r   r   r   r"   r#   r$   types
is_float32r   r   r   r   rR   from_pydictr2   klassr   rf   r   r   r   test_from_arrays_schema  s,      r  clsc              
   C   sT  |  i }|jdkst|jdks&t|jtg ks:t| i ksJttdt fdt fg}t	ddddgfdddd gfg}|  |}|jd	kst|jd
kst|j|kst| |kstddi}|
|}| j ||d}|j|kst|jj|kst| |kst| j ||d}|j|ks<t|jj|ksNt| |ks`ttt | j |||d W 5 Q R X tt. | j dddd	gitdt fgd W 5 Q R X tjtddJ | j dd	d
gd
ddgdtdt fdt fdt fgd W 5 Q R X tt | j ddd	d
gii d W 5 Q R X d S )Nr   r:  r   r9  r   r   r   r   r   r	   r   r   r   r   r  r   r   zdoesn't contain.* c, dr9   r
   r   r`   r   ru   )r  r   r   r   r   r   r   r   rS   r   with_metadatar   r"   r#   r$   rR   r   r   r;   r\   r   )r  rf   r   r2   r   r   r   r   test_table_from_pydict  sR    






r  c              	   C   sv  t d|| d fd|| d fg} tdt fdt fg}tj| }|jdks\t|j	dksjt|j|ksxttjj| |d}|jdkst|j	dkst|j|ksttdt fdt
 fg}tjj| |d}tj|djst|jdkst|j	dkst|j|ks(ttdt fdtdfg}tttf tjj| |d W 5 Q R X d S )	Nr:  r   r   r   r   r	   r   r  )r   r   r   r   rS   rd   r  r   r   r   r   r  r  r   r   r   r"   r#   r   rR   r  r   r   r   #test_table_from_pydict_arrow_arraysB  s&    $ r  c              	   C   s   t d|| d fd|| d fg} tdt fdt fdt fg}tjtdd tj	j
| |d W 5 Q R X tdt fg}tj	j
| |d}|jdkst|j|kst|jdgkstd S )Nr:  r   r   r   r   r9   r   )r   r   r   r   rS   r;   r"   r#   r   rd   r  r   r   r   r  r   r   r   test_table_from_pydict_schemad  s    $
r  c              	   C   sn  |  g }|jdkst|jdks&t|jtg ks:t| g ksJttdt fdt fg}ddddddd	d dg}|  |}|jd
kst|jdkst|j|kst| |kstddi}|	|}| j ||d}|j|kst|jj
|kst| |kst| j ||d}|j|ks8t|jj
|ksJt| |ks\ttt | j |||d W 5 Q R X tt6 | j ddiddidd
igtdt fgd W 5 Q R X tdt fdt fdt fg}| j dddd
dddddg|d}dd d dd
d d ddd d dg}|j|ksHt| |ksZttt& | j ddidd
iddigi d W 5 Q R X dddddidd	ig}dddd ddd	d dg}|  |}|jd
kst|jdkst| |ks tddiddddd ig}ddiddidd ig}|  |}|jdksHt|jdksXt| |ksjtd S )Nr   r:  r   r9  r   )r:  r   r   r   r   r   r	   r   r   r   r   r  r   r   r`   r   ru   r   r
   )r`   r   ru   )Zfrom_pylistr   r   r   r   r   r    r   rS   r  r   r"   r#   r$   rR   r   r;   r\   r   )r  rf   r   r2   r   r3   r   r   r   test_table_from_pylist|  s    







*

r  c               	   C   s  dd l } | tddddgfdddd gfg}tdt fdt fg}tjj||d	}tj	
|djstt|j |ksttdt fdtd
fg}tttf tjj||d	 W 5 Q R X tdt fdt fdt fg}tjtdd tjj||d	 W 5 Q R X tdt fg}tjj||d	}|jdksVt|j |ksjt|jdgks|td S )Nr   r:  r9  r   r   r   r   r   r   r  r   r9   r   )r   r   r   r   r   r   r   rd   r   r  r  r   r   r   Zremove_metadatar   r"   r#   r   rR   rS   r;   r   r   r   )r   r   r   rf   r   r   r   test_table_from_pandas_schema  s*     
r  c                  C   s   dd l } tddddgfddddgfg}tjdddgdd	tjdddgd
d	d}tdt fdt fg}| |}t|}tj	
|}||sttj||d}tj	j
||d}||stt|}tj	|}||sttj||d}tj	j||d}||std S )Nr   rt   r`   r   r   r   r	   r   r   r\   )rt   r`   r   )r   r   r   rA   r   r\   r   r   rf   rd   r   rC   r   r  )r   ru   Z
d_explicitr   r   rk  rl  r   r   r   test_table_factory_function  s&     


r  c               	   C   s   t t  tjddddgidgd W 5 Q R X tdt fg} tdtjdddgt di| }|	dj
t ksttjdddgddg}dg}t||}|j|ksttdt fg} t|| }|j|kstd S )Nr`   r   r   r	   ra   r   r;   )r"   r#   r$   r   rf   r   r\   rA   r;   r   r   r   r   )r   rf   r2   rb   r   r   r    test_table_factory_function_args  s    $$r  c               	   C   s   dd l } tt& tj| ddddgidgd W 5 Q R X tt( tj| ddddgiddid	 W 5 Q R X tdt fg}t| ddddgi|}|	dj
t kstd S )
Nr   r`   r   r   r	   ra   r   r   r   )r   r"   r#   r$   r   rf   r   r   r\   r   r   r   )r   r   rf   r   r   r   'test_table_factory_function_args_pandas+  s    *,r  c                	   C   sL   t jtdd td W 5 Q R X t jtdd td W 5 Q R X d S )Nz!Expected pandas DataFrame, pythonr9   zinvalid inputzExpected pandas DataFrame)r"   r#   rR   r   rf   r   r   r   r   r   $test_factory_functions_invalid_input<  s    r  c                  C   s   t jt jdt  ddidt dt  gddid} t jt jdd	d
dgddt jddddgddg| d}t|dks|t|j	dddkst|j	dddkst|j	dddkstd S )Nr   r   r   r   r   r   r   r   r   r	   r
   r   r   r_      r!  (   r\   r   zJpyarrow.Table
c0: int16
c1: int32
----
c0: [[1,2,3,4]]
c1: [[10,20,30,40]]Tr   zhpyarrow.Table
c0: int16
  -- field metadata --
  key: 'value'
c1: int32
-- schema metadata --
foo: 'bar'r   )Zpreview_colsz:pyarrow.Table
c0: int16
c1: int32
----
c0: [[1,2,3,4]]
...)
r   r   r   r   r\   rf   rA   rU   r   r   r   tabr   r   r   test_table_repr_to_stringD  s     	r  c                  C   s   t jt jdt  ddidt dt  gddid} t jt jdd	d
dgd ddt jddddgd ddg| d}t|dkstd S )Nr   r   r   r   r   r   r   r   r   r	   r
   r_   r   r   r  r!  r  r\   r   zppyarrow.Table
c0: int16
c1: int32
----
c0: [[1,2,3,4,1,...,4,1,2,3,4]]
c1: [[10,20,30,40,10,...,40,10,20,30,40]])	r   r   r   r   r\   rf   rA   rU   r   r  r   r   r   "test_table_repr_to_string_ellipsisn  s    r  c                  C   s   d} d}t |dddgf| dddgfg}t| t f|t fg}tj||d	}|d
 d
tjdddgdds|t	|d d
tjdddgddst	d S )Nu   ääähu   öööfr`   rt   r   r   r   r	   r   r   r\   r   r   )
r   r   r   r\   r   rf   r   rC   rA   r   )Zcol_aZcol_bru   r   r~   r   r   r   "test_table_function_unicode_schema  s     (r  c               
   C   s^   t t dddd dgt ddddd	ggd
dg} | t ddg| ddsZtd S r   )r   rf   rA   r   rC   rc   r   r?  r   r   r   %test_table_take_vanilla_functionality  s    r  c               
   C   s~   t t dddd dgt ddddd	ggd
dg} t t dd gt dd ggd
dg}| t dd g|sztd S )Nr   r   r	   r   r`   rt   r   ru   r   r   r   r   r   rf   rA   r   rC   r   )rf   Zresult_with_null_indexr   r   r   test_table_take_null_index  s    r  c               
   C   s~   t t dddd dgt ddddd	ggd
dg} t t dd gt ddggd
dg}| t ddg|sztd S r   r  )rf   Zresult_non_consecutiver   r   r   test_table_take_non_consecutive  s    r  c               	   C   s4  t dddd dg} t ddddd	g}t ddgdd
gddgd ddgg}t | ||gdddg}|dg}t | gdg}||st|ddg}t ||gddg}||st|dg}t | gdg}||st|ddg}t ||gddg}||st|ddi}|ddg}d|jjksHtt	j
tdd |dg W 5 Q R X t	j
tdd |dg W 5 Q R X |ddg}t ||gddg}||stt | ||gdddg}t	j
tdd |dg W 5 Q R X |dg}t |gdg}||s0td S r   )r   rA   rf   r   rC   r   r  r   r   r"   r#   r   r[   )r  r  r  rf   r~   rF   rl  r   r   r   test_table_select  s@    $r  c                  C   s&  dd } t jt dddddgt dddddgt d	d
dddgt dddddggddddgd}|ddg}| | dddgdddgdkst|dddg}| | dddgdddgd
d
d	gdkst|ddg}| | dddgdddgdkst|dddg}| | dddgd
ddgdd dgd!ks^t|dd"dg}| | dddgdddgdddgd#kst|ddgdg}| | ddddgddddgddddgd$kst|dddd"g d%fdg}| | dddgd
ddgdd dgdddgd
d
d	gdddgd&ksTtt jt dddgt d	d d ggddgd}|dgdd'tjd(d)fg}| dgdgd*kst|dgdd'tjd+d)fg}| dgd
gd*kst|dgdd'tjd,d)fg}| dgd	gd*ks<t|dgg d%fdd'tjd,d)fg}| dgdgd	gd-kst|dgg d%fg}| dgdgd.kstt ddddddgt	d/d0}t j
|jdd1}|ddg}| | ddgd/d2gdks"td S )3Nc                 S   s   d| kr t t| d | d }n| d }t|}dt| d i}| D ]@}|dkrRqDtt|| | }|D ]}||g ||  qhqD|S )Nkeys2keys)rz  ri   sortedr   
setdefaultry   )ru   r  Zsorted_keysZsorted_dentryrj  kr   r   r   sorted_by_keys  s    z+test_table_group_by.<locals>.sorted_by_keysr`   rt   r   XYZr   r   r	   r
   r   r_   r  r!  r  2   r  r  rj  	bigvaluesra   )rj  Zhash_sumr   )r  
values_sum)rj  Z
hash_count)r  r  values_count)rj  r*   )rj  max)r  r*   F   )r  
values_maxbigvalues_sum)r  r  )r  r  bigvalues_max)r  r  r  	count_all)r  r  r  r  r  r  countr'   )mode)r  r  Z	only_nullZ
only_valid)r  r  r  )r  r  r   r  rj  r$  r   )r   rf   rA   Zgroup_byZ	aggregater   r   pcZCountOptionsrj   rd   r  r&  )r  rf   rZtable_with_nullsZtable_with_chunksr   r   r   test_table_group_by  s    
















	

r  c                  C   sr   t jddddgi} |  }| j|jks.t| | ks>t| jdd}| jdks\t| jdksntd S )Nrp   r   r   r	   r$  )	r   rd   r  Z	to_readerr   r   Zread_allZread_next_batchr   )rf   readerr   r   r   test_table_to_recordbatchreader  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}| |dd}| t dddgdddgdd
d gdksvt| j|dddd}| dt ddddgdddd gdd
d d	gdkstd S )Nr   r   r   r`   rt   r   colAcol2rY   r  r2  r1  )colBcol3r  r  r  r  r  
full outer	join_typer   rf   joinrB   r   sort_byrp  r   r   r   test_table_join  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}| |d}| t dddgdddgdd
d gdkstt| j|dddd}| dt ddddgdddd gdd
d d	gdkstd S )Nr   r   r   r`   rt   r   r  rY   r  r2  r1  )r  r  r  r  r  _r)r  Zright_suffixr  rp  r   r   r   test_table_join_unique_key  s(    



r  c               
   C   s   t dddgdddgddd	gd
} t dddgdddgdddgd
}| j|ddd}| dt jddddgdddd gddd	d gddd dgddd dggdddddgdkstd S )Nr   r   r   r_   r  <   r`   rt   r   r  r  colValsrY   r  r2  r1  r  r  r  r  r  ra   )r   rf   r  rB   r  r   rp  r   r   r   test_table_join_collisions  s(    




r  c               
   C   s   t dddgdddgddd	gd
} t dddgdddgdddgd
}t | |g}|tddk }| t dddddgdddddgddd	ddgd
kstd S )Nr   r   r   r_   r  r  r`   rt   r   r  rY   r  r2  r1  r  )r   rf   rd  filterr  r   rB   r   )re  rD  rE  r~   r   r   r   test_table_filter_expression  s"    
r  c                  C   s2  t dddgdddgd} t dddgd	d
dgd	d
dgd	d
dgd	d
dgd	d
dg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d
d gdd
d gdkst| j|dddd}| dt ddddgdddd gdd
d d	gdd
d d	gdd
d d	gdd
d d	gdd
d d	gdks.td S )Nr   r   r   r`   rt   r   r  rY   r  r2  r1  )r  r  col4col5col6col7r  r  )r  r  r  r  r  r  r  r  r  r  rp  r   r   r   test_table_join_many_columns  s@    	








r  c               	   C   s   t d dgd dgd} t t jddddt jddd	dg}tt | | W 5 Q R X t d dgd	dgd} | |j|kstd S )
Nr   Tr   r`   r;   r3  rt   boolF)	r   rf   r   r   r"   r#   r$   r|  r   )rf   Z
new_schemar   r   r   test_table_cast_invalid	  s    r  c               	   C   sh  t jt dddddgt dddddggd	d
gd} | d	 dddddgdddddgdksft| dg dddddgdddddgdkstt jjt jddddgt  dt ddddggddgd}|dg}| }|d ddddgkst|d ddddgkst|dg}| }|d ddddgksJt|d ddddgksdtd S )Nr	   r   r
   r   r   rt   r`   r   rj  r  ra   r  )rj  
descendingr   #   r   r   carr   foobarr`   r  r`   	ascending)	r   rf   rA   r  r   r   rd   re   r;   )rf   r  Z
sorted_tabZsorted_tab_dictr   r   r   test_table_sort_by'	  s8    r  c               	   C   s  t jjt jddddgt  dt jddddgt  dt d	d
ddggdddgd} | ddg}| }|d ddddgkst|d ddddgkst|d d
d	ddgkst| ddg}| }|d ddddgkst|d ddddgkst|d ddd	d
gkstd S )Nr   r  r   r   r
   r   r	   r   r   r  r   r  r`   rt   r   ra   r  )rt   r  r  )rt   r  )r   r   re   rA   r;   r  r   r   )r   Z	sorted_rbZsorted_rb_dictr   r   r   test_record_batch_sortG	  s"    r  constructorc                 C   sT  | dddgdddggddgd	}t |}t jdd
gddgddggdd}t j|| t j|dd}t j|| |jdkst|g }t |}t d}t j|| |jdkstt j|dd}t j|| |jdkst|	dd}t |}t d}t j|| |jdks$tt j|dd}t j|| |jdksPtd S )Nr   r   r	   r6   r7   r8   r`   rt   ra   r
   r   r   rS   r   r\   )r	   r   r   r   )
rJ   rK   rA   r   Zassert_allcloser   r   r   emptyrc   )r  rf   r~   rF   rl  rm  r   r   r   test_numpy_asarray[	  s.     
 




r  )collectionsr   collections.abcr   rz   r+   r0   numpyrJ   r"   Zpyarrowr   Zpyarrow.computeZcomputer  r5   r@   rG   rI   rN   rT   rW   r^   rh   rm   rv   markZparametrizerH   r;   rS   r   r   r  r   r   r   r   r   r   r   r   r   Znopandasr   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/  r5  r6  r7  r;  r=  r>  r@  rF  rH  rM  rN  rQ  rR  rT  rU  rX  r^  r`  rb  rf  rh  rn  rq  rr  r{  r}  r  r  r  r&   rA   r   r  rd   r   r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  r  Zaceror  r  r  r  r  r  r  r  r  r  rf   r   r  r   r   r   r   <module>   s\  "-	
0$





(<	
3%



-
%

	 
#
9 
 

U
 

*	3
 




& 