
    }i|<                    @   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 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e0d<   dZ1de0d<   ed         Z2de0d<   eddddddd d!d"d#d$d%d&d'd(d)d*d+d,e2f         Z3de0d-<   	 d.Z4de0d/<   	 eZ5de0d0<   	 d d1d d2Z6d3e0d4<    ed56          dFd;            Z7 G d< d=          Z8 G d> d?ed@dAe3f                   Z9dGdCZ:dHdEZ;dS )I    )annotationsN)	lru_cache)chain)methodcaller)TYPE_CHECKINGAnyClassVarLiteral)EagerGroupBy)issue_warning)!evaluate_output_names_and_aliases)make_group_by_kwargs)
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idxmaxidxminmaxmeanmedianminmodenthnuniqueprodquantilesemsizestdsumvarNativeAggregationz.Callable[[Any], pd.DataFrame | pd.Series[Any]]
_NativeAggNonStrHashable)firstlast	any_valuez,Mapping[NarwhalsAggregation, Literal[0, -1]]_REMAP_ORDERED_INDEX    )maxsizenamekwdsUnpack[ScalarKwargs]returnc                   | dk    rt          | d          S |r|                    d          dk    rt          |           S t          | fi |S )Nr-   F)dropnaddof   )r   get)r?   r@   s     R/home/jrussi/.local/lib/python3.11/site-packages/narwhals/_pandas_like/group_by.py_native_aggrI   E   sb    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
ddZddZddZedd            Zd dZdS )!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aliasesrB   Nonec                >    || _         d| _        d| _        d| _        d S )N  )rM   rN   rO   
_leaf_name)selfrM   s     rH   __init__zAggExpr.__init__]   s#    	57rJ   group_byPandasLikeGroupByc               h    |j         }|j        }t          | j        ||          \  | _        | _        | S )zd**Mutating operation**.

        Stores the results of `evaluate_output_names_and_aliases`.
        )	compliantexcluder   rM   rN   rO   )rU   rW   dfr[   s       rH   with_expand_nameszAggExpr.with_expand_namesc   s=    
 "*KIr7+
 +
'4< rJ   pd.DataFrame | pd.Series[Any]c                    |j         }| j        }|                                 r*|                                 r|                                }n|                                 rQ|                                |j                                                            fd|D                       }n|                                 r|j        }|	                    | j
                  }|                    d          x}dk    rd| d|j         d}t          |          t          |          }	|j        |j        |j        c|                                                    fd|	D                       }n|                                 s(|                                 s|                                 rI |                                 |g |j        |                   }|                    |j        d	           nMt/          |          d
k    r|d         nt          |          }
 |                                 ||
                   }t1          |          rt          | j                  |_        n| j        d         |_        |S )z8Evaluate the wrapped expression as a group_by operation.c                j    g | ]/}                                                   |          j        0S rR   )from_nativealiasnative).0r?   nsresult_singles     rH   
<listcomp>z)AggExpr._getitem_aggs.<locals>.<listcomp>z   s6    TTTd..44T::ATTTrJ   keepr"   z`Expr.mode(keep='z7')` is not implemented in group by context for backend z3

