
    ɯei&                     B   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
 d dlmc mc mc mc mZ d dlmc m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mZmZ d dl m!Z!m"Z"m#Z# d d	l$m%Z%m&Z& d d
l'm(Z(m)Z) d dl*m+Z+m,Z,m-Z-m.Z.m/Z/m0Z0 d dl1m2Z2 d dl3m4Z4m5Z5 d dl6m7Z7m8Z8m9Z9m:Z:m;Z;m<Z<m=Z=m>Z> d dl?m@Z@ d dlAmBZB d dlCmDZD dedefdZEdeFdedeGdefdZHdeFdeGdefdZI G d d      ZJ G d deJ      ZK G d deJ      ZL G d d eJ      ZM G d! d"eJ      ZN G d# d$      ZO G d% d&      ZPy)'    )CallableDictIterableListOptionalTupleUnionN)SnowparkClientExceptionMessages)pandas)unquote_if_quoted)	functions)
ExpressionLiteralNamedExpressionScalarSubquerySnowflakeUDFUnresolvedAttribute)CubeGroupingSetsExpressionRollup)AliasUnresolvedAlias)	AggregatePivot)build_expr_from_python_val+build_expr_from_snowpark_column_or_col_namebuild_proto_from_callablebuild_proto_from_struct_typedebug_check_missing_astwith_src_position)relational_group_df_api_usage)ColumnOrNameLiteralType)check_agg_exprsexperimentalis_valid_tuple_for_aggparse_positional_args_to_list&parse_positional_args_to_list_variadicprepare_pivot_arguments	publicapigenerate_random_alphanumericColumn)	DataFrame)
StructTypeexprreturnc                     t        | t              rt        |       S t        | t        t        f      r| S t        | | j                  j                         j                  dd            S )N" )	
isinstancer   r   r   r   r   sqlupperreplace)r0   s    q/var/www/html/glpi_dashboard/venv/lib/python3.12/site-packages/snowflake/snowpark/relational_grouped_dataframe.py_aliasr:   9   sQ    $+,t$$	D?L9	:T488>>+33C<==    
input_expr	_emit_astc                 h   | j                         }fd}|dv r%t        j                   ||            j                  S |dv r%t        j                   ||            j                  S |dv r%t        j
                   ||            j                  S t        j                  | |      j                  S )Nc                     t        |       S )Nr=   r,   )r<   r=   s    r9   create_columnz$_expr_to_func.<locals>.create_columnE   s    jI66r;   )avgaveragemean)stddevstd)countsizer@   )lowerr   rB   _expressionrE   rG   _call_function)r0   r<   r=   loweredrA   s     `  r9   _expr_to_funcrM   B   s    jjlG7 ,,}}]:67CCC	%	%j 9:FFF	%	%}Z89EEE''*	

+	r;   c                       fdS )Nc                     t        |       S N)rM   )r<   r=   r0   s    r9   <lambda>z_str_to_expr.<locals>.<lambda>U   s    mD*iH r;    )r0   r=   s   ``r9   _str_to_exprrS   T   s	    HHr;   c                       e Zd ZdefdZy)
_GroupTyper1   c                 4    | j                   j                  dd S )N   )	__class____name__)selfs    r9   	to_stringz_GroupType.to_stringY   s    ~~&&q,,r;   N)rZ   
__module____qualname__strr\   rR   r;   r9   rU   rU   X   s    -3 -r;   rU   c                       e Zd Zy)_GroupByTypeNrZ   r]   r^   rR   r;   r9   ra   ra   ]       r;   ra   c                       e Zd Zy)	_CubeTypeNrb   rR   r;   r9   re   re   a   rc   r;   re   c                       e Zd Zy)_RollupTypeNrb   rR   r;   r9   rg   rg   e   rc   r;   rg   c                   @    e Zd Zdedeeee   ef      dee   ddfdZy)
_PivotType	pivot_colvaluesdefault_on_nullr1   Nc                 .    || _         || _        || _        y rP   )rj   rk   rl   )r[   rj   rk   rl   s       r9   __init__z_PivotType.__init__j   s     #.r;   )	rZ   r]   r^   r   r   r	   r   r   rn   rR   r;   r9   ri   ri   i   sF    // tJ/?@A/ "*-	/
 
/r;   ri   c                   D    e Zd ZdZedddeeee   f   deddfd       Z	y)	GroupingSetsa  Creates a :class:`GroupingSets` object from a list of column/expression sets that you pass
    to :meth:`DataFrame.group_by_grouping_sets`. See :meth:`DataFrame.group_by_grouping_sets` for
    examples of how to use this class with a :class:`DataFrame`. See
    `GROUP BY GROUPING SETS <https://docs.snowflake.com/en/sql-reference/constructs/group-by-grouping-sets.html>`_
    for its counterpart in SQL (several examples are shown below).

    =============================================================  ==================================
    Python interface                                               SQL interface
    =============================================================  ==================================
    ``GroupingSets([col("a")], [col("b")])``                       ``GROUPING SETS ((a), (b))``
    ``GroupingSets([col("a") , col("b")], [col("c"), col("d")])``  ``GROUPING SETS ((a, b), (c, d))``
    ``GroupingSets([col("a"), col("b")])``                         ``GROUPING SETS ((a, b))``
    ``GroupingSets(col("a"), col("b"))``                           ``GROUPING SETS ((a, b))``
    =============================================================  ==================================
    Tr@   setsr=   r1   Nc          
         d | _         |rt        j                         | _         t        | j                   j                        }t        | \  }|j                  _        |D ]0  }t        |j                  j                  j                         |       2 t        | }t        |d   t              r|n|g}t        |D cg c]  }|D cg c]  }|j                   c} c}}      | _        y c c}w c c}}w )Nr   )_astprotoExprr    grouping_setsr(   rq   variadicr   argsaddr'   r5   listr   rJ   _to_expression)r[   r=   rq   grouping_sets_astset_listsprepared_setscs           r9   rn   zGroupingSets.__init__   s     	

DI 1$))2I2I J 7=!&&/ Q*+<+A+A+F+F+J+J+LaPQ 6t<'a(8$?Mm_ 	 51>?AQ'amm'?
'?s   ;	C/
C*C/
*C/
)
rZ   r]   r^   __doc__r*   r	   r-   r   boolrn   rR   r;   r9   rp   rp   u   sD      DH
64</0
=A
	
 
r;   rp   c                   j   e Zd ZdZ	 d+dedee   dedee	j                     ddf
dZ	 	 d,d	ee   dee	j                     d
edefdZeeddddeeeeef   eeef   f   dee	j                     d
edefd              Zee	 d-deded
edefd              ZeZe	 	 	 d.dedeeee   ef      dee   d
edd f
d       Zeeddded
edefd              ZeZeeddded
edefd              Z eeddded
edefd              Z!eeddded
edefd              Z"eeddded
edefd              Z#eed-d
edefd              Z$ed-ded
edefd        Z%e%Z&edddeded
edefd!       Z'e e(d"#      e	 d-d$ed%ed
edefd&                     Z)eddd'eded
edefd(       Z*d)e	jV                  ddfd*Z,y)/RelationalGroupedDataFramezRepresents an underlying DataFrame with rows that are grouped by common values.
    Can be used to define aggregations on these grouped DataFrames.

    See also:
        :meth:`snowflake.snowpark.DataFrame.agg`.
    Ndfgrouping_exprs
