o
    ò©Zh=  ã                   @   sâ   d dl Z d dlZd dlZejjdd„ ƒZejjdd„ ƒZejjdd„ ƒZejjdd	„ ƒZ	ejj
d
ddd„ ƒZejjdd„ ƒZejjdd„ ƒZejjdd„ ƒZejjdd„ ƒZejjdd„ ƒZejjdd„ ƒZejjdd„ ƒZdS )é    Nc                  C   st  dd l m}  |  ¡ }t dt ¡ ¡}t dt ¡ ¡}t ||g¡}t dt ¡ ¡}| |¡}| |¡}| ¡ |j	ks<J ‚| 
d||gt ¡ ¡}| |||t ¡ ¡}	| |	|¡}
|
 ¡ j	t ¡ ksbJ ‚| jdd}|  ||
gt ¡ d|¡}|j d	¡d
ks~J ‚tjg d¢t ¡ d}tjg d¢t ¡ d}tjg d¢t ¡ d}tjj||gddgd}| |¡\}| |¡s¸J ‚d S )Nr   ÚaÚbÚresÚgreater_thanT©Zdump_irÚNONEú@expr_éÿÿÿÿ)é
   é   éìÿÿÿé   ©Útype)r   é   r   é   )r
   r   r   r   ©Únames)Úpyarrow.gandivaÚgandivaÚTreeExprBuilderÚpaÚfieldÚint32ÚschemaÚ
make_fieldÚreturn_typer   Úmake_functionÚbool_Úmake_ifÚmake_expressionÚresultÚConfigurationÚmake_projectorÚdefault_memory_poolÚllvm_irÚfindÚarrayZRecordBatchÚfrom_arraysÚevaluateÚequals)r   ÚbuilderZfield_aZfield_br   Úfield_resultÚnode_aÚnode_bÚ	conditionÚif_nodeÚexprÚconfigÚ	projectorr   r   ÚeZinput_batchÚr© r6   úQ/var/www/html/lang_env/lib/python3.10/site-packages/pyarrow/tests/test_gandiva.pyÚtest_tree_exp_builder   s4   

ÿÿr8   c                  C   sÞ   dd l m}  tj t ddg¡t ddg¡gddg¡}|  ¡ }| |j 	d¡¡}| |j 	d¡¡}| 
d||gt ¡ ¡}t 	d	t ¡ ¡}| ||¡}|  |j|gt ¡ ¡}| | ¡ d ¡\}	t dd
g¡}
|	 |
¡smJ ‚d S )Nr   ç      ð?ç       @ç      @g      @r   r   ÚaddÚcg      @)r   r   r   ÚTabler(   r'   r   r   r   r   r   Úfloat64r    r#   r$   r)   Ú
to_batchesr*   )r   Útabler+   r-   r.   Úsumr,   r1   r3   r5   r4   r6   r6   r7   Ú
test_tableB   s     ÿÿrC   c            
      C   s  dd l m}  tj t dd„ tdƒD ƒ¡gdg¡}|  ¡ }| |j	 
d¡¡}| dt ¡ ¡}| d||gt ¡ ¡}| |¡}| ¡ jt ¡ ksIJ ‚| jdd	}|  |j	||¡}|j d
¡dksaJ ‚| | ¡ d t ¡ ¡}	|	 ¡  tjtdƒt ¡ d¡sJ ‚d S )Nr   c                 S   s   g | ]}d | ‘qS )r9   r6   )Ú.0Úir6   r6   r7   Ú
<listcomp>a   s    ztest_filter.<locals>.<listcomp>i'  r   g     @@Ú	less_thanTr   r   r	   iè  r   )r   r   r   r>   r(   r'   Úranger   r   r   r   Úmake_literalr?   r   r   Úmake_conditionr!   r   r"   Úmake_filterr%   r&   r)   r@   r$   Úto_arrayr*   Úuint32)
r   rA   r+   r-   ZthousandÚcondr/   r2   Úfilterr!   r6   r6   r7   Útest_filter]   s   ÿ
(rP   c            	      C   sö  dd l m}  t g d¢¡}tj |gdg¡}|  ¡ }| |j 	d¡¡}| 
|ddgt ¡ ¡}| |¡}|  |j|¡}| | ¡ d t ¡ ¡}| ¡  tjddgt ¡ d¡sXJ ‚t g d	¢¡}tj | t ¡ ¡gdg¡}| |j 	d¡¡}| 
|dd
gt ¡ ¡}| |¡}|  |j|¡}| | ¡ d t ¡ ¡}| ¡  tjg d¢t ¡ d¡s«J ‚t g d	¢¡}tj |gdg¡}| |j 	d¡¡}| 
|dd
gt ¡ ¡}| |¡}|  |j|¡}| | ¡ d t ¡ ¡}| ¡  tjg d¢t ¡ d¡sùJ ‚d S )Nr   )ÚgaÚanÚndZdiZivÚvar   rR   rS   é   é   r   )
é   rU   é   rU   r   é	   rV   é   r   rX   r   )rU   rW   rX   é   )r   r   r   r'   r>   r(   r   r   r   r   Úmake_in_expressionÚstringrJ   rK   r)   r@   r$   rL   r*   rM   Úcastr   Úint64)	r   ÚarrrA   r+   r-   rN   r/   rO   r!   r6   r6   r7   Útest_in_expru   s4   
$
$
(ra   z>Gandiva C++ did not have *real* binary, time and date support.)Úreasonc                  C   sÒ  dd l m}  t g d¢¡}tj |gdg¡}|  ¡ }| |j 	d¡¡}| 
|ddgt ¡ ¡}| |¡}|  |j|¡}| | ¡ d t ¡ ¡}| ¡  tjddgt ¡ d¡sXJ ‚tj d	¡}	tj d
¡}
tj d¡}t |	|
|g¡}tj |gdg¡}|  ¡ }| |j 	d¡¡}| 
||
gt d¡¡}| |¡}|  |j|¡}| | ¡ d t ¡ ¡}t| ¡ ƒdgks¶J ‚|	 ¡ }|
 ¡ }| ¡ }t |||g¡}tj |gdg¡}|  ¡ }| |j 	d¡¡}| 
