U
    +if                     @  st  d Z ddlmZ ddlZddlZddlmZmZmZm	Z	m
Z
mZmZ ddlZddlmZmZ ddlm  mZ ddlm  mZ ddlmZmZmZmZmZ ddlm Z  ddl!m"Z" dd	l#m$Z$m%Z% dd
l&m'Z'm(Z(m)Z)m*Z*m+Z+m,Z,m-Z-m.Z.m/Z/m0Z0m1Z1m2Z2m3Z3m4Z4 ddl5m6Z6 ddl7m8Z8m9Z9 ddl:m;Z;m<Z<m=Z=m>Z> ddl?m@Z@ ddlAmBZBmCZC ddlDmEZEmFZF ddlGmHZHmIZI ddlJmKZK ddlLmMZM ddlNmOZO ddlPmQZQ ddlRmSZSmTZTmUZUmVZV ddlWmXZX ddlYmZZZm[Z[m\Z\m]Z]m^Z^m_Z_ G dd dZ`G dd dZaG dd deaZbdd d!d"d#ZcG d$d% d%ee ZdG d&d' d'edZeG d(d) d)edZfd/d*d+ddd%d,d-d.ZgdS )0a  
Provide classes to perform the groupby aggregate operations.

These are not exposed to the user and provide implementations of the grouping
operations, primarily in cython. These classes (BaseGrouper and BinGrouper)
are contained *in* the SeriesGroupBy and DataFrameGroupBy objects.
    )annotationsN)CallableGenericHashableIteratorSequencefinaloverload)NaTlib)	ArrayLikeDtypeObjNDFrameTShapenptAbstractMethodError)cache_readonly)maybe_cast_pointwise_resultmaybe_downcast_to_dtype)ensure_float64ensure_int64ensure_platform_intis_1d_only_ea_objis_bool_dtypeis_categorical_dtypeis_complex_dtypeis_datetime64_any_dtypeis_float_dtypeis_integer_dtypeis_numeric_dtype	is_sparseis_timedelta64_dtypeneeds_i8_conversion)ExtensionDtype)isna
maybe_fill)DatetimeArrayExtensionArrayPeriodArrayTimedeltaArray)BooleanDtype)Float64DtypeFloatingDtype)
Int64Dtype_IntegerDtype)BaseMaskedArrayBaseMaskedDtype)StringDtype)	DataFrame)NDFrame)grouper)CategoricalIndexIndex
MultiIndexensure_index)Series)compress_group_indexdecons_obs_group_idsget_flattened_listget_group_indexget_group_index_sorterget_indexer_dictc                   @  s  e Zd ZdZedddddgZdddd	d
Zddddddddddd
dddddddZdddd hZd!d"iZ	e
ejd#d$ddd%d&d'd(d)Zd*d&d+d,d-ZdYd/d&d0d1d2Zd3d*d4d5d6d7Zd%d%d8d9d:Zed%d%d8d;d<Zed=d=d8d>d<Zd/d/d8d?d<Zd&d@dAdBZedCd3d3d*dDdEdFdGZdHdI ZedJd3d3d*dJdEdKdLZed#d#dMd*d3d3d*dNdNd*dOdPdQZed*d3d3d*dNdNd*dOdRdSZedTdUdDd3d3d*d3dDdVdWdXZd#S )ZWrappedCythonOpz?
    Dispatch logic for functions defined in _libs.groupby
    rankcountsizeidxminidxmaxstrkindhowc                 C  s   || _ || _d S NrH   )selfrI   rJ    rM   J/home/mars/bis/venv/lib/python3.8/site-packages/pandas/core/groupby/ops.py__init__s   s    zWrappedCythonOp.__init__Z	group_addZ
group_prodZ	group_minZ	group_maxZ
group_meanZgroup_medianZ	group_varZ	group_nthZ
group_lastZ
group_ohlc)
addprodminmaxmeanmedianvarfirstlastohlcZgroup_cumprodZgroup_cumsumZgroup_cumminZgroup_cummaxZ
group_rank)cumprodcumsumcummincummaxrB   )	aggregate	transformr\   r]   rR   rS   rY      N)maxsizeznp.dtypebool)rI   rJ   dtype
is_numericc                 C  sV   |j }| j| | }tt|}|r&|S |tkrRd|jkrNtd| d| d|S d S )Nobject2function is not implemented for this dtype: [how->,dtype->])name_CYTHON_FUNCTIONSgetattr
libgroupbyre   Z__signatures__NotImplementedError)clsrI   rJ   rc   rd   Z	dtype_strZftypefrM   rM   rN   _get_cython_function   s    

z$WrappedCythonOp._get_cython_function
np.ndarray)valuesrd   c                 C  s   | j }| j}|dkrV|r"t|}ntd| d|jj dttd| d}||fS | |||j|}|jjdkr|dkrt|}||fS )	z
        Find the appropriate cython function, casting if necessary.

        Parameters
        ----------
        values : np.ndarray
        is_numeric : bool

        Returns
        -------
        func : callable
        values : np.ndarray
        )rU   rZ   rf   rg   rh   Zgroup_Z_float64)iu)rP   rV   rQ   rT   rY   )	rJ   rI   r   rm   rc   ri   rk   rl   rp   )rL   rr   rd   rJ   rI   funcrM   rM   rN   get_cython_func_and_vals   s    
z(WrappedCythonOp.get_cython_func_and_valsFr   )rc   rd   c                 C  s   | j }|rdS t|rB|dkr2t| d| dt| dnZt|rZt| dnBt|r||dkrtd| dn t|r|dkrtd| ddS )	z
        Check if we can do this operation with our cython functions.

        Raises
        ------
        NotImplementedError
            This is either not a valid function for this dtype, or
            valid but not implemented in cython.
        N)rP   rQ   r[   rZ   z type does not support z operationsz dtype not supportedz!datetime64 type does not support )rQ   rZ   z"timedelta64 type does not support )rJ   r   	TypeErrorrm   r!   r   r"   )rL   rc   rd   rJ   rM   rM   rN   _disallow_invalid_ops   s    
z%WrappedCythonOp._disallow_invalid_opsintr   )ngroupsrr   returnc                 C  sf   | j }| j}| j|d}|dkr,|df}n6|dkr>tdn$|dkrN|j}n|f|jdd   }|S )N   rY   r`   z<arity of more than 1 is not supported for the 'how' argumentr_   )rJ   rI   _cython_aritygetrm   shape)rL   rz   rr   rJ   rI   arity	out_shaperM   rM   rN   _get_output_shape   s    
z!WrappedCythonOp._get_output_shape)rc   r{   c                 C  s<   | j }|dkrd}nt|r.|j |j }nd}t|S )NrB   float64re   )rJ   r    rI   itemsizenprc   )rL   rc   rJ   	out_dtyperM   rM   rN   get_out_dtype  s    zWrappedCythonOp.get_out_dtypec                 C  s   d S rK   rM   rL   rc   rM   rM   rN   _get_result_dtype  s    z!WrappedCythonOp._get_result_dtyper$   c                 C  s   d S rK   rM   r   rM   rM   rN   r     s    c                 C  s   | j }|dkr>|ttkr(ttjS t|ttfrt S nD|dkrt|ttfrZt	 S t
|sjt|rn|S t|rttjS |S )aI  
        Get the desired dtype of a result based on the
        input dtype and how it was computed.

        Parameters
        ----------
        dtype : np.dtype or ExtensionDtype
            Input dtype.

        Returns
        -------
        np.dtype or ExtensionDtype
            The desired dtype of the result.
        )rP   r[   sumrQ   )rT   rU   rV   )rJ   r   rc   rb   int64