group_type	_ast_stmtr1   c                 r    || _         || _        || _        d | _        ||j                  | _        y d | _        y rP   )
_dataframe_grouping_exprs_group_type_df_api_calluid_ast_id)r[   r   r   r   r   s        r9   rn   z#RelationalGroupedDataFrame.__init__   s9     -% (1(=y}}4r;   	agg_exprsr=   c           	         |j                  dd      }g }|s| j                  D ]w  }t        |t              rTt	               }|j
                  D 	
cg c]!  }	|	D ]  }
|
|vr|j                  |
      s	 |
 # }}	}
|j                  |       g|j                  |       y |j                  |       t	               }|D 
cg c]  }
|
|vs|j                  |
      s |
 }}
|D 
cg c]  }
t        |
       }}
| j                  D cg c]   }t        |t              r|j                  n|" }}t        | j                  t              r;t        ||| j                  j                   xs | j                  j"                        }nt        | j                  t$              rEt        t'        |      g|| j                  j                   xs | j                  j"                        }n9t        | j                  t(              rDt        t+        |      g|| j                  j                   xs | j                  j"                        }nt        | j                  t,              rt/        |      dk7  r"t/        |      dk(  rt1        j2                         t5        || j                  j6                  | j                  j8                  || j                  j:                  | j                  j                   xs | j                  j"                        }nt=        d| j                         | j                  j                   r| j                  j>                  j@                  jC                  | j                  j>                  j@                  jE                  || j                  j>                  j@                        | j                  j>                  j@                        }tG        | j                  j>                  |||      S c c}
}	w c c}
w c c}
w c c}w )	Nexclude_grouping_columnsFrW   r   zWrong group by type )analyzer)from_r   r   r=   )$getr   r5   r   setrx   ry   extendappendr:   r   childr   ra   r   r   _select_statement_planrg   r   re   r   ri   lenr
   "DF_PIVOT_ONLY_SUPPORT_ONE_AGG_EXPRr   rj   rk   rl   	TypeError_session	_analyzercreate_select_statementcreate_select_snowflake_planr.   )r[   r   r   r=   kwargsr   aliased_agggrouping_exprgr_usedargagr_uniquseduniquer0   unaliased_grouping