Hint: Use `nw.col(...).mode(keep='any')` instead.c                   g | ]}  j         g |fi                                                     d                               |          j         fi |                             d                                          S )F)	ascendingrF   )groupbyr1   sort_valuesreset_indexhead
sort_index)rd   colkeyskwargsrc   s     rH   rg   z)AggExpr._getitem_aggs.<locals>.<listcomp>   s     	 	 	 NFN<T<3<::6::TVV [5[11 [%%T	- - &,	- - .1	2
 T!WWZ\\	 	 	rJ   TinplacerF   r   )_groupedrN   is_lenis_top_level_functionr1   rZ   __narwhals_namespace___concat_horizontalis_mode_kwargsrM   rG   _implementationNotImplementedErrorlistrc   _keys_group_by_kwargsis_lastis_firstis_any_value
native_agg	set_indexlenr   rO   columnsr?   )rU   rW   groupednamesresultrZ   node_kwargsrh   msgcolsselectrq   rr   rc   re   rf   s              @@@@@rH   _getitem_aggszAggExpr._getitem_aggso   s   #!;;== +	8T7799 +	8\\^^FF[[]] )	8#LLNNM#::<<B**TTTTTeTTT FF \\^^ #	8 *I"**4955K#///E99H H H(8H H H 
 *#...;;D%F#>8+DLD& 1133B**	 	 	 	 	 	  $	 	 	 FF \\^^ 	8t}} 	8$2C2C2E2E 	8&T__&&w/H/H%/H'IJJFX^T::::!$UqU1XXd5kkF&T__&&wv77F#F++ 	*!$,//FNN,q/FKrJ   boolc                    | j         dk    S )Nr   	leaf_namerU   s    rH   rv   zAggExpr.is_len   s    ~&&rJ   c                    | j         dk    S )Nr:   r   r   s    rH   r   zAggExpr.is_last       ~''rJ   c                    | j         dk    S )Nr9   r   r   s    rH   r   zAggExpr.is_first   s    ~((rJ   c                    | j         dk    S )Nr+   r   r   s    rH   rz   zAggExpr.is_mode   r   rJ   c                    | j         dk    S )Nr;   r   r   s    rH   r   zAggExpr.is_any_value   s    ~,,rJ   c                z    t          t          | j        j                                                            dk    S )NrF   )r   r~   rM   	_metadataop_nodes_reversedr   s    rH   rw   zAggExpr.is_top_level_function   s.    4	+==??@@AAQFFrJ   NarwhalsAggregation | Anyc                n    | j         x}r|S t                               | j                  | _         | j         S N)rT   rX   rM   )rU   r?   s     rH   r   zAggExpr.leaf_name   s5    ?"4 	K+66tyAArJ   r6   c                n   t                               | j                  }t          | j        j                                                  }| j        t          v rL|j        	                    d          rd}t          |          t          dt          | j                           S t          |fi |j        S )z@Return a partial `DataFrameGroupBy` method, missing only `self`.ignore_nullszd`Expr.any_value(ignore_nulls=True)` is not supported in a `group_by` context for pandas-like backendr,   )n)rX   _remap_expr_namer   nextrM   r   r   r<   rr   rG   r}   r   rI   )rU   native_name	last_noder   s       rH   r   zAggExpr.native_agg   s    '88HH,>>@@AA	>111##N33 /6  *#...)=dn)MNNNN;;;)*:;;;rJ   N)rM   r   rB   rP   )rW   rX   rB   rL   )rW   rX   rB   r^   )rB   r   )rB   r   )rB   r6   )__name__
__module____qualname____doc____annotations__rV   r]   r   rv   r   r   rz   r   rw   propertyr   r   rR   rJ   rH   rL   rL   N   s3          8 8 8 8
 
 
 
