SNOW-1754876: Fix inconsistent case in StructFields for structured ty…

…pes (#2502)
snowflakedb · Oct 25, 2024 · 8bc778a · 8bc778a
1 parent aaf2cdc
commit 8bc778a
Show file tree

Hide file tree

Showing 4 changed files with 105 additions and 4 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -25,6 +25,7 @@
 
 - Fixed a bug where the automatic cleanup of temporary tables could interfere with the results of async query execution.
 - Fixed a bug in `DataFrame.analytics.time_series_agg` function to handle multiple data points in same sliding interval.
+- Fixed a bug that created inconsistent casing in field names of structured objects in iceberg schemas.
 
 #### Deprecations:
 

diff --git a/src/snowflake/snowpark/_internal/type_utils.py b/src/snowflake/snowpark/_internal/type_utils.py
@@ -34,6 +34,7 @@
 from snowflake.connector.constants import FIELD_ID_TO_NAME
 from snowflake.connector.cursor import ResultMetadata
 from snowflake.connector.options import installed_pandas, pandas
+from snowflake.snowpark._internal.utils import quote_name
 from snowflake.snowpark.types import (
     LTZ,
     NTZ,
@@ -156,7 +157,8 @@ def convert_metadata_to_sp_type(
             return StructType(
                 [
                     StructField(
-                        field.name, convert_metadata_to_sp_type(field, max_string_size)
+                        quote_name(field.name, keep_case=True),
+                        convert_metadata_to_sp_type(field, max_string_size),
                     )
                     for field in metadata.fields
                 ],
@@ -289,7 +291,7 @@ def convert_sp_to_sf_type(datatype: DataType) -> str:
     if isinstance(datatype, StructType):
         if datatype.structured:
             fields = ", ".join(
-                f"{field.name.upper()} {convert_sp_to_sf_type(field.datatype)}"
+                f"{field.name} {convert_sp_to_sf_type(field.datatype)}"
                 for field in datatype.fields
             )
             return f"OBJECT({fields})"

diff --git a/tests/integ/scala/test_datatype_suite.py b/tests/integ/scala/test_datatype_suite.py
@@ -472,6 +472,104 @@ def test_structured_dtypes_iceberg(
         Utils.drop_table(structured_type_session, table_name)
 
 
+@pytest.mark.skipif(
+    "config.getoption('local_testing_mode', default=False)",
+    reason="local testing does not fully support structured types yet.",
+)
+def test_iceberg_nested_fields(
+    structured_type_session, local_testing_mode, structured_type_support
+):
+    if not (
+        structured_type_support
+        and iceberg_supported(structured_type_session, local_testing_mode)
+    ):
+        pytest.skip("Test requires iceberg support and structured type support.")
+
+    table_name = Utils.random_table_name()
+    transformed_table_name = Utils.random_table_name()
+
+    expected_schema = StructType(
+        [
+            StructField(
+                "NESTED_DATA",
+                StructType(
+                    [
+                        StructField('"camelCase"', StringType(), nullable=True),
+                        StructField('"snake_case"', StringType(), nullable=True),
+                        StructField('"PascalCase"', StringType(), nullable=True),
+                        StructField(
+                            '"nested_map"',
+                            MapType(
+                                StringType(),
+                                StructType(
+                                    [
+                                        StructField(
+                                            '"inner_camelCase"',
+                                            StringType(),
+                                            nullable=True,
+                                        ),
+                                        StructField(
+                                            '"inner_snake_case"',
+                                            StringType(),
+                                            nullable=True,
+                                        ),
+                                        StructField(
+                                            '"inner_PascalCase"',
+                                            StringType(),
+                                            nullable=True,
+                                        ),
+                                    ],
+                                    structured=True,
+                                ),
+                                structured=True,
+                            ),
+                            nullable=True,
+                        ),
+                    ],
+                    structured=True,
+                ),
+                nullable=True,
+            )
+        ],
+        structured=False,
+    )
+
+    try:
+        structured_type_session.sql(
+            f"""
+        CREATE OR REPLACE ICEBERG TABLE {table_name} (
+            "NESTED_DATA" OBJECT(
+                camelCase STRING,
+                snake_case STRING,
+                PascalCase STRING,
+                nested_map MAP(
+                    STRING,
+                    OBJECT(
+                        inner_camelCase STRING,
+                        inner_snake_case STRING,
+                        inner_PascalCase STRING
+                    )
+                )
+            )
+        ) EXTERNAL_VOLUME = 'python_connector_iceberg_exvol' CATALOG = 'SNOWFLAKE' BASE_LOCATION = 'python_connector_merge_gate';
+        """
+        ).collect()
+        df = structured_type_session.table(table_name)
+        assert df.schema == expected_schema
+
+        # Round tripping will fail if the inner fields has incorrect names.
+        df.write.mode("overwrite").save_as_table(
+            table_name=transformed_table_name, iceberg_config=ICEBERG_CONFIG
+        )
+        assert (
+            structured_type_session.table(transformed_table_name).schema
+            == expected_schema
+        )
+    finally:
+        Utils.drop_table(structured_type_session, table_name)
+        Utils.drop_table(structured_type_session, transformed_table_name)
+
+
 @pytest.mark.skipif(
     "config.getoption('local_testing_mode', default=False)",
     reason="local testing does not fully support structured types yet.",

diff --git a/tests/integ/test_stored_procedure.py b/tests/integ/test_stored_procedure.py
@@ -387,8 +387,8 @@ def test_stored_procedure_with_structured_returns(
                 "OBJ",
                 StructType(
                     [
-                        StructField("A", StringType(16777216), nullable=True),
-                        StructField("B", DoubleType(), nullable=True),
+                        StructField('"a"', StringType(16777216), nullable=True),
+                        StructField('"b"', DoubleType(), nullable=True),
                     ],
                     structured=True,
                 ),