group_plans                    r9   _to_dfz!RelationalGroupedDataFrame._to_df   sq    $*::.H%#P '!%!5!5 6m-CD!eG $1#5#5!$ G+Q4 G   &&w/&&}56 	9% u(TATMtxx{!TT*01Qvay11 ,,
 %T51DJJt;
 

 d&&5""11JT__5J5JJ
 ((+6"*+,11JT__5J5JJ
 (()4"()*11JT__5J5JJ
 ((*59~"s+='>!'C5XXZZ"  **  ''  0011JT__5J5JJ 243C3C2DEFF??,,11;;SSoo..88UU)A)A)K)K V  11;;	 T J OO$$	
 	
 U1
s   &O,<	O2O2!O7%O<Tr   exprsc                   |j                  dd      }t        | \  }}t        |      r|g}t        |       d}|r|| j                  j
                  j                  j                         }t        |j                  j                  |      }| j                  |j                         ||j                  _        |D ]0  }	t        |j                  j                   j#                         |	       2 n|}g }
t%        |      dkD  rct'        |d   t(              rP|d   j+                         D ]9  \  }}|
j-                   t/        ||      t1        |      j2                               ; n|D ]  }	t'        |	t0              r|
j-                  |	j2                         /t'        |	t4        t6        f      sFt        |	      sRt'        |	d   t0              r|	d   j2                  nt1        |	d         j2                  }|
j-                   t/        |	d   |      |              | j9                  |
|d      }| j:                  r
t=               ndh|_        |r|j@                  |_!        |S )a  Returns a :class:`DataFrame` with computed aggregates. See examples in :meth:`DataFrame.group_by`.

        Args:
            exprs: A variable length arguments list where every element is

                - A Column object
                - A tuple where the first element is a column object or a column name and the second element is the name of the aggregate function
                - A list of the above

                or a ``dict`` maps column names to aggregate function names.

        Note:
            The name of the aggregate function to compute must be a valid Snowflake `aggregate function
            <https://docs.snowflake.com/en/sql-reference/functions-aggregation.html>`_.

        See also:
            - :meth:`DataFrame.agg`
            - :meth:`DataFrame.group_by`
        r   FNr   rW   r   r=   limit)"r   r(   r&   r$   r   r   
_ast_batchbindr    r0    relational_grouped_dataframe_agg_set_ast_ref
grouped_dfr   rw   r   rx   ry   r   r5   dictitemsr   rS   r-   rJ   rz   tupler   r   r   _ops_after_aggr   r   )r[   r   r=   r   r   r   is_variadicstmtaster   kvcol_exprr   s                  r9   aggzRelationalGroupedDataFrame.agg  s"   8 $*::.H%#P CUK{ "%(GE  //::??A'II>> !!#..1%0		" HA.syy~~/A/A/CQGH !	u:>jq48a( T1  !;a!;F1I<Q<Q!RST  	Na($$Q]]3D%=16LQ6O &adF3 !((#AaD\55 
 $$%B\!A$	%B8%LM	N [[%=  
 &*%9%9CEyBJ	r;   funcoutput_schemac                 .   | j                   D cg c]  }t        |d       }}dg | j                  }| j                  j                  }t	        |      dt
        j                  rf| j                  j                  8| j                  j                  j                  D 	cg c]  }	|	j                   c}	t	        t        j                        j                        |D ]  }
