moar changes

Author: sfc-gh-batur

src/snowflake/snowpark/_internal/ast/utils.py

@@ -1490,10 +1490,14 @@ def clear_symbols_and_udfs(message: proto.Request) -> None:
     for stmt in message.body:
         if hasattr(stmt, "assign"):
             stmt.assign.ClearField("symbol")
-        if str(stmt.assign.expr.udf):
-            stmt.assign.expr.ClearField("udf")
-        if str(stmt.assign.expr.stored_procedure):
-            stmt.assign.expr.ClearField("stored_procedure")
+            if str(stmt.assign.expr.udf):
+                stmt.assign.expr.ClearField("udf")
+            if str(stmt.assign.expr.stored_procedure):
+                stmt.assign.ClearField("expr")
+            if str(stmt.assign.expr.apply_expr):
+                if str(stmt.assign.expr.apply_expr.fn):
+                    if str(stmt.assign.expr.apply_expr.fn.stored_procedure):
+                        stmt.assign.ClearField("expr")
 
 
 def base64_str_to_request(base64_str: str) -> proto.Request:

src/snowflake/snowpark/_internal/proto/ast.proto

@@ -22,10 +22,6 @@ message List_SpDataType {
   repeated SpDataType list = 1;
 }
 
-message List_SpStructField {
-  repeated SpStructField list = 1;
-}
-
 message List_String {
   repeated string list = 1;
 }
@@ -178,7 +174,7 @@ message SpStructField {
 
 // sp-type.ir:46
 message SpStructType {
-  List_SpStructField fields = 1;
+  repeated SpStructField fields = 1;
   bool structured = 2;
 }
 

tests/ast/decoder.py

@@ -551,8 +551,6 @@ def get_statement_params(self, d: Dict):
         return statement_params
 
     def decode_expr(self, expr: proto.Expr, **kwargs) -> Any:
-        if not hasattr(expr, "WhichOneof"):
-            breakpoint()
         match expr.WhichOneof("variant"):
             # COLUMN BINARY OPERATIONS
             case "add":
@@ -614,14 +612,19 @@ def decode_expr(self, expr: proto.Expr, **kwargs) -> Any:
                                     fn_name, *pos_args, **named_args
                                 )
                         case "stored_procedure":
-                            return self.session.call(fn_name, *pos_args, **named_args)
+                            self.session._ast_batch.assign()
+                            return None
+                            # return self.session.call(fn_name, *pos_args, **named_args, _emit_ast=False)
 
                         case _:
                             raise ValueError(
                                 "Unknown function reference type: %s"
                                 % expr.apply_expr.fn.WhichOneof("variant")
                             )
 
+                if isinstance(fn, snowflake.snowpark.stored_procedure.StoredProcedure):
+                    return None
+
                 result = fn(*pos_args, **named_args)
                 if hasattr(expr, "var_id"):
                     self.symbol_table[expr.var_id.bitfield1] = (
@@ -1646,7 +1649,6 @@ def decode_expr(self, expr: proto.Expr, **kwargs) -> Any:
                 return self.decode_expr(expr.sp_session_table_function.fn, **kwargs)
 
             case "sp_table":
-                breakpoint()
                 assert expr.sp_table.HasField("name")
                 table_name = self.decode_name_expr(expr.sp_table.name)
                 is_temp_table_for_cleanup = expr.sp_table.is_temp_table_for_cleanup
@@ -1957,29 +1959,31 @@ def decode_expr(self, expr: proto.Expr, **kwargs) -> Any:
                 return df.write
 
             case "stored_procedure":
-                input_types = [
-                    self.decode_data_type_expr(input_type)
-                    for input_type in expr.stored_procedure.input_types.list
-                ]
-                execute_as = expr.stored_procedure.execute_as
-                comment = expr.stored_procedure.comment.value
+                self.session._ast_batch.assign()
                 registered_object_name = self.decode_name_expr(
                     expr.stored_procedure.func.object_name
                 )
-                return_type = self.decode_data_type_expr(
-                    expr.stored_procedure.return_type
-                )
-                name = self.decode_name_expr(expr.stored_procedure.name)
-                ret_sproc = sproc(
-                    self.session.sproc._registry[registered_object_name],
-                    name=name,
-                    return_type=return_type,
-                    input_types=input_types,
-                    execute_as=execute_as,
-                    comment=comment,
-                    _registered_object_name=registered_object_name,
-                )
-                return ret_sproc
+                return self.session.sproc._registry[registered_object_name]
+                """input_types = [
+                        self.decode_data_type_expr(input_type)
+                        for input_type in expr.stored_procedure.input_types.list
+                    ]
+                    execute_as = expr.stored_procedure.execute_as
+                    comment = expr.stored_procedure.comment.value
+                    return_type = self.decode_data_type_expr(
+                        expr.stored_procedure.return_type
+                    )
+                    name = self.decode_name_expr(expr.stored_procedure.name)
+                    ret_sproc = sproc(
+                        self.session.sproc._registry[registered_object_name],
+                        name=name,
+                        return_type=return_type,
+                        input_types=input_types,
+                        execute_as=execute_as,
+                        comment=comment,
+                        _registered_object_name=registered_object_name,
+                    )
+                    return ret_sproc"""
 
             case "sp_flatten":
                 input = self.decode_expr(expr.sp_flatten.input)
@@ -2023,7 +2027,7 @@ def decode_expr(self, expr: proto.Expr, **kwargs) -> Any:
                     )
 
             case "sp_sql":
-                params = [self.decode_expr(param) for parm in expr.sp_sql.params]
+                params = [self.decode_expr(param) for param in expr.sp_sql.params]
                 query = expr.sp_sql.query
                 return self.session.sql(query=query, params=params)
 
@@ -2045,8 +2049,8 @@ def decode_expr(self, expr: proto.Expr, **kwargs) -> Any:
                     case "sp_save_mode_error_if_exists":
                         mode = "error_if_exists"
 
-                    case "sp_save_mode_truncate":
-                        mode = "truncate"
+                    case "sp_save_mode_ignore":
+                        mode = "ignore"
 
                     case "_":
                         mode = None