SNOW-1853347: Add mechanism to allow changing type strs when printing schema (#2819)

Author: sfc-gh-jrose

src/snowflake/snowpark/dataframe.py

@@ -5655,7 +5655,10 @@ def convert(col: ColumnOrName) -> Expression:
         return exprs
 
     def _format_schema(
-        self, level: Optional[int] = None, translate_columns: Optional[dict] = None
+        self,
+        level: Optional[int] = None,
+        translate_columns: Optional[dict] = None,
+        translate_types: Optional[dict] = None,
     ) -> str:
         def _format_datatype(name, dtype, nullable=None, depth=0):
             if level is not None and depth >= level:
@@ -5669,6 +5672,10 @@ def _format_datatype(name, dtype, nullable=None, depth=0):
             extra_lines = []
             type_str = dtype.__class__.__name__
 
+            translated = None
+            if translate_types:
+                translated = translate_types.get(type_str, type_str)
+
             # Structured Type format their parameters on multiple lines.
             if isinstance(dtype, ArrayType):
                 extra_lines = [
@@ -5695,7 +5702,7 @@ def _format_datatype(name, dtype, nullable=None, depth=0):
 
             return "\n".join(
                 [
-                    f"{prefix} |-- {name}: {type_str}{nullable_str}",
+                    f"{prefix} |-- {name}: {translated or type_str}{nullable_str}",
                 ]
                 + [f"{line}" for line in extra_lines if line]
             )

tests/integ/scala/test_datatype_suite.py

@@ -1077,6 +1077,19 @@ def test_structured_type_print_schema(
         == 'root\n |-- "map": MapType (nullable = True)'
     )
 
+    # Check that column types can be translated
+    assert (
+        df._format_schema(
+            2,
+            translate_types={
+                "MapType": "dict",
+                "StringType": "str",
+                "ArrayType": "list",
+            },
+        )
+        == 'root\n |-- "MAP": dict (nullable = True)\n |   |-- key: str\n |   |-- value: list'
+    )
+
 
 @pytest.mark.skipif(
     "config.getoption('local_testing_mode', default=False)",