SNOW-1794510 refactor decoder

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

@@ -4,6 +4,7 @@
 syntax = "proto3";
 
 option java_package = "com.snowflake.snowpark.proto";
+option java_outer_classname = "JavaProto";
 
 package ast;
 
@@ -14,6 +15,10 @@ enum __Version__ {
   MAX_VERSION = 1;
 }
 
+message InternedValueTable {
+  map<int32, string> string_values = 1;
+}
+
 message List_Expr {
   repeated Expr list = 1;
 }
@@ -22,6 +27,10 @@ message List_SpDataType {
   repeated SpDataType list = 1;
 }
 
+message List_SpStructField {
+  repeated SpStructField list = 1;
+}
+
 message List_String {
   repeated string list = 1;
 }
@@ -59,7 +68,7 @@ message Tuple_String_String {
   string _2 = 2;
 }
 
-// ast.ir:15
+// ast.ir:17
 message Language {
   oneof sealed_value {
     JavaLanguage java_language = 1;
@@ -68,17 +77,17 @@ message Language {
   }
 }
 
-// ast.ir:16
+// ast.ir:18
 message PythonLanguage {
   Version version = 1;
 }
 
-// ast.ir:17
+// ast.ir:19
 message ScalaLanguage {
   Version version = 1;
 }
 
-// ast.ir:18
+// ast.ir:20
 message JavaLanguage {
   Version version = 1;
 }
@@ -174,7 +183,7 @@ message SpStructField {
 
 // sp-type.ir:46
 message SpStructType {
-  repeated SpStructField fields = 1;
+  List_SpStructField fields = 1;
   bool structured = 2;
 }
 
@@ -358,7 +367,7 @@ message SpWindowRelativePosition_Position {
 message SrcPosition {
   int64 end_column = 1;
   int64 end_line = 2;
-  string file = 3;
+  int32 file = 3;
   int64 start_column = 4;
   int64 start_line = 5;
 }
@@ -386,12 +395,12 @@ message UdtfSchema_Names {
   repeated string schema = 1;
 }
 
-// ast.ir:53
+// ast.ir:55
 message VarId {
   uint64 bitfield1 = 1;
 }
 
-// ast.ir:21
+// ast.ir:23
 message Version {
   string label = 1;
   int64 major = 2;
@@ -421,7 +430,7 @@ message ApplyExpr {
   SrcPosition src = 4;
 }
 
-// ast.ir:34
+// ast.ir:36
 message Assign {
   Expr expr = 1;
   google.protobuf.StringValue symbol = 2;
@@ -559,13 +568,13 @@ message Error {
   }
 }
 
-// ast.ir:42
+// ast.ir:44
 message Eval {
   int64 uid = 1;
   VarId var_id = 2;
 }
 
-// ast.ir:64
+// ast.ir:66
 message EvalOk {
   EvalResult data = 1;
   int64 uid = 2;
@@ -1236,15 +1245,17 @@ message Ref {
 
 // ast.ir:4
 message Request {
-  repeated Stmt body = 1;
-  int64 client_ast_version = 2;
-  Language client_language = 3;
-  Version client_version = 4;
+  InternedValueTable interned_value_table = 1;
+  repeated Stmt body = 2;
+  int64 client_ast_version = 3;
+  Language client_language = 4;
+  Version client_version = 5;
 }
 
-// ast.ir:11
+// ast.ir:12
 message Response {
-  repeated Result body = 1;
+  InternedValueTable interned_value_table = 1;
+  repeated Result body = 2;
 }
 
 message Result {
@@ -1261,13 +1272,13 @@ message SeqMapVal {
   SrcPosition src = 2;
 }
 
-// ast.ir:90
+// ast.ir:92
 message SessionResetRequiredError {
   int64 uid = 1;
   VarId var_id = 2;
 }
 
-// ast.ir:79
+// ast.ir:81
 message SfQueryResult {
   string uuid = 1;
 }

src/snowflake/snowpark/types.py

@@ -784,7 +784,7 @@ def json_value(self) -> Dict[str, Any]:
     def _fill_ast(self, ast: proto.SpDataType) -> None:
         ast.sp_struct_type.structured = self.structured
         for field in self.fields:
-            field._fill_ast(ast.sp_struct_type.fields.add())
+            field._fill_ast(ast.sp_struct_type.fields.list.add())
 
 
 class VariantType(DataType):

tests/ast/data/sproc.test

@@ -173,11 +173,11 @@ res33 = sproc("select_sp", return_type=StructType([StructField("A", IntegerType(
 
 session.sql("SELECT 1 as A, 2 as B").show()
 
-res37 = sproc("select_sp", return_type=StructType([], structured=False), input_types=[IntegerType(), IntegerType()], source_code_display=False, _registered_object_name="\"MOCK_DATABASE\".\"MOCK_SCHEMA\".\"SNOWPARK_TEMP_PROCEDURE_xxx\"")(1, 2)
+res37 = sproc("select_sp", return_type=StructType(structured=False), input_types=[IntegerType(), IntegerType()], source_code_display=False, _registered_object_name="\"MOCK_DATABASE\".\"MOCK_SCHEMA\".\"SNOWPARK_TEMP_PROCEDURE_xxx\"")(1, 2)
 
 session.sql("SELECT 1 as A, 2 as B").show()
 
-res41 = sproc("select_sp", return_type=StructType([], structured=False), input_types=[LongType(), LongType()], _registered_object_name="\"MOCK_DATABASE\".\"MOCK_SCHEMA\".\"SNOWPARK_TEMP_PROCEDURE_xxx\"")(1, 2)
+res41 = sproc("select_sp", return_type=StructType(structured=False), input_types=[LongType(), LongType()], _registered_object_name="\"MOCK_DATABASE\".\"MOCK_SCHEMA\".\"SNOWPARK_TEMP_PROCEDURE_xxx\"")(1, 2)
 
 session.sql("SELECT 1 as A, 2 as B").show()
 
@@ -1603,22 +1603,24 @@ body {
         return_type {
           sp_struct_type {
             fields {
-              column_identifier {
-                name: "A"
-              }
-              data_type {
-                sp_integer_type: true
-              }
-              nullable: true
-            }
-            fields {
-              column_identifier {
-                name: "B"
+              list {
+                column_identifier {
+                  name: "A"
+                }
+                data_type {
+                  sp_integer_type: true
+                }
+                nullable: true
               }
-              data_type {
-                sp_integer_type: true
+              list {
+                column_identifier {
+                  name: "B"
+                }
+                data_type {
+                  sp_integer_type: true
+                }
+                nullable: true
               }
-              nullable: true
             }
           }
         }