isinstancer+   r/   r.   r,   r   r   r    r   )rL   rc   rJ   rM   rM   rN   r     s    r{   c                 C  s   | j | jkS rK   )rJ   _MASKED_CYTHON_FUNCTIONSrL   rM   rM   rN   	uses_mask:  s    zWrappedCythonOp.uses_maskr(   r   )rr   	min_countrz   comp_idsr{   c                 K  s   t |tr,|  r,| j|f|||d|S t |tttfrJ|jd}nrt |j	t
tfrl|jdtjd}nPt |j	tr|j|j	jtjd}n.t |j	tr|jttjd}ntd|j	 | j|f|||dd|}| jdkr|S | ||S )	zs
        If we have an ExtensionArray, unwrap, call _cython_operation, and
        re-wrap if appropriate.
        r   rz   r   zM8[ns]r   )Zna_valuez,function is not implemented for this dtype: Nr   rz   r   mask)rB   )r   r0   r    _masked_ea_wrap_cython_operationr'   r)   r*   Z_ndarrayviewrc   r+   r/   Zto_numpyr   nanr-   Znumpy_dtyper2   re   rm   _cython_op_ndim_compatrJ   _reconstruct_ea_result)rL   rr   r   rz   r   kwargsnpvalues
res_valuesrM   rM   rN   _ea_wrap_cython_operation=  sD    
	