5 5 5 5n' ' ' '( ( ( () ) ) )( ( ( (- - - -G G G G    X< < < < < <rJ   rL   c                     e Zd ZU i 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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?d1Z
d@d5ZdAd7ZdBd9Zd:S )CrX   r3   r(   r)   r'   r*   r+   r2   r4   r   r1   n_uniquer-   r$   r/   r#   r"   r9   r,   r:   r;   z9ClassVar[Mapping[NarwhalsAggregation, NativeAggregation]]_REMAP_AGGStuple[str, ...]_original_columnsz	list[str]r   _output_key_nameszMapping[str, bool]r   rB   c                    | j         S )z>Group keys to ignore when expanding multi-output aggregations.)_excluder   s    rH   r[   zPandasLikeGroupBy.exclude   s     }rJ   r\   r   rq   (Sequence[PandasLikeExpr] | Sequence[str]drop_null_keysr   rP   c                 t          |j                  | _        || _        |                     ||          \  | _        | _        | _        g | j        | j        R | _        t          |          | _
        | j        j        | _        t          | j        j        j                                      | j        j                  r"| j                            d          | _        d S d S )N)r   T)drop)tupler   r   _drop_null_keys_parse_keys_compliant_framer   r   r   r   r   rZ   rc   _nativesetindexr   intersectionrm   )rU   r\   rq   r   s       rH   rV   zPandasLikeGroupBy.__init__   s     "'rz!2!2-DHDTDTE
 E
Atz4+A *P4:)O8N)O)O 4N S S S ~,t|!'((55dn6LMM 	?<333>>DLLL	? 	?rJ   exprsr   c                h   d}g }d}|D ]}|                     t          |                              |                      |                     |          sd}t	          |j                                                  }|j                            dd          x}r"|r||k    rd| d| d}t          |          |}|rR | j
                            t          |          d	          j        | j                                        fi | j        }	n/ | j
        j        | j                                        fi | j        }	|	| _        |r|rB| j                                        }
|
                    |                     |                    }n{| j                                                            t          |	j                  | j        
          }n5| j        j        j        rt7                      |                     |	|          }|                    d           |                     ||          S )NTrR   Forder_byz?Only one `order_by` can be specified in `group_by`. Found both z and .r9   )na_position)r   rs   )appendrL   r]   