||gt d¡¡}| |¡}|  |j|¡}| | ¡ d t ¡ ¡}t| ¡ ƒdgksJ ‚|	 ¡ }|
 ¡ }| ¡ }t |||g¡}tj |gdg¡}|  ¡ }| |j 	d¡¡}| 
||gt ¡ ¡}| |¡}|  |j|¡}| | ¡ d t ¡ ¡}t| ¡ ƒdgksgJ ‚d S )Nr   )s   gaó   anó   nds   dis   ivs   var   rc   rd   rU   rV   r   gÕÌç¹,ûÖAgÕÌç¯,ûÖAgÕÌçØ+ûÖAÚms)r   r   r   r'   r>   r(   r   r   r   r   r\   ÚbinaryrJ   rK   r)   r@   r$   rL   r*   rM   ÚdatetimeÚutcfromtimestampÚ	timestampÚlistÚtimeZtime64ÚdateZdate32)r   r`   rA   r+   r-   rN   r/   rO   r!   Z
datetime_1Z
datetime_2Z
datetime_3Ztime_1Ztime_2Ztime_3Zdate_1Zdate_2Zdate_3r6   r6   r7   Útest_in_expr_todoš   s\   
$


rm   c                  C   s4  dd l m}  tj t g d¢¡t g d¢¡gddg¡}|  ¡ }| |j 	d¡¡}| |j 	d¡¡}| 
dt ¡ ¡}| 
dt ¡ ¡}| d||gt ¡ ¡}| d	||gt ¡ ¡}| d||gt ¡ ¡}	| | ||g¡|	g¡}
| |
¡}|  |j|¡}| | ¡ d t ¡ ¡}| ¡  tjg d
¢t ¡ d¡s˜J ‚d S )Nr   )r9   g      ?@g      G@r;   g     €L@g      F@g      6@)g      @g     €F@g      B@g     @R@g     ÀT@g      7@g      S@r   r   g      I@g      &@rG   r   )r   rV   r   r   )r   r   r   r>   r(   r'   r   r   r   r   rI   r?   r   r   Úmake_orÚmake_andrJ   rK   r)   r@   r$   rL   r*   rM   )r   rA   r+   r-   r.   ZfiftyZelevenZcond_1Zcond_2Zcond_3rN   r/   rO   r!   r6   r6   r7   Útest_booleanä   s,   þýÿ
(rp   c                  C   sö  dd l m}  |  ¡ }| dt ¡ ¡ | dt ¡ ¡ | dt ¡ ¡ | dt ¡ ¡ | dt 	¡ ¡ | dt 
¡ ¡ | dt ¡ ¡ | dt ¡ ¡ | d	t ¡ ¡ | d
t ¡ ¡ | dt ¡ ¡ | dt ¡ ¡ | dt ¡ ¡ | dd¡ | dd¡ | dd¡ | dd¡ | dd¡ | dd¡ | dd¡ | dd¡ | d	d¡ | d
d¡ | dd¡ | dd¡ | dd¡ t t¡ | dt ¡ ¡ W d   ƒ n1 sØw   Y  t t¡ | dd ¡ W d   ƒ d S 1 sôw   Y  d S )Nr   TrU   rV   rW   rX   r   rZ   é   g       @g      "@Zhellos   worldÚboolÚuint8Úuint16rM   Úuint64Úint8Úint16r   r_   Úfloat32r?   r]   rf   )r   r   r   rI   r   r   rs   rt   rM   ru   rv   rw   r   r_   rx   r?   r]   rf   ÚpytestÚraisesÚ	TypeError)r   r+   r6   r6   r7   Útest_literalsÿ   sD   ÿ"ÿr|   c                  C   sè   dd l m}  g d¢}tj|t ¡ d}tjj|gdgd}|  ¡ }| |j	 
d¡¡}| dt ¡ ¡}| d||gt ¡ ¡}t 
dt ¡ ¡}| ||¡}	|  |j	|	gt ¡ ¡}
|
 | ¡ d ¡\}tjg d	¢t ¡ d}| |¡srJ ‚d S )