z)WrappedCythonOp._ea_wrap_cython_operationc                 C  sd   t |jttttfr6| |j}| }|j||dS t	|jr\|
d}t|||jdS tdS )zL
        Construct an ExtensionArray result from an ndarray result.
        rc   i8N)r   rc   r+   r/   r-   r2   r   construct_array_typeZ_from_sequencer#   r   typerm   )rL   rr   r   rc   rn   Zi8valuesrM   rM   rN   r   v  s     


z&WrappedCythonOp._reconstruct_ea_resultr0   c                 K  s   |}|j  }tj|td}|j}	| j|	f|||||d|}
| |j}t	|t
sZt| }| jdkr||
j|jdd|S ||
j|jdd|S dS )z
        Equivalent of `_ea_wrap_cython_operation`, but optimized for masked EA's
        and cython algorithms which accept a mask.
        r   r   rz   r   r   result_maskr^   FcopyN)Z_maskr   r   zerosrb   _datar   r   rc   r   r1   AssertionErrorr   rI   astyper   )rL   rr   r   rz   r   r   orig_valuesr   r   Zarrr   rc   rn   rM   rM   rN   r     s*    


z0WrappedCythonOp._masked_ea_wrap_cython_operation)r   r   znp.ndarray | None)rr   r   rz   r   r   r   r{   c          
      K  s   |j dkr|d d d f }|d k	r2|d d d f }|d k	rJ|d d d f }| j|f|||||d|}	|	jd dkr~|	d S |	jS | j|f|||||d|S )Nr|   r   r   )ndim_call_cython_opr   T)
rL   rr   r   rz   r   r   r   r   Zvalues2dresrM   rM   rN   r     s>    
	z&WrappedCythonOp._cython_op_ndim_compatc             
   K  s
  |}|j }	t|	}
t|	}|r.|d}d}
nLt|	rB|d}n8t|	rf|	tjkrz|jtj	dd}n|
rzt
|	szt|}|j}|d k	r|j}|d k	r|j}| ||}| ||
\}}| |j }ttj||d}| jdkrLtj|tj	d}| jdkr|||||||||d n2| jd	kr:|||||||d
 n|||||| n<|  rt||||||fd|i| n||||||f| | jdkrt|j r|std|}||k }| r|d}tj||< |j}| j| jkr| |j }t||}n|}|S )Nr   TFr   r   r^   )rR   rS   rT   )r   r   is_datetimelike)rP   )Zdatetimeliker   r|   r   )rc   r    r#   r   r   r   r   r   Zuint64r   r   r   r   r   rv   r   r&   emptyrI   r   rJ   r   rS   anyr   cast_blocklistr   r   )rL   rr   r   rz   r   r   r   r   r   rc   rd   r   r   ru   r   resultcountscutoffZempty_groupsZ	res_dtype	op_resultrM   rM   rN   r     s    


	





zWrappedCythonOp._call_cython_op)r   )rr   axisr   r   rz   r{   c          	      K  s   |j dkrtdn0|j dkr0|dksDt|nt|sD|dksDt|j}t|}| || t|tj	s| j
|f|||d|S | j|f|||dd|S )zW
        Call our cython function, with appropriate pre- and post- processing.
           z.number of dimensions is currently limited to 2r|   r   r   Nr   )r   rm   r   r   rc   r    rx   r   r   ndarrayr   r   )	rL   rr   r   r   r   rz   r   rc   rd   rM   rM   rN   cython_operationJ  s:    