_is_simpler   r   r   rr   rG   r}   r   rl   r~   rk   r   copyr   ru   rZ   rx   ry   r   __native_namespace__	DataFramegroupsrc   emptyempty_results_error_apply_aggsrm   _select_results)rU   r   all_aggs_are_simple	agg_exprsr   rM   md_current_order_byr   r   re   r   s               rH   aggzPandasLikeGroupBy.agg  sZ   "#%	 		- 		-DWT]]<<TBBCCC??4(( ,&+#dn668899B$&IMM*b$A$AA  - 3 1X = =\dk|C-c222, 	W&T\%=%=XG &> & &djoo''&B &B+/+@&B &BGG +dl*4:??+<+<VV@UVVG 	6 ^::<<..t/A/A)/L/LMM<<>>HH(($* I   ^"( 	6%'''%%gu55F 	4(((##FI666rJ   r   Sequence[AggExpr]pd.DataFramec               
   t          j        d |D                       } | j                            |d          j        g | j        |R                      t          t          | j        | j	                                      S )zgResponsible for remapping temp column names back to original.

        See `ParseKeysGroupBy`.
        c              3  $   K   | ]}|j         V  d S r   )rO   )rd   es     rH   	<genexpr>z4PandasLikeGroupBy._select_results.<locals>.<genexpr><  s$      'E'Ea	'E'E'E'E'E'ErJ   F)validate_column_names)
r   from_iterablerZ   _with_nativesimple_selectr   renamedictzipr   )rU   r\   r   	new_namess       rH   r   z!PandasLikeGroupBy._select_results5  s     ''E'E9'E'E'EEE	DN''%'HH4 J4)24 4 4VDTZ)?@@AABB	
rJ   Iterable[AggExpr]#list[pd.DataFrame | pd.Series[Any]]c                      fd|D             S )Nc                :    g | ]}|                               S rR   )r   )rd   r   rU   s     rH   rg   z3PandasLikeGroupBy._getitem_aggs.<locals>.<listcomp>F  s%    555!%%555rJ   rR   )rU   r   s   ` rH   r   zPandasLikeGroupBy._getitem_aggsC  s     6555u5555rJ   r   r   Iterable[PandasLikeExpr]c                    t                       | j        j        }|                     |          }|j        }|                                r%|                                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_byrZ   r|   _apply_exprs_functionapply	is_pandas_backend_version)rU   r   r   implfuncr   s         rH   r   zPandasLikeGroupBy._apply_aggsH  s     	~-))%00>> 	5 5 5 7 76 A A5e4444uT{{rJ   r   c                j      j                                         j        j        d fd}|S )Nr\   r   rB   pd.Series[Any]c                    j                             |           fdD             }|r	t          | ng g f\  }} ||          j        S )Nc              3  f   K   | ]+} |          D ]}|j         j        d          |j        fV  ,dS )r   N)rc   ilocr?   )rd   rM   rq   rZ   s      rH   r   zFPandasLikeGroupBy._apply_exprs_function.<locals>.fn.<locals>.<genexpr>d  sg         DOO   !!$di0      rJ   )r   context)rZ   r   r   rc   )	r\   results	out_group	out_namesrZ   r   into_seriesre   rU   s	       @rH   fnz3PandasLikeGroupBy._apply_exprs_function.<locals>.fnb  s{    33B77I   !  G
 <C#P:w#7#7R Iy;y	2FFFMMrJ   )r\   r   rB   r   )rZ   rx   _seriesr   )rU   r   r   r   re   s   `` @@rH   r   z'PandasLikeGroupBy._apply_exprs_function^  s`    ^2244j.	N 	N 	N 	N 	N 	N 	N 	N 	N 	rJ   )Iterator[tuple[Any, PandasLikeDataFrame]]c              #  Z  K    | j         j        | j                                        fi | j        }t          j                    5  t          j        ddt                     | j	        j
        }|D ]!\  }}|  ||          j        | j         fV  "	 d d d            d S # 1 swxY w Y   d S )Nignorez#.*a length 1 tuple will be returned)messagecategory)r   rk   r   r   r   warningscatch_warningsfilterwarningsFutureWarningrZ   r   r   r   )rU   r   with_nativekeygroups        rH   __iter__zPandasLikeGroupBy.__iter__n  s.     &$,&tz'8'8RRD<QRR$&& 	W 	W#=&   
 .5K% W W
U<KK..<d>TUVVVVVW	W 	W 	W 	W 	W 	W 	W 	W 	W 	W 	W 	W 	W 	W 	W 	W 	W 	Ws   AB  B$'B$N)rB   r   )r\   r   rq   r   r   r   rB   rP   )r   r   rB   r   )r   r   r\   r   rB   r   )r   r   rB   r   )r   r   r   r   rB   r   )r   r   rB   r   )rB   r   )r   r   r   r   r   r   r[   rV   r   r   r   r   r   r
  rR   rJ   rH   rX   rX      s        NuNN 	(N 	u	N
 	uN 	N 	uN 	uN 	vN 	IN 	N 	JN 	uN 	uN 	N  	!N" 	U#NK    & '&&&EO    8((((K   X? ? ? ?,'7 '7 '7 '7R
 
 
 
6 6 6 6
   ,    
W 
W 
W 
W 
W 
WrJ   rX   r   r   
ValueErrorc                 $    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())

)r  )r   s    rH   r   r   {  s    	^  c??rJ   rP   c                 0    t          dt                     d S )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   UserWarningrR   rJ   rH   r   r     s)    	W
 	    rJ   )r?   r5   r@   rA   rB   r6   )rB   r  )rB   rP   )<
__future__r   r  	functoolsr   	itertoolsr   operatorr   typingr   r   r	   r
   narwhals._compliantr   narwhals._exceptionsr   narwhals._expression_parsingr   narwhals._pandas_like.utilsr   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   r<   rI   rL   rX   r   r   rR   rJ   rH   <module>r#     sa   " " " " " " "              ! ! ! ! ! ! 8 8 8 8 8 8 8 8 8 8 8 8 , , , , , , . . . . . . J J J J J J < < < < < < & & & & & & : : : : : : POOOOOOOOOOOOOODDDDDD33333333LLLLLLLLCCCCCC999999OMOOOOC C C C C*1-*@  @ @ @ @&			
	
	
	
			 '"      , hH
 H H H H @       6 F F      2& & & &~< ~< ~< ~< ~< ~< ~< ~<BiW iW iW iW iW&(8:KKLiW iW iWX        rJ   