
    Wh4                       U d dl 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
mZmZ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 e
r;d dlmZmZmZmZmZ d dlZd dl m!Z" d dl#m$Z$m%Z% d dl&m'Z'm(Z( d dl)m*Z* d dl+m,Z, dZ-de.d<   dZ/de.d<   ed   Z0de.d<   edddddddd d!d"d#d$d%d&d'd(d)d*d+d,e0f   Z1de.d-<   	 d.Z2de.d/<   	 eZ3de.d0<   	  ed12      d<d3       Z4 G d4 d5      Z5 G d6 d7ed8d9e1f         Z6d=d:Z7d>d;Z8y)?    )annotationsN)	lru_cache)chain)methodcaller)TYPE_CHECKINGAnyClassVarLiteral)EagerGroupBy)issue_warning)!evaluate_output_names_and_aliases)
zip_strict)is_pandas_like_dataframe)CallableIterableIteratorMappingSequence)DataFrameGroupBy)	TypeAliasUnpack)NarwhalsAggregationScalarKwargs)PandasLikeDataFrame)PandasLikeExprz._NativeGroupBy[tuple[str, ...], Literal[True]]r   NativeGroupByz(Callable[[pd.DataFrame], pd.Series[Any]]NativeApply)covskewInefficientNativeAggregationanyallcountfirstidxmaxidxminlastmaxmeanmedianminmodenuniqueprodquantilesemsizestdsumvarNativeAggregationz.Callable[[Any], pd.DataFrame | pd.Series[Any]]
_NativeAggNonStrHashable    )maxsizec                   | dk(  rt        | d      S |r|j                  d      dk(  rt        |       S t        | fi |S )Nr-   F)dropnaddof   )r   get)namekwdss     X/var/www/html/jupyter_env/lib/python3.12/site-packages/narwhals/_pandas_like/group_by.py_native_aggrB   ?   sE    yD//488F#q(D!!%%%    c                      e Zd ZU dZded<   ded<   ded<   ddZddZ	 	 	 	 dd	Zdd
ZddZ	ddZ
edd       Zedd       ZddZy)AggExpraM  Wrapper storing the intermediate state per-`PandasLikeExpr`.

    There's a lot of edge cases to handle, so aim to evaluate as little
    as possible - and store anything that's needed twice.

    Warning:
        While a `PandasLikeExpr` can be reused - this wrapper is valid **only**
        in a single `.agg(...)` operation.
    r   exprzSequence[str]output_namesaliasesc                <    || _         d| _        d| _        d| _        y )N  )rF   rG   rH   
_leaf_name)selfrF   s     rA   __init__zAggExpr.__init__W   s    	57rC   c               ~    |j                   }|j                  }t        | j                  ||      \  | _        | _        | S )zd**Mutating operation**.

        Stores the results of `evaluate_output_names_and_aliases`.
        )	compliantexcluder   rF   rG   rH   )rM   group_bydfrQ   s       rA   with_expand_nameszAggExpr.with_expand_names]   s@    
 ""*KIIr7+
'4< rC   c                  | j                   }| j                         r,| j                         r|j                  j	                         }n| j                         r}|j                  j	                         }|j
                  j                         }|j                  |D cg c],  }|j                  |      j                  |      j                  . c}      }nn| j                         r|j
                  }| j                  j                  d      x}dk7  rd| d|j                   d}	t        |	      t!        |      }
|j                  }|j"                  |j$                  }}|j                         }|j                  |
D cg c]w  }  |j&                  g ||fi |j	                         j)                  d      j+                  |      j&                  |fi ||   j-                  d      j/                         y c}      }nAt1        |      dk(  r|d	   n
t!        |      } | j3                         |j                  |         }t5        |      rt!        | j6                        |_        |S | j6                  d	   |_        |S c c}w c c}w )
z8Evaluate the wrapped expression as a group_by operation.keepr!   z`Expr.mode(keep='z7')` is not implemented in group by context for backend z3