z WrappedCythonOp.cython_operation)F)__name__
__module____qualname____doc__	frozensetr   rO   rj   r   r}   classmethod	functools	lru_cacherp   rv   rx   r   r   r	   r   r   r   r   r   r   r   r   r   rM   rM   rM   rN   rA   j   sd   
&$8& +mrA   c                	   @  s@  e Zd ZU dZded< dbdddddd	dd
ddZeddddZeddddZdd Z	eddddZ
dcddddddZeddddd dd!d"Zd#d$ Zeed%d& Zeded'd(dd)d*d+d,Zed-dd.d/Zeed0dd1d2Zed3dd4d5Zed6dd7d8Zed9dd:d;Zed<dd=d>Zeeddd?d@ZedAddBdCZeedDddEdFZedGddHdIZeedddJdKZed0ddLdMZeedNddOdPZedddQdRZedSddTdUZ edfdWdWdddNdXdYdZZ!edgd9d'ddNd[d\d]Z"ed9d'd^d_d`daZ#dS )hBaseGroupera  
    This is an internal Grouper class, which actually holds
    the generated groups

    Parameters
    ----------
    axis : Index
    groupings : Sequence[Grouping]
        all the grouping instances to handle in this grouper
        for example for grouper list to groupby, need to pass the list
    sort : bool, default True
        whether this grouper will give sorted result or not
    group_keys : bool, default True
    mutated : bool, default False
    indexer : np.ndarray[np.intp], optional
        the indexer created by Grouper
        some groupers (TimeGrouper) will sort its axis and its
        group_info is also sorted, so need the indexer to reorder

    r7   r   TFNzSequence[grouper.Grouping]rb   znpt.NDArray[np.intp] | None)r   	groupingssort
group_keysmutatedindexerdropnac                 C  sD   t |tst||| _t|| _|| _|| _|| _|| _	|| _
d S rK   )r   r7   r   r   list
_groupings_sortr   r   r   r   )rL   r   r   r   r   r   r   r   rM   rM   rN   rO     s    

zBaseGrouper.__init__list[grouper.Grouping]r   c                 C  s   | j S rK   )r   r   rM   rM   rN   r     s    zBaseGrouper.groupingsr   c                 C  s   t dd | jD S )Nc                 s  s   | ]}|j V  qd S rK   )rz   .0pingrM   rM   rN   	<genexpr>  s     z$BaseGrouper.shape.<locals>.<genexpr>)tupler   r   rM   rM   rN   r     s    zBaseGrouper.shapec                 C  s
   t | jS rK   )iterindicesr   rM   rM   rN   __iter__  s    zBaseGrouper.__iter__ry   c                 C  s
   t | jS rK   )lenr   r   rM   rM   rN   nkeys  s    zBaseGrouper.nkeysr   r   z#Iterator[tuple[Hashable, NDFrameT]])datar   r{   c                 c  s@   | j ||d}| j}t||D ]\}}||j|ddfV  qdS )
        Groupby iterator

        Returns
        -------
        Generator yielding sequence of (name, subsetted object)
        for each group
        r   groupbymethodN)_get_splittergroup_keys_seqzip__finalize__)rL   r   r   splitterkeyskeygrouprM   rM   rN   get_iterator  s    zBaseGrouper.get_iteratorr4   DataSplitterc                 C  s   | j \}}}t||||dS )z
        Returns
        -------
        Generator yielding subsetted objects

        __finalize__ has not been called for the subsetted objects returned.
        r   )
group_infoget_splitter)rL   r   r   ids_rz   rM   rM   rN   r     s    	zBaseGrouper._get_splitterc                 C  s   | j d jS )
        We are a grouper as part of another's groupings.

        We have a specific method of grouping, so cannot
        convert to a Index for our grouper.
        r   )r   grouping_vectorr   rM   rM   rN   _get_grouper  s    zBaseGrouper._get_grouperc                 C  s:   t | jdkr| jd S | j\}}}t||| j| jS d S Nr|   r   )r   r   levelsr   r=   codes)rL   r   r   rz   rM   rM   rN   r     s    
