o
    $&]i/'                  	   @  s  d dl mZ d dlZd dlmZ d dlmZ d dlmZm	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mZ d dlmZ d dlmZmZ d dlmZ er|d dl m!Z!m"Z" d dl#m$Z$m%Z%m&Z& d dlm'Z' d dl(m)Z)m*Z* G dd deeeeej+df Z,dS )    )annotationsN)reduce)chain)TYPE_CHECKINGLiteralArrowDataFrame	ArrowExpr)ArrowSelectorNamespaceArrowSeries)cast_to_comparable_string_types)EagerNamespace)combine_alias_output_namescombine_evaluate_output_names)Implementation)IteratorSequence)ChunkedArrayAny
Incomplete	ScalarAny)Version)	IntoDTypeNonNestedLiteralc                   @  s   e Zd ZejZedGddZedHddZedId	d
Z	dJddZ
dKddZdLddZdMddZdNd!d"ZdNd#d$ZdOd%d&ZdOd'd(ZdOd)d*ZdOd+d,ZdPd0d1ZdPd2d3ZdPd4d5ZedQd7d8ZdRd;d<ZdOd=d>Z	?dSdTdEdFZd?S )UArrowNamespacereturntype[ArrowDataFrame]c                 C     t S Nr   self r"   W/var/www/html/IGF-ODF-V3/venv/lib/python3.10/site-packages/narwhals/_arrow/namespace.py
_dataframe$      zArrowNamespace._dataframetype[ArrowExpr]c                 C  r   r   r	   r    r"   r"   r#   _expr(   r%   zArrowNamespace._exprtype[ArrowSeries]c                 C  r   r   r   r    r"   r"   r#   _series,   r%   zArrowNamespace._seriesversionr   Nonec                C  s
   || _ d S r   )_version)r!   r*   r"   r"   r#   __init__0   s   
zArrowNamespace.__init__seriesr   %Iterator[ChunkedArrayAny | ScalarAny]c                 G  s   dd |D S )Nc                 s  s&    | ]}|j r|jd  n|jV  qdS )r   N)
_broadcastnative.0sr"   r"   r#   	<genexpr>6   s   $ z0ArrowNamespace.extract_native.<locals>.<genexpr>r"   )r!   r.   r"   r"   r#   extract_native3   s   zArrowNamespace.extract_nativer
   c                   s     j  fdddd d  jdS )Nc                   s   t jt| jgd dgS )Nlen)namecontext)r   from_iterabler7   r1   dfr    r"   r#   <lambda>;   s   z$ArrowNamespace.len.<locals>.<lambda>c                 S     dgS )Nr7   r"   _dfr"   r"   r#   r=   >       evaluate_output_namesalias_output_namesr*   r'   r,   r    r"   r    r#   r7   8   s   
zArrowNamespace.lenvaluer   dtypeIntoDType | Nonec                   s2   dfdd j  fddd	d d jd
S )N_r   r   r   c                   s$   t jgdd} r| S |S )Nliteral)datar8   r9   )r   r:   cast)rI   arrow_series)rG   r!   rF   r"   r#   _lit_arrow_seriesD   s   
z-ArrowNamespace.lit.<locals>._lit_arrow_seriesc                   s
    | gS r   r"   r;   )rN   r"   r#   r=   M   s   
 z$ArrowNamespace.lit.<locals>.<lambda>c                 S  r>   )NrJ   r"   r?   r"   r"   r#   r=   N   rA   rB   )rI   r   r   r   rE   )r!   rF   rG   r"   )rN   rG   r!   rF   r#   litC   s   