|
j                         }|rt        |      nd}|sNdt        d       j                         }|j!                  |
j#                  |             j!                  |       rt%        |
j&                  t(        t*        f      r#|j!                  |
       j!                  |       j!                  |        } | j                  j,                  | }t	        |       G fdd      }t.        j0                  |j2                  _        |j7                  d|j8                  j:                  D cg c]  }|j<                   c}      |d<   |j7                  d	|j8                  j:                  D cg c]  }|j>                   c}      |d	<    |j@                  jB                  jD                  |f||d
|}|j-                   ||j                   jG                  |d      d      }| j                   r
tI               ndh|_%        |r |j@                  jL                  jO                         }tQ        |jR                  jT                  |      }| jW                  |jX                         t[        |j\                  |j@                  jL                         t_        ||j`                         |jc                         D ]<  \  }}|jd                  jg                         }||_4        tk        |jl                  |       > |jn                  |_8        |S c c}w c c}	w c c}w c c}w )a?  Maps each grouped dataframe in to a pandas.DataFrame, applies the given function on
        data of each grouped dataframe, and returns a pandas.DataFrame. Internally, a vectorized
        UDTF with input ``func`` argument as the ``end_partition`` is registered and called. Additional
        ``kwargs`` are accepted to specify arguments to register the UDTF. Group by clause used must be
        column reference, not a general expression.

        Requires ``pandas`` to be installed in the execution environment and declared as a dependency by either
        specifying the keyword argument `packages=["pandas]` in this call or calling :meth:`~snowflake.snowpark.Session.add_packages` beforehand.

        Args:
            func: A Python native function that accepts a single input argument - a ``pandas.DataFrame``
                object and returns a ``pandas.Dataframe``. It is used as input to ``end_partition`` in
                a vectorized UDTF.
            output_schema: A :class:`~snowflake.snowpark.types.StructType` instance that represents the
                table function's output columns.
            input_names: A list of strings that represents the table function's input column names. Optional,
                if unspecified, default column names will be ARG1, ARG2, etc.
            kwargs: Additional arguments to register the vectorized UDTF. See
                :meth:`~snowflake.snowpark.udtf.UDTFRegistration.register` for all options.

        Examples::
            Call ``apply_in_pandas`` using temporary UDTF:

                >>> import pandas as pd
                >>> from snowflake.snowpark.types import StructType, StructField, StringType, FloatType
                >>> def convert(pandas_df):
                ...     return pandas_df.assign(TEMP_F = lambda x: x.TEMP_C * 9 / 5 + 32)
                >>> df = session.createDataFrame([('SF', 21.0), ('SF', 17.5), ('SF', 24.0), ('NY', 30.9), ('NY', 33.6)],
                ...         schema=['location', 'temp_c'])
                >>> df.group_by("location").apply_in_pandas(convert,
                ...     output_schema=StructType([StructField("location", StringType()),
                ...                               StructField("temp_c", FloatType()),
                ...                               StructField("temp_f", FloatType())])).order_by("temp_c").show()
                ---------------------------------------------
                |"LOCATION"  |"TEMP_C"  |"TEMP_F"           |
                ---------------------------------------------
                |SF          |17.5      |63.5               |
                |SF          |21.0      |69.8               |
                |SF          |24.0      |75.2               |
                |NY          |30.9      |87.61999999999999  |
                |NY          |33.6      |92.48              |
                ---------------------------------------------
                <BLANKLINE>

            Call ``apply_in_pandas`` using permanent UDTF with replacing original UDTF:

                >>> from snowflake.snowpark.types import IntegerType, DoubleType
                >>> _ = session.sql("create or replace temp stage mystage").collect()
                >>> def group_sum(pdf):
                ...     return pd.DataFrame([(pdf.GRADE.iloc[0], pdf.DIVISION.iloc[0], pdf.VALUE.sum(), )])
                ...
                >>> df = session.createDataFrame([('A', 2, 11.0), ('A', 2, 13.9), ('B', 5, 5.0), ('B', 2, 12.1)],
                ...                              schema=["grade", "division", "value"])
                >>> df.group_by([df.grade, df.division] ).applyInPandas(
                ...     group_sum,
                ...     output_schema=StructType([StructField("grade", StringType()),
                ...                                        StructField("division", IntegerType()),
                ...                                        StructField("sum", DoubleType())]),
                ...                is_permanent=True, stage_location="@mystage", name="group_sum_in_pandas", replace=True
                ...            ).order_by("sum").show()
                --------------------------------
                |"GRADE"  |"DIVISION"  |"SUM"  |
                --------------------------------
                |B        |5           |5.0    |
                |B        |2           |12.1   |
                |A        |2           |24.9   |
                --------------------------------
                <BLANKLINE>

        See Also:
            - :class:`~snowflake.snowpark.udtf.UDTFRegistration`
            - :func:`~snowflake.snowpark.functions.pandas_udtf`
        Fr@   Nsas_col_alias_   c                   R    e Zd Zdej                  dej                  f fdZy)BRelationalGroupedDataFrame.apply_in_pandas.<locals>._ApplyInPandaspdfr1   c                 <   dk(  r?dd l }	D cg c]  }||   j                  d    }}|j                  |      }t        |      }
