From 42e0ad5e34e01cdd58d6173f4871451396bb85a7 Mon Sep 17 00:00:00 2001 From: Leonhard Spiegelberg Date: Thu, 24 Oct 2024 11:08:38 -0700 Subject: [PATCH] rename ast -> _ast --- src/snowflake/snowpark/column.py | 132 ++++++++++++++++------------ src/snowflake/snowpark/dataframe.py | 6 +- src/snowflake/snowpark/functions.py | 16 ++-- src/snowflake/snowpark/udaf.py | 2 +- src/snowflake/snowpark/window.py | 2 +- 5 files changed, 87 insertions(+), 71 deletions(-) diff --git a/src/snowflake/snowpark/column.py b/src/snowflake/snowpark/column.py index 3ad311dc98..95135d0858 100644 --- a/src/snowflake/snowpark/column.py +++ b/src/snowflake/snowpark/column.py @@ -248,10 +248,10 @@ def __init__( self, expr1: Union[str, Expression], expr2: Optional[str] = None, - ast: Optional[proto.Expr] = None, + _ast: Optional[proto.Expr] = None, _emit_ast: bool = True, ) -> None: - self._ast = ast + self._ast = _ast if expr2 is not None: if not (isinstance(expr1, str) and isinstance(expr2, str)): @@ -313,7 +313,7 @@ def __getitem__(self, field: Union[str, int]) -> "Column": ast.col.CopyFrom(self._ast) ast.field = field return Column( - SubfieldString(self._expression, field), ast=expr, _emit_ast=_emit_ast + SubfieldString(self._expression, field), _ast=expr, _emit_ast=_emit_ast ) elif isinstance(field, int): if _emit_ast: @@ -322,7 +322,7 @@ def __getitem__(self, field: Union[str, int]) -> "Column": ast.col.CopyFrom(self._ast) ast.idx = field return Column( - SubfieldInt(self._expression, field), ast=expr, _emit_ast=_emit_ast + SubfieldInt(self._expression, field), _ast=expr, _emit_ast=_emit_ast ) else: raise TypeError(f"Unexpected item type: {type(field)}") @@ -337,7 +337,7 @@ def __eq__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": build_expr_from_snowpark_column_or_python_val(ast.rhs, other) right = Column._to_expr(other) - return Column(EqualTo(self._expression, right), ast=expr, _emit_ast=_emit_ast) + return Column(EqualTo(self._expression, right), _ast=expr, _emit_ast=_emit_ast) def __ne__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": """Not equal to.""" @@ -350,7 +350,7 @@ def __ne__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": right = Column._to_expr(other) return Column( - NotEqualTo(self._expression, right), ast=expr, _emit_ast=_emit_ast + NotEqualTo(self._expression, right), _ast=expr, _emit_ast=_emit_ast ) def __gt__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": @@ -363,7 +363,7 @@ def __gt__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": build_expr_from_snowpark_column_or_python_val(ast.rhs, other) return Column( GreaterThan(self._expression, Column._to_expr(other)), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -377,7 +377,7 @@ def __lt__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": build_expr_from_snowpark_column_or_python_val(ast.rhs, other) return Column( LessThan(self._expression, Column._to_expr(other)), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -391,7 +391,7 @@ def __ge__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": build_expr_from_snowpark_column_or_python_val(ast.rhs, other) return Column( GreaterThanOrEqual(self._expression, Column._to_expr(other)), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -405,7 +405,7 @@ def __le__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": build_expr_from_snowpark_column_or_python_val(ast.rhs, other) return Column( LessThanOrEqual(self._expression, Column._to_expr(other)), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -418,7 +418,9 @@ def __add__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": ast.lhs.CopyFrom(self._ast) build_expr_from_snowpark_column_or_python_val(ast.rhs, other) return Column( - Add(self._expression, Column._to_expr(other)), ast=expr, _emit_ast=_emit_ast + Add(self._expression, Column._to_expr(other)), + _ast=expr, + _emit_ast=_emit_ast, ) def __radd__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": @@ -429,7 +431,9 @@ def __radd__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": build_expr_from_snowpark_column_or_python_val(ast.lhs, other) ast.rhs.CopyFrom(self._ast) return Column( - Add(Column._to_expr(other), self._expression), ast=expr, _emit_ast=_emit_ast + Add(Column._to_expr(other), self._expression), + _ast=expr, + _emit_ast=_emit_ast, ) def __sub__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": @@ -442,7 +446,7 @@ def __sub__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": build_expr_from_snowpark_column_or_python_val(ast.rhs, other) return Column( Subtract(self._expression, Column._to_expr(other)), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -455,7 +459,7 @@ def __rsub__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": ast.rhs.CopyFrom(self._ast) return Column( Subtract(Column._to_expr(other), self._expression), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -469,7 +473,7 @@ def __mul__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": build_expr_from_snowpark_column_or_python_val(ast.rhs, other) return Column( Multiply(self._expression, Column._to_expr(other)), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -482,7 +486,7 @@ def __rmul__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": ast.rhs.CopyFrom(self._ast) return Column( Multiply(Column._to_expr(other), self._expression), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -496,7 +500,7 @@ def __truediv__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": build_expr_from_snowpark_column_or_python_val(ast.rhs, other) return Column( Divide(self._expression, Column._to_expr(other)), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -509,7 +513,7 @@ def __rtruediv__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": ast.rhs.CopyFrom(self._ast) return Column( Divide(Column._to_expr(other), self._expression), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -523,7 +527,7 @@ def __mod__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": build_expr_from_snowpark_column_or_python_val(ast.rhs, other) return Column( Remainder(self._expression, Column._to_expr(other)), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -536,7 +540,7 @@ def __rmod__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": ast.rhs.CopyFrom(self._ast) return Column( Remainder(Column._to_expr(other), self._expression), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -549,7 +553,9 @@ def __pow__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": ast.lhs.CopyFrom(self._ast) build_expr_from_snowpark_column_or_python_val(ast.rhs, other) return Column( - Pow(self._expression, Column._to_expr(other)), ast=expr, _emit_ast=_emit_ast + Pow(self._expression, Column._to_expr(other)), + _ast=expr, + _emit_ast=_emit_ast, ) def __rpow__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": @@ -560,7 +566,9 @@ def __rpow__(self, other: Union[ColumnOrLiteral, Expression]) -> "Column": build_expr_from_snowpark_column_or_python_val(ast.lhs, other) ast.rhs.CopyFrom(self._ast) return Column( - Pow(Column._to_expr(other), self._expression), ast=expr, _emit_ast=_emit_ast + Pow(Column._to_expr(other), self._expression), + _ast=expr, + _emit_ast=_emit_ast, ) def __bool__(self) -> bool: @@ -637,7 +645,7 @@ def in_( proto_ast = ast.sp_column_in__seq proto_ast.col.CopyFrom(self._ast) - return Column(Literal(False), ast=ast, _emit_ast=_emit_ast) + return Column(Literal(False), _ast=ast, _emit_ast=_emit_ast) cols = [_to_col_if_lit(col, "in_") for col in cols] @@ -700,7 +708,7 @@ def validate_value(value_expr: Expression): return Column( InExpression(self._expression, value_expressions), - ast=ast, + _ast=ast, _emit_ast=_emit_ast, ) @@ -739,7 +747,7 @@ def bitand( build_expr_from_snowpark_column_or_python_val(ast.rhs, other) return Column( BitwiseAnd(Column._to_expr(other), self._expression), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -756,7 +764,7 @@ def bitor( build_expr_from_snowpark_column_or_python_val(ast.rhs, other) return Column( BitwiseOr(Column._to_expr(other), self._expression), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -773,7 +781,7 @@ def bitxor( build_expr_from_snowpark_column_or_python_val(ast.rhs, other) return Column( BitwiseXor(Column._to_expr(other), self._expression), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -789,7 +797,7 @@ def __neg__(self) -> "Column": ast = with_src_position(expr.neg) ast.operand.CopyFrom(self._ast) - return Column(UnaryMinus(self._expression), ast=expr, _emit_ast=_emit_ast) + return Column(UnaryMinus(self._expression), _ast=expr, _emit_ast=_emit_ast) @publicapi def equal_null(self, other: "Column", _emit_ast: bool = True) -> "Column": @@ -802,7 +810,7 @@ def equal_null(self, other: "Column", _emit_ast: bool = True) -> "Column": build_expr_from_snowpark_column_or_python_val(ast.rhs, other) return Column( EqualNullSafe(self._expression, Column._to_expr(other)), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -814,7 +822,7 @@ def equal_nan(self, _emit_ast: bool = True) -> "Column": expr = proto.Expr() ast = with_src_position(expr.sp_column_equal_nan) ast.col.CopyFrom(self._ast) - return Column(IsNaN(self._expression), ast=expr, _emit_ast=_emit_ast) + return Column(IsNaN(self._expression), _ast=expr, _emit_ast=_emit_ast) @publicapi def is_null(self, _emit_ast: bool = True) -> "Column": @@ -824,7 +832,7 @@ def is_null(self, _emit_ast: bool = True) -> "Column": expr = proto.Expr() ast = with_src_position(expr.sp_column_is_null) ast.col.CopyFrom(self._ast) - return Column(IsNull(self._expression), ast=expr, _emit_ast=_emit_ast) + return Column(IsNull(self._expression), _ast=expr, _emit_ast=_emit_ast) @publicapi def is_not_null(self, _emit_ast: bool = True) -> "Column": @@ -834,7 +842,7 @@ def is_not_null(self, _emit_ast: bool = True) -> "Column": expr = proto.Expr() ast = with_src_position(expr.sp_column_is_not_null) ast.col.CopyFrom(self._ast) - return Column(IsNotNull(self._expression), ast=expr, _emit_ast=_emit_ast) + return Column(IsNotNull(self._expression), _ast=expr, _emit_ast=_emit_ast) # `and, or, not` cannot be overloaded in Python, so use bitwise operators as boolean operators def __and__(self, other: Union["Column", Any]) -> "Column": @@ -847,7 +855,9 @@ def __and__(self, other: Union["Column", Any]) -> "Column": build_expr_from_snowpark_column_or_python_val(ast.rhs, other) return Column( - And(self._expression, Column._to_expr(other)), ast=expr, _emit_ast=_emit_ast + And(self._expression, Column._to_expr(other)), + _ast=expr, + _emit_ast=_emit_ast, ) def __rand__(self, other: Union["Column", Any]) -> "Column": @@ -858,7 +868,9 @@ def __rand__(self, other: Union["Column", Any]) -> "Column": build_expr_from_snowpark_column_or_python_val(ast.lhs, other) ast.rhs.CopyFrom(self._ast) return Column( - And(Column._to_expr(other), self._expression), ast=expr, _emit_ast=_emit_ast + And(Column._to_expr(other), self._expression), + _ast=expr, + _emit_ast=_emit_ast, ) # pragma: no cover def __or__(self, other: Union["Column", Any]) -> "Column": @@ -870,7 +882,7 @@ def __or__(self, other: Union["Column", Any]) -> "Column": ast.lhs.CopyFrom(self._ast) build_expr_from_snowpark_column_or_python_val(ast.rhs, other) return Column( - Or(self._expression, Column._to_expr(other)), ast=expr, _emit_ast=_emit_ast + Or(self._expression, Column._to_expr(other)), _ast=expr, _emit_ast=_emit_ast ) def __ror__(self, other: Union["Column", Any]) -> "Column": @@ -881,7 +893,9 @@ def __ror__(self, other: Union["Column", Any]) -> "Column": build_expr_from_snowpark_column_or_python_val(ast.lhs, other) ast.rhs.CopyFrom(self._ast) return Column( - And(Column._to_expr(other), self._expression), ast=expr, _emit_ast=_emit_ast + And(Column._to_expr(other), self._expression), + _ast=expr, + _emit_ast=_emit_ast, ) # pragma: no cover def __invert__(self) -> "Column": @@ -893,7 +907,7 @@ def __invert__(self) -> "Column": ast = with_src_position(getattr(expr, "not")) ast.operand.CopyFrom(self._ast) - return Column(Not(self._expression), ast=expr, _emit_ast=_emit_ast) + return Column(Not(self._expression), _ast=expr, _emit_ast=_emit_ast) def _cast( self, to: Union[str, DataType], try_: bool = False, _emit_ast: bool = True @@ -912,7 +926,7 @@ def _cast( ) ast.col.CopyFrom(self._ast) to._fill_ast(ast.to) - return Column(Cast(self._expression, to, try_), ast=expr, _emit_ast=_emit_ast) + return Column(Cast(self._expression, to, try_), _ast=expr, _emit_ast=_emit_ast) @publicapi def cast(self, to: Union[str, DataType], _emit_ast: bool = True) -> "Column": @@ -937,7 +951,7 @@ def desc(self, _emit_ast: bool = True) -> "Column": ast = with_src_position(expr.sp_column_desc) ast.col.CopyFrom(self._ast) return Column( - SortOrder(self._expression, Descending()), ast=expr, _emit_ast=_emit_ast + SortOrder(self._expression, Descending()), _ast=expr, _emit_ast=_emit_ast ) @publicapi @@ -952,7 +966,7 @@ def desc_nulls_first(self, _emit_ast: bool = True) -> "Column": ast.nulls_first.value = True return Column( SortOrder(self._expression, Descending(), NullsFirst()), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -968,7 +982,7 @@ def desc_nulls_last(self, _emit_ast: bool = True) -> "Column": ast.nulls_first.value = False return Column( SortOrder(self._expression, Descending(), NullsLast()), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -981,7 +995,7 @@ def asc(self, _emit_ast: bool = True) -> "Column": ast = with_src_position(expr.sp_column_asc) ast.col.CopyFrom(self._ast) return Column( - SortOrder(self._expression, Ascending()), ast=expr, _emit_ast=_emit_ast + SortOrder(self._expression, Ascending()), _ast=expr, _emit_ast=_emit_ast ) @publicapi @@ -996,7 +1010,7 @@ def asc_nulls_first(self, _emit_ast: bool = True) -> "Column": ast.nulls_first.value = True return Column( SortOrder(self._expression, Ascending(), NullsFirst()), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -1012,7 +1026,7 @@ def asc_nulls_last(self, _emit_ast: bool = True) -> "Column": ast.nulls_first.value = False return Column( SortOrder(self._expression, Ascending(), NullsLast()), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -1035,7 +1049,7 @@ def like(self, pattern: ColumnOrLiteralStr, _emit_ast: bool = True) -> "Column": build_expr_from_snowpark_column_or_python_val(ast.pattern, pattern) return Column( Like(self._expression, Column._to_expr(pattern)), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -1075,7 +1089,7 @@ def regexp( Column._to_expr(pattern), None if parameters is None else Column._to_expr(parameters), ), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -1096,7 +1110,7 @@ def startswith(self, other: ColumnOrLiteralStr, _emit_ast: bool = True) -> "Colu other = snowflake.snowpark.functions.lit(other) return Column( snowflake.snowpark.functions.startswith(self, other)._expression, - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -1118,7 +1132,7 @@ def endswith(self, other: ColumnOrLiteralStr, _emit_ast: bool = True) -> "Column other = snowflake.snowpark.functions.lit(other) return Column( snowflake.snowpark.functions.endswith(self, other)._expression, - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -1147,7 +1161,7 @@ def substr( return Column( snowflake.snowpark.functions.substring(self, start_pos, length)._expression, - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -1168,7 +1182,7 @@ def collate(self, collation_spec: str, _emit_ast: bool = True) -> "Column": ast.collation_spec, collation_spec ) return Column( - Collate(self._expression, collation_spec), ast=expr, _emit_ast=_emit_ast + Collate(self._expression, collation_spec), _ast=expr, _emit_ast=_emit_ast ) @publicapi @@ -1186,7 +1200,7 @@ def contains(self, string: ColumnOrName, _emit_ast: bool = True) -> "Column": build_expr_from_snowpark_column_or_python_val(ast.pattern, string) return Column( snowflake.snowpark.functions.contains(self, string)._expression, - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -1228,7 +1242,7 @@ def name( if variant_is_as is not None: ast.variant_is_as.value = variant_is_as return Column( - Alias(self._expression, quote_name(alias)), ast=expr, _emit_ast=_emit_ast + Alias(self._expression, quote_name(alias)), _ast=expr, _emit_ast=_emit_ast ) @publicapi @@ -1334,7 +1348,7 @@ def within_group( self._expression, [order_by_col._expression for order_by_col in order_by_cols], ), - ast=expr, + _ast=expr, _emit_ast=_emit_ast, ) @@ -1362,7 +1376,7 @@ def _to_expr(cls, expr: Union[ColumnOrLiteral, Expression]) -> Expression: @classmethod def _expr(cls, e: str, ast: Optional[proto.Expr] = None) -> "Column": - return cls(UnresolvedAttribute(e, is_sql_text=True), ast=ast) + return cls(UnresolvedAttribute(e, is_sql_text=True), _ast=ast) # Add these alias for user code migration isin = in_ @@ -1404,9 +1418,9 @@ class CaseExpr(Column): def __init__( self, expr: CaseWhen, - ast: Optional[proto.Expr] = None, + _ast: Optional[proto.Expr] = None, ) -> None: - super().__init__(expr, ast=ast) + super().__init__(expr, _ast=_ast) self._branches = expr.branches @publicapi @@ -1437,7 +1451,7 @@ def when( ), ] ), - ast=self._ast, + _ast=self._ast, ) @publicapi @@ -1449,7 +1463,9 @@ def otherwise(self, value: ColumnOrLiteral, _emit_ast: bool = True) -> "CaseExpr if _emit_ast: case_expr = with_src_position(self._ast.sp_column_case_when.cases.add()) build_expr_from_snowpark_column_or_python_val(case_expr.value, value) - return CaseExpr(CaseWhen(self._branches, Column._to_expr(value)), ast=self._ast) + return CaseExpr( + CaseWhen(self._branches, Column._to_expr(value)), _ast=self._ast + ) # This alias is to sync with snowpark scala else_ = otherwise diff --git a/src/snowflake/snowpark/dataframe.py b/src/snowflake/snowpark/dataframe.py index 46dc7c5aa0..31d9c87dd5 100644 --- a/src/snowflake/snowpark/dataframe.py +++ b/src/snowflake/snowpark/dataframe.py @@ -1269,9 +1269,9 @@ def col(self, col_name: str, _emit_ast: bool = True) -> Column: self._set_ast_ref(col_expr_ast.df) col_expr_ast.col_name = col_name if col_name == "*": - return Column(Star(self._output), ast=expr) + return Column(Star(self._output), _ast=expr) else: - return Column(self._resolve(col_name), ast=expr) + return Column(self._resolve(col_name), _ast=expr) @df_api_usage @publicapi @@ -1360,7 +1360,7 @@ def select( col_expr_ast = ast.cols.add() if ast else proto.Expr() fill_ast_for_column(col_expr_ast, e, None) - col = Column(e, ast=col_expr_ast) + col = Column(e, _ast=col_expr_ast) names.append(col._named()) elif isinstance(e, TableFunctionCall): diff --git a/src/snowflake/snowpark/functions.py b/src/snowflake/snowpark/functions.py index bf1d10a8fa..80b705cef1 100644 --- a/src/snowflake/snowpark/functions.py +++ b/src/snowflake/snowpark/functions.py @@ -274,9 +274,9 @@ def col(name1: str, name2: Optional[str] = None, _emit_ast: bool = True) -> Colu ast = create_ast_for_column(name1, name2, "col") if name2 is None: - return Column(name1, ast=ast) + return Column(name1, _ast=ast) else: - return Column(name1, name2, ast=ast) + return Column(name1, name2, _ast=ast) @overload @@ -310,9 +310,9 @@ def column(name1: str, name2: Optional[str] = None, _emit_ast: bool = True) -> C ast = create_ast_for_column(name1, name2, "column") if _emit_ast else None if name2 is None: - return Column(name1, ast=ast) + return Column(name1, _ast=ast) else: - return Column(name1, name2, ast=ast) + return Column(name1, name2, _ast=ast) @publicapi @@ -347,7 +347,7 @@ def lit(literal: LiteralType, _emit_ast: bool = True) -> Column: return ( literal if isinstance(literal, Column) - else Column(Literal(literal), ast=ast, _emit_ast=_emit_ast) + else Column(Literal(literal), _ast=ast, _emit_ast=_emit_ast) ) @@ -851,7 +851,7 @@ def count_distinct(*cols: ColumnOrName, _emit_ast: bool = True) -> Column: cs = [_to_col_if_str(c, "count_distinct") for c in cols] return Column( FunctionExpression("count", [c._expression for c in cs], is_distinct=True), - ast=ast, + _ast=ast, _emit_ast=_emit_ast, ) @@ -7681,7 +7681,7 @@ def when( ) ] ), - ast=ast, + _ast=ast, ) @@ -9459,7 +9459,7 @@ def _call_function( api_call_source=api_call_source, is_data_generator=is_data_generator, ), - ast=ast, + _ast=ast, _emit_ast=_emit_ast, ) diff --git a/src/snowflake/snowpark/udaf.py b/src/snowflake/snowpark/udaf.py index 9020e52287..db697214df 100644 --- a/src/snowflake/snowpark/udaf.py +++ b/src/snowflake/snowpark/udaf.py @@ -112,7 +112,7 @@ def __call__( udaf_expr = proto.Expr() build_udaf_apply(udaf_expr, self._ast_id, *cols) - return Column(self._create_udaf_expression(exprs), ast=udaf_expr) + return Column(self._create_udaf_expression(exprs), _ast=udaf_expr) def _create_udaf_expression(self, exprs: List[Expression]) -> SnowflakeUDF: if len(exprs) != len(self._input_types): diff --git a/src/snowflake/snowpark/window.py b/src/snowflake/snowpark/window.py index 2d71b461eb..0863ad4c57 100644 --- a/src/snowflake/snowpark/window.py +++ b/src/snowflake/snowpark/window.py @@ -465,7 +465,7 @@ def _with_aggregate( ) -> "snowflake.snowpark.column.Column": spec = WindowSpecDefinition(self.partition_spec, self.order_spec, self.frame) return snowflake.snowpark.column.Column( - WindowExpression(aggregate, spec), ast=ast, _emit_ast=_emit_ast + WindowExpression(aggregate, spec), _ast=ast, _emit_ast=_emit_ast ) orderBy = order_by