Nr   )ZparkÚsparklezbright spark and fireZsparkr   r   r   z%spark%Úliker   )FTTT)r   r   r   r'   r]   r>   r(   r   r   r   r   rI   r   r   r    r#   r$   r)   r@   r*   )r   ÚelementsÚdatarA   r+   r-   Úregexr~   r,   r1   r3   r5   r   r6   r6   r7   Ú
test_regex'  s    ÿr‚   c                  C   s\   dd l m}  |  ¡ }t|d  ¡ ƒtju sJ ‚t|d  ¡ ƒtu s#J ‚t	|d dƒs,J ‚d S )Nr   Úname)
r   r   Z"get_registered_function_signaturesr   r   r   ZDataTypeZparam_typesrj   Úhasattr)r   Z
signaturesr6   r6   r7   Ú'test_get_registered_function_signatures?  s
   r…   c                  C   s~  dd l m}  t ¡ }t g d¢t ¡ ¡}t g d¢t ¡ ¡}t g d¢t ¡ ¡}tj |||gg d¢¡}t dt ¡ ¡}|  	¡ }| 
|j d¡¡}| 
|j d¡¡}	| 
|j d	¡¡}
| d
||	gt ¡ ¡}| |¡}| d|	|
gt ¡ ¡}| ||	|
t ¡ ¡}| ||¡}|  |j|¡}|  |j|g|d¡}| | ¡ d |¡}| | ¡ d |¡\}t g d¢t ¡ ¡}| |¡s½J ‚d S )Nr   )r
   r   r   r   é   é   )r   r   r   r   r   rW   )rU   é   é   é   éëÿÿÿN)r   r   r=   r   r   r   r=   r   rG   ZUINT32)rU   r‹   N)r   r   r   r$   r'   r   r>   r(   r   r   r   r   r   r   rJ   r   r    rK   r#   r)   r@   r*   )r   ZmpoolZarray0Zarray1Zarray2rA   r,   r+   r-   r.   Znode_cZgreater_than_functionZfilter_conditionZproject_conditionr0   r1   rO   r3   Zselection_vectorr5   Úexpr6   r6   r7   Útest_filter_projectI  sD   ÿÿÿ
ÿÿÿr   c                  C   s  dd l m}  |  ¡ }t| dt ¡ ¡ƒ d¡sJ ‚t| dt ¡ ¡ƒdks'J ‚t| 	t 
dt ¡ ¡¡ƒdks8J ‚t| 	t 
dt ¡ ¡¡ƒd	ksIJ ‚| 	t 
d
t ¡ ¡¡}| d|gt ¡ ¡}t|ƒdksfJ ‚| 	t 
dt ¡ ¡¡}| ||g¡}t|ƒdks€J ‚d S )Nr   r:   z(const double) 2 raw(rV   z(const int64) 2Úxz
(double) xÚyz
(string) yÚzÚnotzbool not((bool) z)zbool not((bool) z) && (bool) y)r   r   r   ÚstrrI   r   r?   Ú
startswithr_   r   r   r]   r   r   ro   )r   r+   Zfield_zZ	func_nodeZfield_yZand_noder6   r6   r7   Útest_to_stringx  s   ÿ""r”   c                  C   s€  dd l m}  |  ¡ }t dt ¡ ¡}t |g¡}| dt ¡ ¡}t	 
t¡ | d ¡ W d   ƒ n1 s5w   Y  t	 
t¡ | |d d d ¡ W d   ƒ n1 sRw   Y  t	 
t¡ | |d g¡ W d   ƒ n1 snw   Y  t	 
t¡ | d |g¡ W d   ƒ n1 sŠw   Y  t	 
t¡ | d g d¢t ¡ ¡ W d   ƒ n1 sªw   Y  t	 
t¡ | d |¡ W d   ƒ n1 sÅw   Y  t	 
t¡ | d ¡ W d   ƒ n1 sßw   Y  t	 
t¡ | d|d gt ¡ ¡ W d   ƒ n1 sÿw   Y  t	 
t¡ |  |d g¡ W d   ƒ n	1 sw   Y  t	 
t¡ |  |d ¡ W d   ƒ d S 1 s9w   Y  d S )Nr   rŽ   T)rU   rV   rW   rG   )r   r   r   r   r   r   r   rI   r   ry   rz   r{   r   r   ro   rn   r\   r    rJ   r   r#   rK   )r   r+   Zfield_xr   Zliteral_truer6   r6   r7   Útest_rejects_noneŒ  sF   ÿÿÿÿÿÿÿÿÿ$ÿr•   )rg   ry   Zpyarrowr   Úmarkr   r8   rC   rP   ra   Úskiprm   rp   r|   r‚   r…   r   r”   r•   r6   r6   r6   r7   Ú<module>   s6   
)


$
H

'

	
.
