Skip to content

changes for fixing sp_struct_type decoding #2942

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sfc-gh-batur merged 1 commit into from
Jan 27, 2025
Merged

changes for fixing sp_struct_type decoding #2942

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 1 addition & 5 deletions src/snowflake/snowpark/_internal/proto/ast.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down Expand Up @@ -178,7 +174,7 @@ message SpStructField {

// sp-type.ir:46
message SpStructType {
List_SpStructField fields = 1;
repeated SpStructField fields = 1;
bool structured = 2;
}

Expand Down
2 changes: 1 addition & 1 deletion src/snowflake/snowpark/types.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.list.add())
field._fill_ast(ast.sp_struct_type.fields.add())


class VariantType(DataType):
Expand Down
34 changes: 16 additions & 18 deletions tests/ast/data/sproc.test
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down Expand Up @@ -1603,24 +1603,22 @@ body {
return_type {
sp_struct_type {
fields {
list {
column_identifier {
name: "A"
}
data_type {
sp_integer_type: true
}
nullable: true
column_identifier {
name: "A"
}
list {
column_identifier {
name: "B"
}
data_type {
sp_integer_type: true
}
nullable: true
data_type {
sp_integer_type: true
}
nullable: true
}
fields {
column_identifier {
name: "B"
}
data_type {
sp_integer_type: true
}
nullable: true
}
}
}
Expand Down
4 changes: 2 additions & 2 deletions tests/ast/decoder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -536,7 +536,7 @@ def decode_data_type_expr(
# The fields can be a list of Expr, a single Expr, or None.
fields = []
if hasattr(data_type_expr.sp_struct_type, "fields"):
for field in data_type_expr.sp_struct_type.fields.list:
for field in data_type_expr.sp_struct_type.fields:
column_identifier = field.column_identifier.name
data_type = self.decode_data_type_expr(field.data_type)
nullable = field.nullable
Expand Down Expand Up @@ -765,7 +765,7 @@ def decode_struct_type_expr(
The decoded object.
"""
struct_field_list = []
for field in sp_struct_type_expr.fields.list:
for field in sp_struct_type_expr.fields:
column_identifier = field.column_identifier.name
datatype = self.decode_data_type_expr(field.data_type)
nullable = field.nullable
Expand Down
Loading