Hint: Use `nw.col(...).mode(keep='any')` instead.F)	ascendingr=   r   )rG   is_lenis_top_level_function_groupedr1   rP   __narwhals_namespace___concat_horizontalfrom_nativealiasnativeis_modekwargsr>   _implementationNotImplementedErrorlist_keys_kwargsgroupbysort_valuesreset_indexhead
sort_indexlen
native_aggr   rH   columnsr?   )rM   rR   namesresultresult_singlensr?   rP   rV   msgcolsr_   keysra   colselects                   rA   _getitem_aggszAggExpr._getitem_aggsi   sO   
 !!;;=T779&&++-F[[]$--224M##::<B**NSTd.44T:AATF \\^ **I//E9'v .(889 :HH 
 *#..;D%%F#>>8+;+;&D 113B**  $	 NFNN<T<3<:6:TV [5[1 [%WT	- &,	- .1	2
 T!WZ\"	F "%UqU1Xd5kF&T__&x'8'8'@AF#F+!$,,/FN  ,,q/FKM U(	s   1J?A<Jc                     | j                   dk(  S )Nrl   	leaf_namerM   s    rA   rX   zAggExpr.is_len   s    ~~&&rC   c                     | j                   dk(  S )Nr,   rz   r|   s    rA   r`   zAggExpr.is_mode   s    ~~''rC   c                4    | j                   j                  dk(  S )Nr   )rF   _depthr|   s    rA   rY   zAggExpr.is_top_level_function   s    yy1$$rC   c                .    | j                   j                  S N)rF   _scalar_kwargsr|   s    rA   ra   zAggExpr.kwargs   s    yy'''rC   c                    | j                   x}r|S t        j                  | j                        | _         | j                   S r   )rL   PandasLikeGroupByrF   )rM   r?   s     rA   r{   zAggExpr.leaf_name   s6    ??"4"K+66tyyArC   c                h    t        t        j                  | j                        fi | j                  S )z@Return a partial `DataFrameGroupBy` method, missing only `self`.)rB   r   _remap_expr_namer{   ra   r|   s    rA   rm   zAggExpr.native_agg   s-    ..t~~>
BF++
 	
rC   N)rF   r   returnNone)rR   r   r   rE   )rR   r   r   zpd.DataFrame | pd.Series[Any])r   bool)r   r   )r   zNarwhalsAggregation | Any)r   r6   )__name__
__module____qualname____doc____annotations__rN   rT   rx   rX   r`   rY   propertyra   r{   rm   rJ   rC   rA   rE   rE   H   sw     8
2)2	&2h'(% ( (  
rC   rE   c                      e Zd ZU ddddddddd	d
dddddZded<   ded<   	 ded<   	 ded<   	 ded<   	 ed"d       Z	 	 	 	 	 	 	 	 d#dZd$dZ	 	 	 	 	 	 d%dZ		 	 	 	 d&dZ
d'dZd(dZd)d Zy!)*r   r3   r)   r*   r(   r+   r,   r2   r4   r1   r-   r#   r/   r"   r!   )r3   r)   r*   r(   r+   r,   r2   r4   rl   n_uniquer#   r/   r"   r!   z9ClassVar[Mapping[NarwhalsAggregation, NativeAggregation]]_REMAP_AGGStuple[str, ...]_original_columnsz	list[str]re   _output_key_nameszMapping[str, bool]rf   c                    | j                   S )z>Group keys to ignore when expanding multi-output aggregations.)_excluder|   s    rA   rQ   zPandasLikeGroupBy.exclude   s     }}rC   c              0   t        |j                        | _        || _        | j	                  ||      \  | _        | _        | _        g | j                  | j                  | _        | j                  j                  }t        |j                  j                        j                  | j                  j                        r|j                  d      }dd|dd| _         |j"                  | j                  j%                         fi | j                   | _        y )NT)dropF)sortas_indexr;   observed)tuplern   r   _drop_null_keys_parse_keys_compliant_framere   r   r   rP   r_   setindexro   intersectionri   rf   rg   copyrZ   )rM   rS   ru   drop_null_keysr_   s        rA   rN   zPandasLikeGroupBy.__init__   s     "'rzz!2-DHDTDTE
Atz4+A *P4::)O8N8N)O &&v||!!"//0F0FG''T'2F $	
 (6v~~djjoo6G'X4<<'XrC   c                t   d}g }|D ]?  }|j                  t        |      j                  |              | j                  |      r>d}A |r|r;| j                  j                         }|j                  | j                  |            }n| j                  j                         j                  t        | j                  j                        | j                        }n;| j                  j                  j                  r
t!               | j#                  |      }|j%                  d       | j'                  ||      S )NTF)rn   )inplace)appendrE   rT   
_is_simplerP   r[   r\   rx   __native_namespace__	DataFramerd   rZ   groupsre   r_   emptyempty_results_error_apply_aggsri   _select_results)rM   exprsall_aggs_are_simple	agg_exprsrF   rr   rp   s          rA   aggzPandasLikeGroupBy.agg   s   "#%	 	,DWT]<<TBC??4(&+#	,
 ^^::<..t/A/A)/LM<<>HH--.

 I  ^^""((%''%%e,F 	4(##FI66rC   c                  t        j                  d |D              } | j                  j                  |d      j                  g | j
                  | j                  t        t        | j
                  | j                                    S )zgResponsible for remapping temp column names back to original.

        See `ParseKeysGroupBy`.
        c              3  4   K   | ]  }|j                     y wr   )rH   ).0es     rA   	<genexpr>z4PandasLikeGroupBy._select_results.<locals>.<genexpr>  s     'Ea		'Es   F)validate_column_names)
r   from_iterablerP   _with_nativesimple_selectre   renamedictzipr   )rM   rS   r   	new_namess       rA   r   z!PandasLikeGroupBy._select_results  sx     '''E9'EE	DNN''%'H]4 JJ4)24VDTZZ)?)?@AB	
rC   c               J    |D cg c]  }|j                  |        c}S c c}w r   )rx   )rM   r   r   s      rA   rx   zPandasLikeGroupBy._getitem_aggs!  s!     055!%555s    c                    t                | j                  j                  }| j                  |      }| j                  j
                  }|j                         r|j                         dk\  r
 ||d      S  ||      S )a"  Stub issue for `include_groups` [pandas-dev/pandas-stubs#1270].

        - [User guide] mentions `include_groups` 4 times without deprecation.
        - [`DataFrameGroupBy.apply`] doc says the default value of `True` is deprecated since `2.2.0`.
        - `False` is explicitly the only *non-deprecated* option, but entirely omitted since [pandas-dev/pandas-stubs#1268].

        [pandas-dev/pandas-stubs#1270]: https://github.com/pandas-dev/pandas-stubs/issues/1270
        [User guide]: https://pandas.pydata.org/pandas-docs/stable/user_guide/groupby.html
        [`DataFrameGroupBy.apply`]: https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.core.groupby.DataFrameGroupBy.apply.html
        [pandas-dev/pandas-stubs#1268]: https://github.com/pandas-dev/pandas-stubs/pull/1268
        )   r   F)include_groups)warn_complex_group_byrP   rb   _apply_exprs_functionrZ   apply	is_pandas_backend_version)rM   r   implfuncr   s        rA   r   zPandasLikeGroupBy._apply_aggs&  si     	~~--))%0##>> 5 5 76 Ae44T{rC   c                      j                   j                         j                  j                  d fd}|S )Nc                    j                   j                  |       fdD        }|rt        | ng g f\  }} ||      j                  S )Nc              3     K   | ]6  } |      D ])  }|j                   j                  d    |j                  f + 8 yw)r   N)r_   ilocr?   )r   rF   ru   rP   s      rA   r   zFPandasLikeGroupBy._apply_exprs_function.<locals>.fn.<locals>.<genexpr>@  sJ       O  !!!$dii00s   <?)r   context)rP   r   r   r_   )	rS   results	out_group	out_namesrP   r   into_seriesrr   rM   s	       @rA   fnz3PandasLikeGroupBy._apply_exprs_function.<locals>.fn>  sU    33B7I!G
 <C:w#7R Iyy	2FMMMrC   )rS   pd.DataFramer   zpd.Series[Any])rP   r[   _seriesr   )rM   r   r   r   rr   s   `` @@rA   r   z'PandasLikeGroupBy._apply_exprs_function:  s7    ^^224jj..	N 	N 	rC   c              #  2  K   t        j                         5  t        j                  ddt               | j                  j
                  }| j                  D ](  \  }}|  ||      j                  | j                   f * 	 d d d        y # 1 sw Y   y xY ww)Nignorez#.*a length 1 tuple will be returned)messagecategory)	warningscatch_warningsfilterwarningsFutureWarningrP   r   rZ   r   r   )rM   with_nativekeygroups       rA   __iter__zPandasLikeGroupBy.__iter__J  s     $$& 	W##=&
 ..55K"mm W
U<K.<<d>T>TUVVW	W 	W 	Ws   BA*B	BBBN)r   r   )rS   r   ru   z(Sequence[PandasLikeExpr] | Sequence[str]r   r   r   r   )r   r   r   r   )r   zSequence[AggExpr]rS   r   r   r   )r   zIterable[AggExpr]r   z#list[pd.DataFrame | pd.Series[Any]])r   Iterable[PandasLikeExpr]r   r   )r   r   r   r   )r   z)Iterator[tuple[Any, PandasLikeDataFrame]])r   r   r   r   r   r   rQ   rN   r   r   rx   r   r   r   rJ   rC   rA   r   r      s     NKJ   '&EO  8K YY 7Y Y 
Y874
.?

	
6&6	,6
( 	WrC   r   r   r   c                     d} t        |       S )zJDon't even attempt this, it's way too inconsistent across pandas versions.au  No results for group-by aggregation.

Hint: you were probably trying to apply a non-elementary aggregation with a pandas-like API.
Please rewrite your query such that group-by aggregations are elementary. For example, instead of:

    df.group_by('a').agg(nw.col('b').round(2).mean())

use:

    df.with_columns(nw.col('b').round(2)).group_by('a').agg(nw.col('b').mean())

)
ValueError)rs   s    rA   r   r   V  s    	^  c?rC   c                 $    t        dt               y )Na)  Found complex group-by expression, which can't be expressed efficiently with the pandas API. If you can, please rewrite your query such that group-by aggregations are simple (e.g. mean, std, min, max, ...). 

Please see: https://narwhals-dev.github.io/narwhals/concepts/improve_group_by_operation/)r   UserWarningrJ   rC   rA   r   r   e  s    	W
 	rC   )r?   r5   r@   zUnpack[ScalarKwargs]r   r6   )r   r   )r   r   )9
__future__r   r   	functoolsr   	itertoolsr   operatorr   typingr   r   r	   r
   narwhals._compliantr   narwhals._exceptionsr   narwhals._expression_parsingr   narwhals._utilsr   narwhals.dependenciesr   collections.abcr   r   r   r   r   pandaspdpandas.api.typingr   _NativeGroupBytyping_extensionsr   r   narwhals._compliant.typingr   r   narwhals._pandas_like.dataframer   narwhals._pandas_like.exprr   r   r   r   r    r5   r6   r7   rB   rE   r   r   r   rJ   rC   rA   <module>r      s1   "    ! 8 8 , . J & :OOD3LC9OM9OCY C*1-*@ i @&		
	
	

	
			 )"  9 . hH
I H @  	  6 2& &n
 n
bZW&(8:KKLZWzrC   