zArrowNamespace.litexprsignore_nullsboolc                  ,   d fdd}| j j|t  t  | dS )	Nr<   r   r   list[ArrowSeries]c                   8   t  fddD }rdd |D }ttj|gS )Nc                 3      | ]}| V  qd S r   r"   r3   er;   r"   r#   r5   U       z>ArrowNamespace.all_horizontal.<locals>.func.<locals>.<genexpr>c                 s      | ]
}| d ddV  qdS )TN	fill_nullr2   r"   r"   r#   r5   W       )r   r:   r   operatorand_r<   r.   rP   rQ   r;   r#   funcT      z+ArrowNamespace.all_horizontal.<locals>.funcrb   rC   rD   r9   r<   r   r   rT   r'   _from_callabler   r   r!   rQ   rP   rb   r"   ra   r#   all_horizontalS      zArrowNamespace.all_horizontalc                  rS   )	Nr<   r   r   rT   c                   rU   )Nc                 3  rV   r   r"   rW   r;   r"   r#   r5   c   rY   z>ArrowNamespace.any_horizontal.<locals>.func.<locals>.<genexpr>c                 s  rZ   )FNr[   r2   r"   r"   r#   r5   e   r]   )r   r:   r   r^   or_r`   ra   r;   r#   rb   b   rc   z+ArrowNamespace.any_horizontal.<locals>.funcrd   re   rf   rh   r"   ra   r#   any_horizontala   rj   zArrowNamespace.any_horizontalc                   s*   d fdd}| j j|t  t  | dS )	Nr<   r   r   rT   c                   s4   t  fddD }dd |D }ttj|gS )Nc                 3  rV   r   r"   r3   exprr;   r"   r#   r5   q   rY   z>ArrowNamespace.sum_horizontal.<locals>.func.<locals>.<genexpr>c                 s  s     | ]}|j d dddV  qdS r   N)strategylimitr[   r2   r"   r"   r#   r5   r   s    )r   r:   r   r^   add)r<   itr.   rP   r;   r#   rb   p   s   z+ArrowNamespace.sum_horizontal.<locals>.funcrd   re   rf   r!   rP   rb   r"   rt   r#   sum_horizontalo   s   zArrowNamespace.sum_horizontalc                   s8   | j j d fdd}| jj|t  t  | dS )	Nr<   r   r   rT   c                   sV   t t fddD }dd |D }fdd|D }ttj|ttj| gS )Nc                 3  rV   r   r"   rm   r;   r"   r#   r5      rY   z?ArrowNamespace.mean_horizontal.<locals>.func.<locals>.<genexpr>c                 S  s   g | ]
}|j d dddqS ro   r[   r2   r"   r"   r#   
<listcomp>   s    z@ArrowNamespace.mean_horizontal.<locals>.func.<locals>.<listcomp>c                   s   g | ]}d |    qS )   )is_nullrL   r2   )int_64r"   r#   rw      s    )tupler   r:   r   r^   rr   )r<   expr_resultsr.   non_narP   rz   r;   r#   rb      s   z,ArrowNamespace.mean_horizontal.<locals>.funcrd   re   )r,   dtypesInt64r'   rg   r   r   ru   r"   r~   r#   mean_horizontal|   s   zArrowNamespace.mean_horizontalc                   ,   d fdd}j j|t  t  dS )	Nr<   r   r   rT   c                   N   t t fddD ^}}ttjdd |D |j}t||jj	dgS )Nc                 3  rV   r   r"   rm   r;   r"   r#   r5      rY   z>ArrowNamespace.min_horizontal.<locals>.func.<locals>.<genexpr>c                 S     g | ]}|j qS r"   r1   r2   r"   r"   r#   rw          z?ArrowNamespace.min_horizontal.<locals>.func.<locals>.<listcomp>r8   r*   )
r{   r   r:   r   pcmin_element_wiser1   r   r8   r,   r<   init_seriesr.   native_seriesrP   r!   r;   r#   rb          z+ArrowNamespace.min_horizontal.<locals>.funcrd   re   rf   ru   r"   r   r#   min_horizontal      	zArrowNamespace.min_horizontalc                   r   )	Nr<   r   r   rT   c                   r   )Nc                 3  rV   r   r"   rm   r;   r"   r#   r5      rY   z>ArrowNamespace.max_horizontal.<locals>.func.<locals>.<genexpr>c                 S  r   r"   r   r2   r"   r"   r#   rw      r   z?ArrowNamespace.max_horizontal.<locals>.func.<locals>.<listcomp>r   )
r{   r   r:   r   r   max_element_wiser1   r   r8   r,   r   r   r;   r#   rb      r   z+ArrowNamespace.max_horizontal.<locals>.funcrd   re   rf   ru   r"   r   r#   max_horizontal   r   zArrowNamespace.max_horizontaldfsSequence[pa.Table]pa.Tablec                C  s&   | j dkrtj|ddS tj|ddS )N)   default)promote_optionsT)promote)_backend_versionpaconcat_tables)r!   r   r"   r"   r#   _concat_diagonal   s   
zArrowNamespace._concat_diagonalc                C  s@   t tdd |D }ttdd |D }tjj||dS )Nc                 s      | ]}|j V  qd S r   )column_namesr3   r<   r"   r"   r#   r5          z4ArrowNamespace._concat_horizontal.<locals>.<genexpr>c                 s  s    | ]}|  V  qd S r   )itercolumnsr   r"   r"   r#   r5      rY   )names)listr   r:   r{   r   Tablefrom_arrays)r!   r   r   arraysr"   r"   r#   _concat_horizontal   s   z!ArrowNamespace._concat_horizontalc                C  s`   |d j }t|dd  ddD ]\}}|j }||kr*d| d| d| d}t|qt|S )Nr   rx   )startz>unable to vstack, column names don't match:
   - dataframe 0: z
   - dataframe z: 
)r   	enumerate	TypeErrorr   r   )r!   r   cols_0ir<   cols_currentmsgr"   r"   r#   _concat_vertical   s   