zBaseGrouper.group_keys_seqr   zDataFrame | Seriesztuple[list, bool])ro   r   r   r{   c                 C  s   | j }| j||d}| j}g }t||}|D ]P\}	}
|
j|dd}
t|
d|	 |
j}||
}|srt|||srd}|	| q,t
|dkrt|dd dkr||jd d  ||fS )	Nr   r   r   ri   Tr   r   )rE   rF   Z	nanargminZ	nanargmax)r   r   r   r   r   re   __setattr__axes_is_indexed_likeappendr   rk   iloc)rL   ro   r   r   r   r   r   Zresult_valuesZzippedr   r   Z
group_axesr   rM   rM   rN   apply  s     
	zBaseGrouper.applyz$dict[Hashable, npt.NDArray[np.intp]]c                 C  sP   t | jdkr&t| jtr&| jd jS dd | jD }dd | jD }t||S )z"dict {group name -> group indices}r|   r   c                 S  s   g | ]
}|j qS rM   r   r   rM   rM   rN   
<listcomp>  s     z'BaseGrouper.indices.<locals>.<listcomp>c                 S  s   g | ]
}|j qS rM   group_indexr   rM   rM   rN   r     s     )r   r   r   result_indexr6   r   r@   )rL   Z
codes_listr   rM   rM   rN   r     s
    zBaseGrouper.indiceslist[np.ndarray]c                 C  s   dd | j D S )Nc                 S  s   g | ]
}|j qS rM   r   r   rM   rM   rN   r     s     z%BaseGrouper.codes.<locals>.<listcomp>r   r   rM   rM   rN   r     s    zBaseGrouper.codeslist[Index]c                 C  s   dd | j D S )Nc                 S  s   g | ]
}|j qS rM   r   r   rM   rM   rN   r     s     z&BaseGrouper.levels.<locals>.<listcomp>r   r   rM   rM   rN   r     s    zBaseGrouper.levelslist[Hashable]c                 C  s   dd | j D S )Nc                 S  s   g | ]
}|j qS rM   )ri   r   rM   rM   rN   r   !  s     z%BaseGrouper.names.<locals>.<listcomp>r   r   rM   rM   rN   names  s    zBaseGrouper.namesr:   c                 C  s<   | j \}}}|r(tj||dk |d}ng }t|| jddS )z&
        Compute group sizes.
        r   )Z	minlengthr   )indexrc   )r   r   Zbincountr:   r   )rL   r   r   rz   outrM   rM   rN   rD   #  s
    zBaseGrouper.sizezdict[Hashable, np.ndarray]c                 C  sF   t | jdkr| jd jS tdd | jD  }t|}| j|S dS )!dict {group name -> group labels}r|   r   c                 s  s   | ]}|j V  qd S rK   )r   r   rM   rM   rN   r   6  s     z%BaseGrouper.groups.<locals>.<genexpr>N)r   r   groupsr   r7   r   r   )rL   Z
to_groupbyr  rM   rM   rN   r  0  s
    zBaseGrouper.groupsc                 C  s   t | jd jS Nr   )r7   r   is_monotonicr   rM   rM   rN   r	  :  s    zBaseGrouper.is_monotonic6tuple[npt.NDArray[np.intp], npt.NDArray[np.intp], int]c                 C  s&   |   \}}t|}t|}|||fS rK   )_get_compressed_codesr   r   )rL   r   obs_group_idsrz   rM   rM   rN   r   @  s    zBaseGrouper.group_infonpt.NDArray[np.intp]c                 C  s:   | j \}}}| jd k	r6t|| jf}|| }t|}|S rK   )r   r   r   Zlexsortr   )rL   r   r   ZsorterrM   rM   rN   
codes_infoI  s    
zBaseGrouper.codes_infoz'tuple[np.ndarray, npt.NDArray[np.intp]]c                 C  sV   t | jdkr0t| j| jddd}t|| jdS | jd }|jtjt |j	tj
dfS )Nr|   T)r   xnull)r   r   r   )r   r   r>   r   r   r;   r   r   aranger   intp)rL   r   r   rM   rM   rN   r  V  s
    
