a
    ~bm=                     @   s   d dl Z d dlZzd dlmZ W n ey>   d dlmZ Y n0 ddlmZ ddlm	Z	 ddl
mZmZmZ ddlmZ dd	lmZ dd
lmZ G dd deZee dS )    N)Iterable   )util)Element)	NdMapping
item_checksorted_context   )	Interface)pandas)cachedc                   @   s  e Zd ZdZdZdZdgZedd Zedd Z	ed	d
 Z
edd Zedd Zedd Zedd Zeedd Zeedd Zeedd ZeedBddZedCddZeed d! Zeed"d# ZeZeg dfd$d%Zed&d' ZejjZejjZed(d) Zed*d+ Zed,d- Z ed.d/ Z!ed0d1 Z"ed2d3 Z#eed4d5 Z$edDd6d7Z%ed8d9 Z&eg fd:d;Z'ed<d= Z(eee)j*fd>d?Z+eed@dA Z,dS )EIbisInterface ibisd   zibis.backends.omniscidb.clientc                 C   s   dd l }t|jjdS )Nr   ZRowID)ibis.expr.operationshasattrexpr
operations)clsr   r   r   7lib/python3.9/site-packages/holoviews/core/data/ibis.py	has_rowid   s    zIbisInterface.has_rowidc                 C   sR   z&ddl m}m} |t||\}W n ty@   | }Y n0 t|j| jv S )Nr   )find_backendsvalidate_backends)	Zibis.clientr   r   list	ExceptionZ_find_backendtype
__module__zero_indexed_backend_modules)r   datar   r   Zbackendr   r   r   is_rowid_zero_indexed#   s    z#IbisInterface.is_rowid_zero_indexedc                 C   s
   dt jv S )Nr   )sysmodules)r   r   r   r   loaded,   s    zIbisInterface.loadedc                 C   s"   |   sdS ddlm} t||S )NFr   )Expr)r#   Zibis.expr.typesr$   
isinstance)r   objr$   r   r   r   applies0   s    zIbisInterface.appliesc           
         s@  |j  }|d }|d }t|jd trDt|jd t|jg}nd }t|jd trb|jd nd }	 rd u r fdd|jD nr d u rfdd|jD d |  n| d u rt	|jd |  d u r fdd|j||	r||	 nd  D n0 g kr.d u r.t	|jd |	r$|	nd  |t
 di fS )	Nkdimsvdimsr	   c                    s   g | ]}| vr|qS r   r   .0ckeysr   r   