;t	        
      z
  }|dkD  r!|j                  |j                  | d  d      }
|_        dk(  r	 |      S  |      S c c}w )N   r   rW   )axis)numpyilocarrayr   r   dropcolumns)r[   r   npkeykey_listnumpy_arraykeysto_drop_columnsr   key_columnsoriginal_columnsparameters_countworking_columns_counts           r9   end_partitionzPRelationalGroupedDataFrame.apply_in_pandas.<locals>._ApplyInPandas.end_partition  s     $q(&<GHSCa 0HHH"$((8"4K -D#/ '<cBR>S&SO&*!hhs{{O3C3D'EAhN"2CK#q(c?*Cy   Is   BN)rZ   r]   r^   r   r.   r   )r   r   r   r   r   s   r9   _ApplyInPandasr     s)    !!++!!!! !r;   r   input_typesinput_names)r   r=   )partition_byr=   r   )9r   r-   r   r   r   context$_is_snowpark_connect_compatible_mode_column_map
spark_nameinspect	signature
parametersget_namer   r+   r7   r   aliasr5   rJ   r   r   selectr   r.   r   _sf_vectorized_inputr   schemafieldsdatatypenamer   udtfregisteroverr   r   r   r   r    r0   ,relational_grouped_dataframe_apply_in_pandasr   r   r   r   r   r   r   r   ry   _1r   _2r   r   )r[   r   r   r=   r   r0   r   working_dataframeworking_columnscolumncolcolumn_nameunquoted_namealiased_namer   field_apply_in_pandas_udtfr   r   r   r   r   entryr   r   r   r   s    `                     @@@@r9   apply_in_pandasz*RelationalGroupedDataFrame.apply_in_pandasZ  s   f CGBVBVW$tu5WW .2!# OO48OO4K4K #O 477**648OO4O4O4W4W$*0F%%$   #7#4#4T#:#E#EF $ 6!llnBM 1+ >SW" ))Ea)H(IJPPR ! $**399\+BC&&|4%1IJ
 $**3/&&}5&&}5'6,   !7 6 6 ! %($8!	! 	!0 =C<L<L$$9 !'

8I8P8P8W8WXuENNX!
} !'

4E4L4L4S4ST5EJJT!
} !I 1 : : ? ? H H!
'!
 	!
 %%!#4#<#<=BB)U C  	 & 
 &*%9%9CEy$--88==?D#		FFC cnn-%$ 1 : : E E )8I8IJ 81

(*588Q78 BJ	u X$H Y Us   P&P'P
(P
rj   rk   rl   c                    t        | j                  d|||      \  | _        }}}t        |d   ||      | _        |r| j                  j                  j
                  j                         }t        |j                  j                  |      }	|t        |	j                  |       t        |	j                  |       t        |	j                  |       | j                  |	j                          |j"                  | _        | S )aj  Rotates this DataFrame by turning unique values from one column in the input
        expression into multiple columns and aggregating results where required on any
        remaining column values.

        Only one aggregate is supported with pivot.

        Args:
            pivot_col: The column or name of the column to use.
            values: A list of values in the column,
                or dynamic based on the DataFrame query,
                or None (default) will use all values of the pivot column.
            default_on_null: Expression to replace empty result values.

        Example::

            >>> create_result = session.sql('''create or replace temp table monthly_sales(empid int, team text, amount int, month text)
            ... as select * from values
            ... (1, 'A', 10000, 'JAN'),
            ... (1, 'B', 400, 'JAN'),
            ... (2, 'A', 4500, 'JAN'),
            ... (2, 'A', 35000, 'JAN'),
            ... (1, 'B', 5000, 'FEB'),
            ... (1, 'A', 3000, 'FEB'),
            ... (2, 'B', 200, 'FEB') ''').collect()
            >>> df = session.table("monthly_sales")
            >>> df.group_by("empid").pivot("month", ['JAN', 'FEB']).sum("amount").sort(df["empid"]).show()
            -------------------------------
            |"EMPID"  |"'JAN'"  |"'FEB'"  |
            -------------------------------
            |1        |10400    |8000     |
            |2        |39500    |200      |
            -------------------------------
            <BLANKLINE>

            >>> df.group_by(["empid", "team"]).pivot("month", ['JAN', 'FEB']).sum("amount").sort("empid", "team").show()
            ----------------------------------------
            |"EMPID"  |"TEAM"  |"'JAN'"  |"'FEB'"  |
            ----------------------------------------
            |1        |A       |10000    |3000     |
            |1        |B       |400      |5000     |
            |2        |A       |39500    |NULL     |
            |2        |B       |NULL     |200      |
            ----------------------------------------
            <BLANKLINE>

            >>> df = session.table("monthly_sales")
            >>> df.group_by(["empid", "team"]).pivot("month").sum("amount").sort("empid", "team").show()
            ----------------------------------------
            |"EMPID"  |"TEAM"  |"'FEB'"  |"'JAN'"  |
            ----------------------------------------
            |1        |A       |3000     |10000    |
            |1        |B       |5000     |400      |
            |2        |A       |NULL     |39500    |
            |2        |B       |200      |NULL     |
            ----------------------------------------
            <BLANKLINE>

            >>> from snowflake.snowpark.functions import col
            >>> subquery_df = session.table("monthly_sales").select("month").filter(col("month") == "JAN")
            >>> df = session.table("monthly_sales")
            >>> df.group_by(["empid", "team"]).pivot("month", values=subquery_df, default_on_null=999).sum("amount").sort("empid", "team").show()
            ------------------------------
            |"EMPID"  |"TEAM"  |"'JAN'"  |
            ------------------------------
            |1        |A       |10000    |
            |1        |B       |400      |
            |2        |A       |39500    |
            |2        |B       |999      |
            ------------------------------
            <BLANKLINE>
        z RelationalGroupedDataFrame.pivotr   )r)   r   ri   r   r   r   r   r    r0   "relational_grouped_dataframe_pivotr   rl   r   rj   rk   r   r   r   r   )
r[   rj   rk   rl   r=   pcpivot_valuespivot_default_on_nullr   r   s
             r9   pivotz RelationalGroupedDataFrame.pivot+  s    j $OO.
	
O! &be\;PQ ??++66;;=D#DII$P$PRVWC**3+>+>P7yQ&szz6:cnn-  88DLr;   r@   colsc                T    |j                  dd      } | j                  dg|||dS )zReturn the average for the specified numeric columns.

        Args:
            cols: The columns to calculate average for.
        r   FrB   r   r   _non_empty_argument_functionr[   r=   r
  r   r   s        r9   rB   zRelationalGroupedDataFrame.avg  C     $*::.H%#P 0t00

 &>	
 	
r;   c                T    |j                  dd      } | j                  dg|||dS )z}Return the sum for the specified numeric columns.

        Args:
            cols: The columns to calculate sum for.
        r   Fsumr   r  r  s        r9   r  zRelationalGroupedDataFrame.sum  r  r;   c                T    |j                  dd      } | j                  dg|||dS )zReturn the median for the specified numeric columns.

        Args:
            cols: The columns to calculate median for.
        r   Fmedianr   r  r  s        r9   r  z!RelationalGroupedDataFrame.median  sC     $*::.H%#P 0t00

 &>	
 	
r;   c                T    |j                  dd      } | j                  dg|||dS )z}Return the min for the specified numeric columns.

        Args:
            cols: The columns to calculate min for.
        r   Fminr   r  r  s        r9   r  zRelationalGroupedDataFrame.min  r  r;   c                T    |j                  dd      } | j                  dg|||dS )z}Return the max for the specified numeric columns.

        Args:
            cols: The columns to calculate max for.
        r   Fmaxr   r  r  s        r9   r  zRelationalGroupedDataFrame.max  r  r;   c           
         |j                  dd      }| j                  t        t        j                  dt        d      d      j                  d      g|d      }| j                  r
t               ndh|_	        |r| j                  j                  j                  j                         }t        |j                  j                   |      }| j#                  |j$                         d|_        |j(                  |_        |S )z)Return the number of rows for each group.r   FrG   rW   r@   r   r   )r   r   r   r   rK   r   rJ   r   r   r   r   r   r   r   r    r0   $relational_grouped_dataframe_builtinr   r   agg_namer   r   )r[   r=   r   r   r   r   r   s          r9   rG   z RelationalGroupedDataFrame.count  s     $*::.H%#P [[,,u!k	 &>  
 &*%9%9CEy ??++66;;=D#		>>C cnn-"CLBJ	r;   r  c                 >     |j                  dd       fdS )a%  Computes the builtin aggregate ``agg_name`` over the specified columns. Use
        this function to invoke any aggregates not explicitly listed in this class.
        See examples in :meth:`DataFrame.group_by`.

        Args:
            agg_name: The name of the aggregate function.
        r   Fc                  2     j                   g| dS )Nr   )	_function)r
  r=   r  r   r[   s    r9   rQ   z5RelationalGroupedDataFrame.function.<locals>.<lambda>  s*    ^T^^

 &>	
 r;   )r   )r[   r  r=   r   r   s   ``` @r9   functionz#RelationalGroupedDataFrame.function  s!     $*::.H%#P 
 	
r;   c                   |j                  dd      }g }|D ]f  }t        |t              rt        |      j                  n|j                  }t        j                  ||d      j                  }	|j                  |	       h | j                  ||      }
| j                  r
t               ndh|
_        |r| j                  j                  j                  j                         }t!        |j"                  j$                  |      }| j'                  |j(                         ||_        t-        | \  }}||j.                  _        |D ]0  }t3        |j.                  j4                  j7                         |       2 |j8                  |
_        |
S )Nr   Fr@   )r   r   )r   r5   r_   r-   rJ   r   rK   r   r   r   r   r   r   r   r   r   r    r0   r  r   r   r  r(   r
  rw   r   rx   ry   r   r   )r[   r  r=   r
  r   r   r   r   c_exprr0   r   r   r   r   r   r   s                   r9   r  z$RelationalGroupedDataFrame._function!  sR    $*::.H%#P 	 	#A.8C.@VAY**ammF++&Ek  T"	# [[=U[V &*%9%9CEy??++66;;=D#		>>C cnn-#CL!G!NE; +CHH C*388==+<+<+>BC BJ	r;   z1.39.0)versionr0   task_descriptionc                    |j                  dd      }t        |t              rt        |      j                  n|j                  }t        j                  t        |d      t        j                  |d      d      j                  }| j                  |g|d      }| j                  r
t               ndh|_        |r| j                  j                  j                  j                         }	t!        |	j"                  j$                  |	      }
| j'                  |
j(                         t+        |
j"                  |       ||
_        |	j.                  |_        |S )a  Aggregate a column of text data using a natural language task description.

        This method reduces a column of text by performing a natural language aggregation
        as described in the task description for each group. For instance, it can summarize
        large datasets or extract specific insights per group.

        Args:
            expr: The column (Column object or column name as string) containing the text data
                on which the aggregation operation is to be performed.
            task_description: A plain English string that describes the aggregation task, such as
                "Summarize the product reviews for a blog post targeting consumers" or
                "Identify the most positive review and translate it into French and Polish, one word only".

        Returns:
            A DataFrame with one row per group containing the aggregated result.

        Example::

            >>> df = session.create_dataframe([
            ...     ["electronics", "Excellent product, highly recommend!"],
            ...     ["electronics", "Great quality and fast shipping"],
            ...     ["clothing", "Perfect fit and great material"],
            ...     ["clothing", "Poor quality, very disappointed"],
            ... ], schema=["category", "review"])
            >>> summary_df = df.group_by("category").ai_agg(
            ...     expr="review",
            ...     task_description="Summarize these product reviews for a blog post targeting consumers"
            ... )
            >>> summary_df.count()
            2

        Note:
            For optimal performance, follow these guidelines:

                - Use plain English text for the task description.

                - Describe the text provided in the task description. For example, instead of a task
                  description like "summarize", use "Summarize the phone call transcripts".

                - Describe the intended use case. For example, instead of "find the best review",
                  use "Find the most positive and well-written restaurant review to highlight on
                  the restaurant website".

                - Consider breaking the task description into multiple steps.
        r   Fr@   r   r   )r   r5   r_   r-   rJ   r   ai_agglitr   r   r   r   r   r   r   r   r    r0   #relational_grouped_dataframe_ai_aggr   r   r   r"  r   r   )r[   r0   r"  r=   r   r   expr_colagg_exprr   r   r   s              r9   r$  z!RelationalGroupedDataFrame.ai_aggC  s   n $*::.H%#P  )34(=F4L$$4CSCS 	
 ##8u-MM*e<
 +	 	 [[J%=  
 &*%9%9CEy??++66;;=D#DII$Q$QSWXCcnn-&sxx6#3C BJ	r;   	func_namec                t    |j                  dd      }|st        d|        | j                  |||      | S )Nr   Fz9You must pass a list of one or more Columns to function: r   )r   
ValueErrorbuiltin)r[   r)  r=   r
  r   r   s         r9   r  z7RelationalGroupedDataFrame._non_empty_argument_function  s]     $*::.H%#P KI;W 4<<)A#    	 r;   expr_builderc                     t        | j                  | j                  j                  | j                         | j                  |j                  _        y)z
        Given a field builder expression of the AST type Expr, points the builder to reference this RelationalGroupedDataFrame.
        N)r   r   r   r    relational_grouped_dataframe_refid)r[   r-  s     r9   r   z'RelationalGroupedDataFrame._set_ast_ref  s4     	 doo.F.FX;?<<558r;   rP   )NF)T)NNT)-rZ   r]   r^   r   r.   r   r   rU   r   rt   Bindrn   r   r   r!   r*   r	   r-   r   r"   r_   r   r   r   r/   r  applyInPandasr   r#   r	  rB   rD   r  r  r  r  rG   r  r,  r  r%   r$  r  ru   r   rR   r;   r9   r   r      s!    +/HH Z(H 	H
 EJJ'H 
H  +/	S

#S
 EJJ'S
 	S
 
S
j # +/	NfeL#$56S#XFGN EJJ'N 	N 
N  #N` #
 	KK "K 	K 
K  #KZ $M EI15kk x4i?@Ak "+.	k
 k 
&k kZ #9= 
 
$ 
I 
  #
 D"9= 
 
$ 
I 
  #
 #59
!
.2
	
  #
  #9= 
 
$ 
I 
  #
 #9= 
 
$ 
I 
  #
 #t    #> 
 
 
8 
 
  GDH$0=A	 B #(#
 	UU U 	U 
U  $ #Un EI%1>B	 H H Hr;   r   )Qtypingr   r   r   r   r   r   r	   r   *snowflake.snowpark._internal.error_messager
   4snowflake.snowpark._internal.proto.generated.ast_pb2snowpark	_internalrt   	generatedast_pb2snowflake.snowpark.contextr   snowflake.connector.optionsr   4snowflake.snowpark._internal.analyzer.analyzer_utilsr   snowflake.snowparkr   0snowflake.snowpark._internal.analyzer.expressionr   r   r   r   r   r   2snowflake.snowpark._internal.analyzer.grouping_setr   r   r   6snowflake.snowpark._internal.analyzer.unary_expressionr   r   5snowflake.snowpark._internal.analyzer.unary_plan_noder   r   &snowflake.snowpark._internal.ast.utilsr   r   r   r   r   r    &snowflake.snowpark._internal.telemetryr!   'snowflake.snowpark._internal.type_utilsr"   r#   "snowflake.snowpark._internal.utilsr$   r%   r&   r'   r(   r)   r*   r+   snowflake.snowpark.columnr-   snowflake.snowpark.dataframer.   snowflake.snowpark.typesr/   r:   r_   r   rM   rS   rU   ra   re   rg   ri   rp   r   rR   r;   r9   <module>rI     s   
 J I I  V D D D , , . R (  
 S  Q M	 	 	 - 2 /> > >    $Is It I I- -
	: 		
 		* 		/ 	/&
 &
RTH THr;   