zArrowNamespace._concat_verticalr   c                 C  s
   t | S r   )r   from_namespacer    r"   r"   r#   	selectors   s   
zArrowNamespace.selectors	separatorstrc                  s0   d fdd}j j|t  t  dS )	Nr<   r   r   rT   c                   s|   t t fddD }|d j}rdnd}tj| di\}}tj}j|g ||R d|i|j	d}|gS )	Nc                 3  rV   r   r"   rm   r;   r"   r#   r5      rY   z:ArrowNamespace.concat_str.<locals>.func.<locals>.<genexpr>r   skip	emit_nullr   null_handlingr   )
r   r   r:   r8   r   r6   r   binary_join_element_wiser)   r,   )r<   r.   r8   r   rs   separator_scalar
concat_str	compliantrP   rQ   r!   r   r;   r#   rb      s    


z'ArrowNamespace.concat_str.<locals>.funcrd   re   rf   )r!   r   rQ   rP   rb   r"   r   r#   r      s   zArrowNamespace.concat_strc                   r   )	Nr<   r   r   rT   c                   sV   j j}|t fddD  ^}}ttj|jgdd |D R  |jj	dgS )Nc                 3  rV   r   r"   rm   r;   r"   r#   r5      rY   z8ArrowNamespace.coalesce.<locals>.func.<locals>.<genexpr>c                 s  r   r   r   r2   r"   r"   r#   r5      r   r   )
r)   _align_full_broadcastr   r:   r   r   coalescer1   r8   r,   )r<   alignr   r.   r   r;   r#   rb      s    z%ArrowNamespace.coalesce.<locals>.funcrd   re   rf   ru   r"   r   r#   r      s   zArrowNamespace.coalesceNwhenr   then	otherwiseChunkedArrayAny | Nonec                 C  s,   |d u rt t||jn|}t|||S r   )r   nullsr7   typer   if_else)r!   r   r   r   r"   r"   r#   _if_then_else   s   zArrowNamespace._if_then_else)r   r   )r   r&   )r   r(   )r*   r   r   r+   )r.   r   r   r/   )r   r
   )rF   r   rG   rH   r   r
   )rP   r
   rQ   rR   r   r
   )rP   r
   r   r
   )r   r   r   r   )r   r   )rP   r
   r   r   rQ   rR   r   r
   r   )r   r   r   r   r   r   r   r   )__name__
__module____qualname__r   PYARROW_implementationpropertyr$   r'   r)   r-   r6   r7   rO   ri   rl   rv   r   r   r   r   r   r   r   r   r   r   r"   r"   r"   r#   r      s6    














r   r   )-
__future__r   r^   	functoolsr   	itertoolsr   typingr   r   pyarrowr   pyarrow.computecomputer   narwhals._arrow.dataframer   narwhals._arrow.exprr
   narwhals._arrow.selectorsr   narwhals._arrow.seriesr   narwhals._arrow.utilsr   narwhals._compliantr   narwhals._expression_parsingr   r   narwhals._utilsr   collections.abcr   r   narwhals._arrow.typingr   r   r   r   narwhals.typingr   r   r   r   r"   r"   r"   r#   <module>   s.    