<listcomp>C       z&IbisInterface.init.<locals>.<listcomp>c                    s   g | ]}| vr|qS r   r   r*   )valuesr   r   r/   E   r0   c                    s   g | ]}| vr|qS r   r   )r+   keyr-   r   r   r/   I   s   )r(   r)   )ZparamZobjectsr%   Zboundsintminlendefaultcolumnsr   dict)
r   Zeltyper   r.   r1   paramsindexr7   ndimZnvdimr   )r.   r1   r   init7   s(    

zIbisInterface.initc                 C   s   | |j S N)Zcloner   executer   datasetr   r   r   computeR   s    zIbisInterface.computec                 C   s
   |  |S r=   )rA   r?   r   r   r   persistV   s    zIbisInterface.persistc                 C   s   |j g    S r=   )r   countr>   )selfr@   r   r   r   lengthZ   s    zIbisInterface.lengthc                 C   s   t t|jg  d S )Nr	   )boolr5   r   headr>   r?   r   r   r   nonzero`   s    zIbisInterface.nonzeroc                 C   st   |j |dd}| ||jdv r$dS |jd ur:t||S |j|j }t|j	|
 | g jdd d f S )NTstrictZSUO)NNr   )get_dimensiondtypekindZnodatar
   ranger   nametuple	aggregater4   maxr>   r1   )r   r@   	dimensioncolumnr   r   r   rN   f   s    
(zIbisInterface.rangeTFc                 C   s<   |j |dd}|j|j }|s&| }|s.|s2|S | jS )NTrI   )rK   r   rO   distinctr>   r1   )r   r@   rS   ZexpandedZflatrA   Z
keep_indexr   r   r   r   r1   s   s
    zIbisInterface.valuesNc           
         s   t    fdd D  | d}t t d }| d }t	|d |d D ] \}}	t 
|rvqb|	|t|< qb|d urtd|r||   }| fS )Nc                    s(   g | ] } j jd v rt|nt|qS )iu)rL   rM   r3   float)r+   vbinsr   r   r/      r0   z+IbisInterface.histogram.<locals>.<listcomp>bucketr	   rC   z@Weighted histograms currently not implemented for IbisInterface.)numpyZasarrayr[   rO   Zzerosr5   Zvalue_countssort_byr>   zipZisnanr3   NotImplementedErrorrC   )
r   r   rZ   ZdensityZweightsZbinnedZhistZ	hist_binsbrX   r   rY   r   	histogram   s    

zIbisInterface.histogramc                 C   s   |  |t|jjfS r=   )rE   r5   r   r7   r?   r   r   r   shape   s    zIbisInterface.shapec                 C   s"   | |}|jd j|j S Nr   )rK   r   rG   r>   ZdtypesrO   )r   r@   rS   r   r   r   rL      s    
zIbisInterface.dtypec                    s    j  fdd|D S )Nc                    s   g | ]}  |j fqS r   rK   rO   r+   xr@   reverser   r   r/      r0   z&IbisInterface.sort.<locals>.<listcomp>)r   r]   )r   r@   Zbyrh   r   rg   r   sort   s    zIbisInterface.sortc                    s$    j jf i  fdd| D S )Nc                    s   i | ]\}}|j  j| qS r   )rO   r   r+   krX   r@   r   r   
<dictcomp>   r0   z'IbisInterface.redim.<locals>.<dictcomp>)r   mutateitemsr   r@   
dimensionsr   rl   r   redim   s    
zIbisInterface.redimc                 C   s^   dd l }|  std|j d|jv r,|S | |rF|j| dS |j| d dS d S )Nr   z7iloc expressions are not supported for ibis version %s.hv_row_id__)rs   r	   )r   r   
ValueError__version__r7   r    rn   Zrowid)r   r   r   r   r   r   _index_ibis_table   s    

zIbisInterface._index_ibis_tablec                    sL  |\}}t ttj|}t|tr:dd   | D }n,t|rT |j	g}n fdd|D }| 
 j| }|r||j|k| d jd S t|trtdd |j|j|jfD r@g }|jr||j|jkg7 }|jr||j|jk g7 }||dgS n.t|ts$|g}||j|gdgS |dgS )	Nc                 S   s   g | ]
}|j qS r   rO   re   r   r   r   r/      r0   z&IbisInterface.iloc.<locals>.<listcomp>c                    s   g | ]}  |jqS r   rd   r+   drl   r   r   r/      r0   r	   r   r   c                 s   s   | ]}|d uV  qd S r=   r   re   r   r   r   	<genexpr>   r0   z%IbisInterface.iloc.<locals>.<genexpr>rs   )allmapr   isscalarr%   slicerq   r\   rK   rO   rv   r   filterrs   rG   r>   ilocanystartstopstepdropr   isin)r   r@   r:   Zrowsr7   Zscalarr   
predicatesr   rl   r   r      s6    

 zIbisInterface.ilocc                 C   s4   t |jdks"|g    dkr&|S | jd S )z~
        Given a dataset object and data in the appropriate format for
        the interface, return a simple scalar.
        r	   rz   )r5   r7   rC   r>   Ziat)r   r@   r   r   r   r   unpack_scalar   s    "zIbisInterface.unpack_scalarc           
   	      s$   fdd|D fdd j D }i dkrNttrNtt |d|  jd< dd D } j	|
  } fdd| D }	t|trtd	R td	* ||	dW  d    W  d    S 1 s0    Y  W d    n1 s0    Y  n||	S d S )
Nc                    s   g | ]} j |d dqS )TrI   )rK   rx   rl   r   r   r/      r0   z)IbisInterface.groupby.<locals>.<listcomp>c                    s   g | ]}| vr|qS r   r   )r+   Zkdim)
index_dimsr   r   r/      r0   raw)r(   r@   c                 S   s   g | ]
}|j qS r   rw   rx   r   r   r   r/      r0   c              	      sL   g | ]D\}}t |j  j fd d|  D fi fqS )c                    s   g | ]\}} j | |kqS r   r   rj   rl   r   r   r/     r0   z4IbisInterface.groupby.<locals>.<listcomp>.<listcomp>)rP   r1   tolistr   r   Zto_dictro   )r+   is)r@   group_kwargs
group_typer   r   r/     s   
F)r(   
issubclassr   r8   r   Zget_param_valuesupdater@   r   groupbyrQ   r>   Ziterrowsr   r   r   )
r   r@   rq   Zcontainer_typer   kwargsZelement_dimsZgroup_bygroupsr   r   )r@   r   r   r   r   r      s     


ZzIbisInterface.groupbyc                 C   s   |j jf i |S r=   )r   rn   )r   r@   Znew_datar   r   r   assign  s    zIbisInterface.assignc                 C   s\   dd l }|j}|j|jvrXt||jsBt|sBtdt	|j
 |jf i |j|i}|S )Nr   zSCannot assign %s type as a Ibis table column, expecting either ibis.Expr or scalar.)r   r   rO   r7   r%   r$   r\   r~   rt   r   __name__rn   )r   r@   rS   Zdim_posr1   Zvdimr   r   r   r   r   add_dimension  s    zIbisInterface.add_dimensionc                 C   s&   |j |j|ddj    dkS )NTrI   r	   )r   rK   rO   rU   rC   rA   )r   r@   dimr   r   r   r~   %  s     zIbisInterface.isscalarc                 K   s   |d u r|  ||}| ||}|j}t|tjr| |}|jtdkrZt|d }|	|d 
ttt|dg}n |d urt|tr|r|	|}|r|  dkrt|jdkr||jd j  jd S |S )NrF   r   rs   r	   )select_maskindexedr   r%   r\   Zndarrayrv   rL   wherer   r   r   r}   r3   r   rC   r>   r5   r)   rO   r   )r   r@   Zselection_mask	selectionr   r   r   r   r   select0  s$    

"zIbisInterface.selectc                 C   s,  dd l }g }| D ]\}}t|tr0t| }||j}|j| }t|tr|jd urtt	
|j}	||	|k |jd urt	
|j}	|||	k  qt|ttfrd }
|D ] }||k}|
d u r|n|
|B }
q|
d ur||
 qt|r||| qt||jr|| q|||k q|S rc   )r   ro   r%   rP   r   rK   rO   r   r   r   numpy_scalar_to_pythonappendr   setr   callabler$   )r   r@   r   r   r   r   objectaliasrT   boundZ	conditionid	predicater   r   r   r   E  s8    




zIbisInterface.select_maskc                 C   s   dd l }| }|j}tdd |D rNdd |D }|||d j | S d }|D ]z}t|rj|g}|spqVd }	t|D ]<\}
}|||
 j |	t
|k}|	d u r|}	q||	|M }	q||d u r|	}qV||	O }qV|d u r|S ||S )Nr   c                 s   s$   | ]}t |pt|d kV  qdS )r	   N)r   r~   r5   r+   r   r   r   r   r{   m  r0   z'IbisInterface.sample.<locals>.<genexpr>c                 S   s"   g | ]}t |tr|d  n|qS )r   )r%   rP   r   r   r   r   r/   n  r0   z(IbisInterface.sample.<locals>.<listcomp>)r   rq   r   r|   rO   r   r   r~   	enumerateliteralr   r   )r   r@   Zsamplesr   Zdimsr   ro   r   sampler   r   rX   pr   r   r   r   h  s,    


zIbisInterface.samplec           	         s  dd l |jfdd|jD  |jddd} |  tjjjjtj	jjjtj
jjjtjjjjtjjjjtjjjjtjjjjtjjjjtjjjjtjjjjtjjjjtjjjjtjjjitrN }tju r(|jf i  fddjD }n$|jf i  fd	djD }n"jf i fd
djD }fdd|D }||fS )Nr   c                    s   g | ]}| v r|j qS r   rw   rx   )rq   r   r   r/     r0   z+IbisInterface.aggregate.<locals>.<listcomp>valuerO   )Zlabelc                    s8   i | ]0}| vr|j jj| | d kd qS )r   )r   )r   r   Countto_exprre   )r7   r   newr   r   rm     s   z+IbisInterface.aggregate.<locals>.<dictcomp>c                    s&   i | ]}| vr||   qS r   r   re   )r7   functionr   r   r   rm     s   c                    s   i | ]}| |   qS r   r   re   )r   r   r   r   rm     r0   c                    s   g | ]}| j vr|qS r   )r7   re   r   r   r   r/     r0   ) r   r   r(   rq   r\   r4   r   r   ZMinZnanminrR   ZMaxZnanmaxZmeanZMeanZnanmeanZstdZStandardDevZnanstdsumZSumZnansumvarZVarianceZnanvarr5   r   getr   Zcount_nonzerorQ   r7   )	r   r@   rq   r   r   r1   r   ZaggregationZdroppedr   )r7   r   rq   r   r   r   r   rQ     sN    


zIbisInterface.aggregatec                 C   s   t dd S )Nz*Mask is not implemented for IbisInterface.)r_   )r   r@   maskZ
mask_valuer   r   r   r     s    zIbisInterface.maskc                 C   s   |j |  S r=   )r   r>   rp   r   r   r   dframe  s    zIbisInterface.dframe)TTTF)TN)N)-r   r   __qualname__typesZdatatypeZdefault_partitionsr   classmethodr   r    r#   r'   r<   rA   rB   r   rE   rH   rN   r1   ra   rb   rL   Zdimension_typeri   rr   r   ZPandasInterfaceZvalidateZreindexrv   r   r   r   r   r   r~   r   r   r   rQ   r\   nanr   r   r   r   r   r   r      s   






    


&
	
#

	
"
1r   )r!   r\   collections.abcr   ImportErrorcollections r   elementr   Z	ndmappingr   r   r   Z	interfacer
   r   r   r   registerr   r   r   r   <module>   s       2