z!BaseGrouper._get_compressed_codesc                 C  s
   t | jS rK   )r   r   r   rM   rM   rN   rz   `  s    zBaseGrouper.ngroupsc                 C  s&   | j }| j\}}}t||| j|ddS )NT)r  )r   r   r<   r   )rL   r   r   Zobs_idsr   rM   rM   rN   reconstructed_codese  s    zBaseGrouper.reconstructed_codesr   c                 C  s"   t | jdkr| jd jS | jjS )z
        Analogous to result_index, but returning an ndarray/ExtensionArray
        allowing us to retain ExtensionDtypes not supported by Index.
        r|   r   )r   r   group_arrayliker   _valuesr   rM   rM   rN   result_arraylikek  s    	zBaseGrouper.result_arraylikec                 C  sN   t | jdkr&| jd j| jd S | j}dd | jD }t||d| jdS )Nr|   r   c                 S  s   g | ]
}|j qS rM   )r   r   rM   rM   rN   r     s     z,BaseGrouper.result_index.<locals>.<listcomp>F)r   r   Zverify_integrityr  )r   r   r   renamer  r  r8   )rL   r   r   rM   rM   rN   r   z  s       zBaseGrouper.result_indexzlist[ArrayLike]c                 C  sZ   t | jdkr| jd jgS g }t| j| jD ]&\}}t|}|j|}|| q.|S r   )r   r   r  r   r  r   taker   )rL   	name_listr   r   r   rM   rM   rN   get_group_levels  s    zBaseGrouper.get_group_levelsr   rG   )rI   rJ   r   r   r{   c                 K  sF   |dkst t||d}| j\}}	}	| j}
|jf |||||
d|S )z;
        Returns the values of a cython operation.
        )r_   r^   rH   )rr   r   r   r   rz   )r   rA   r   rz   r   )rL   rI   rr   rJ   r   r   r   Zcy_opr   r   rz   rM   rM   rN   _cython_operation  s    zBaseGrouper._cython_operation)objru   preserve_dtyper{   c                 C  sr   t |dkr| ||}n,t|jtjs:| ||}d}n| ||}tj|dd}|rjt||j	dd}n|}|S )a1  
        Parameters
        ----------
        obj : Series
        func : function taking a Series and returning a scalar-like
        preserve_dtype : bool
            Whether the aggregation is known to be dtype-preserving.

        Returns
        -------
        np.ndarray or ExtensionArray
        r   TF)Z	try_float)Znumeric_only)
r   _aggregate_series_pure_pythonr   r  r   r   r   Zmaybe_convert_objectsr   rc   )rL   r  ru   r  r   r   r  rM   rM   rN   
agg_series  s    zBaseGrouper.agg_seriesznpt.NDArray[np.object_]r  ru   r{   c                 C  s   | j \}}}tj|td}tj|dd}d}t|||dd}	t|	D ]T\}
}|j|dd}||}t	|}|st
||j d}|jd ||
< |||
< qD|S )	Nr   OFr   r   r   r   T)r   r   r   ry   r   r   	enumerater   libreductionZextract_resultZcheck_result_arrayrc   r   )rL   r  ru   r   r   rz   r   r   Zinitializedr   rs   r   r   rM   rM   rN   r    s    

z)BaseGrouper._aggregate_series_pure_python)TTFNT)r   )r   )r   )r   )F)$r   r   r   r   __annotations__rO   propertyr   r   r   r   r   r   r   r   r   r   r   r   r   r   r  rD   r  r	  r   r  r  rz   r  r  r   r  r  r  r  rM   rM   rM   rN   r   |  s   
      		 $			
  )r   c                   @  s   e Zd ZU dZded< ded< ded< d3dd
ddZedd ZeddddZ	dd Z
d4dddddZedd ZeddddZeddd d!Zed"d# Zed$dd%d&Zed'dd(d)Zed*dd+d,Zd-d.d/d0d1d2Zd	S )5
BinGrouperaU  
    This is an internal Grouper class

    Parameters
    ----------
    bins : the split index of binlabels to group the item of axis
    binlabels : the label list
    mutated : bool, default False
    indexer : np.ndarray[np.intp]

    Examples
    --------
    bins: [2, 4, 6, 8, 10]
    binlabels: DatetimeIndex(['2005-01-01', '2005-01-03',
        '2005-01-05', '2005-01-07', '2005-01-09'],
        dtype='datetime64[ns]', freq='2D')

    the group_info, which contains the label of each item in grouped
    axis, the index of label in label list, group number, is

    (array([0, 0, 1, 1, 2, 2, 3, 3, 4, 4]), array([0, 1, 2, 3, 4]), 5)

    means that, the grouped axis has 10 items, can be grouped into 5
    labels, the first and second items belong to the first label, the
    third and forth items belong to the second label, and so on

    znpt.NDArray[np.int64]binsr7   	binlabelsrb   r   FN)r   c                 C  s<   t || _t|| _|| _|| _t| jt| jks8td S rK   )r   r&  r9   r'  r   r   r   r   )rL   r&  r'  r   r   rM   rM   rN   rO     s
    

zBinGrouper.__init__c                 C  s   dd t | j| jD }|S )r  c                 S  s   i | ]\}}|t k	r||qS rM   )r
   )r   r   valuerM   rM   rN   
<dictcomp>1  s    z%BinGrouper.groups.<locals>.<dictcomp>)r   r'  r&  )rL   r   rM   rM   rN   r  ,  s    zBinGrouper.groupsry   r   c                 C  s   dS )Nr|   rM   r   rM   rM   rN   r   8  s    zBinGrouper.nkeysc                 C  s   | S )r   rM   r   rM   rM   rN   r   =  s    zBinGrouper._get_grouperr   r4   )r   r   c                 #  s   |dkr fdd}n fdd}t  j| }d}t| j| jD ]$\}}|tk	rb||||fV  |}qB||k r| jd ||dfV  dS )r   r   c                   s    j | | S rK   r   startedger   rM   rN   <lambda>P      z)BinGrouper.get_iterator.<locals>.<lambda>c                   s    j d d | |f S rK   r*  r+  r.  rM   rN   r/  R  r0  r   N)r   r   r   r&  r'  r
   )rL   r   r   Zslicerlengthr,  r-  labelrM   r.  rN   r   F  s    	zBinGrouper.get_iteratorc                 C  sP   t t}d}t| j| jD ].\}}||k r|tk	rFtt||||< |}q|S r  )collectionsdefaultdictr   r   r'  r&  r
   range)rL   r   rs   r2  binrM   rM   rN   r   _  s    
zBinGrouper.indicesr
  c                 C  s   | j }tj|tjd}ttjd| jf }t|}|t| jkrVt	t||}nt	tjdt|f |}t|||fS )Nr   r   r   )
rz   r   r  r  diffr_r&  r   r   repeat)rL   rz   r  repr   rM   rM   rN   r   k  s    zBinGrouper.group_infor   c                 C  s2   t jdt | jdd  | jd d kd f gS )Nr   r|   r   )r   r8  Zflatnonzeror&  r   rM   rM   rN   r  }  s    zBinGrouper.reconstructed_codesc                 C  s0   t | jdkr*t| jd r*| jdd  S | jS )Nr   r|   )r   r'  r%   r   rM   rM   rN   r     s    zBinGrouper.result_indexr  c                 C  s   | j gS rK   )r'  r   rM   rM   rN   r     s    zBinGrouper.levelsr  c                 C  s
   | j jgS rK   )r'  ri   r   rM   rM   rN   r    s    zBinGrouper.namesr   c                 C  s   | j }tj||dd d}|gS )NF)Zin_axislevel)r'  r5   ZGrouping)rL   Zlevr   rM   rM   rN   r     s    zBinGrouper.groupingsr:   r   rq   r  c                 C  s   t dd S )Nz=This should not be reached; use _aggregate_series_pure_python)rm   )rL   r  ru   rM   rM   rN   _aggregate_series_fast  s    z!BinGrouper._aggregate_series_fast)FN)r   )r   r   r   r   r#  rO   r   r  r$  r   r   r   r   r   r  r   r   r  r   r<  rM   rM   rM   rN   r%    s8   
  
	

r%  ry   rb   )r   r{   c                 C  sP   t | tr.t|dkrdS | j| || S t | trL| j| || S dS )Nr|   F)r   r:   r   r   equalsr3   )r  r   r   rM   rM   rN   r     s    

r   c                   @  sp   e Zd ZddddddddZeddd	d
ZeddddZdd ZeddddZdddddZ	dS )r   r   r   r  ry   )r   labelsrz   r   c                 C  s2   || _ t|| _|| _|| _t|ts.t|d S rK   )r   r   r>  rz   r   r   ry   r   )rL   r   r>  rz   r   rM   rM   rN   rO     s
    
zDataSplitter.__init__r   c                 C  s   | j | jS rK   )r>  r  	_sort_idxr   rM   rM   rN   slabels  s    zDataSplitter.slabelsc                 C  s   t | j| jS rK   )r?   r>  rz   r   rM   rM   rN   r?    s    zDataSplitter._sort_idxc                 c  sT   | j }| jdkrd S t| j| j\}}t||D ]\}}| |t||V  q2d S r  )sorted_datarz   r   Zgenerate_slicesr@  r   _chopslice)rL   sdataZstartsZendsr,  endrM   rM   rN   r     s    
zDataSplitter.__iter__c                 C  s   | j j| j| jdS )Nr   )r   r  r?  r   r   rM   rM   rN   rA    s    zDataSplitter.sorted_datarC  r4   )	slice_objr{   c                 C  s   t | d S rK   r   )rL   rD  rF  rM   rM   rN   rB    s    zDataSplitter._chopN)r   )
r   r   r   rO   r   r@  r?  r   rA  rB  rM   rM   rM   rN   r     s    r   c                   @  s   e Zd ZddddddZdS )SeriesSplitterr:   rC  rD  rF  r{   c                 C  s   |j |}|j||jddS )NT)ri   Zfastpath)_mgr	get_slice_constructorri   rL   rD  rF  ZmgrrM   rM   rN   rB    s    zSeriesSplitter._chopNr   r   r   rB  rM   rM   rM   rN   rG    s   rG  c                   @  s   e Zd ZddddddZdS )FrameSplitterr3   rC  rH  c                 C  s    |j j|d| j d}||S )Nr|   r   )rI  rJ  r   rK  rL  rM   rM   rN   rB    s    zFrameSplitter._chopNrM  rM   rM   rM   rN   rN    s   rN  r4   rq   )r   r>  rz   r   r{   c                 C  s"   t | trt}nt}|| |||S rK   )r   r:   rG  rN  )r   r>  rz   r   klassrM   rM   rN   r     s    
r   )r   )hr   
__future__r   r3  r   typingr   r   r   r   r   r   r	   numpyr   Zpandas._libsr
   r   Zpandas._libs.groupbyZ_libsr   rl   Zpandas._libs.reductionZ	reductionr"  Zpandas._typingr   r   r   r   r   Zpandas.errorsr   Zpandas.util._decoratorsr   Zpandas.core.dtypes.castr   r   Zpandas.core.dtypes.commonr   r   r   r   r   r   r   r   r   r   r    r!   r"   r#   Zpandas.core.dtypes.dtypesr$   Zpandas.core.dtypes.missingr%   r&   Zpandas.core.arraysr'   r(   r)   r*   Zpandas.core.arrays.booleanr+   Zpandas.core.arrays.floatingr,   r-   Zpandas.core.arrays.integerr.   r/   Zpandas.core.arrays.maskedr0   r1   Zpandas.core.arrays.string_r2   Zpandas.core.framer3   Zpandas.core.genericr4   Zpandas.core.groupbyr5   Zpandas.core.indexes.apir6   r7   r8   r9   Zpandas.core.seriesr:   Zpandas.core.sortingr;   r<   r=   r>   r?   r@   rA   r   r%  r   r   rG  rN  r   rM   rM   rM   rN   <module>   sX   $
@ 
        $. 