Skip to content

Commit

Permalink
changes for fixing sp_struct_type decoding (#2942)
Browse files Browse the repository at this point in the history
<!---
Please answer these questions before creating your pull request. Thanks!
--->

1. Which Jira issue is this PR addressing? Make sure that there is an
accompanying issue to your PR.

   <!---
   In this section, please add a Snowflake Jira issue number.

Note that if a corresponding GitHub issue exists, you should still
include
   the Snowflake Jira issue number. For example, for GitHub issue
#1400, you should
   add "SNOW-1335071" here.
    --->

   Fixes SNOW-NNNNNNN

2. Fill out the following pre-review checklist:

- [ ] I am adding a new automated test(s) to verify correctness of my
new code
- [ ] If this test skips Local Testing mode, I'm requesting review from
@snowflakedb/local-testing
   - [ ] I am adding new logging messages
   - [ ] I am adding a new telemetry message
   - [ ] I am adding new credentials
   - [ ] I am adding a new dependency
- [ ] If this is a new feature/behavior, I'm adding the Local Testing
parity changes.
- [ ] I acknowledge that I have ensured my changes to be thread-safe.
Follow the link for more information: [Thread-safe Developer
Guidelines](https://github.com/snowflakedb/snowpark-python/blob/main/CONTRIBUTING.md#thread-safe-development)

3. Please describe how your code solves the related issue.

Please write a short description of how your code change solves the
related issue.
  • Loading branch information
sfc-gh-batur authored Jan 27, 2025
1 parent 6d16961 commit bb8a2d1
Show file tree
Hide file tree
Showing 4 changed files with 20 additions and 26 deletions.
6 changes: 1 addition & 5 deletions src/snowflake/snowpark/_internal/proto/ast.proto
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
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
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
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

0 comments on commit bb8a2d1

Please sign in to comment.