SNOW-3215485: Fix colon escape in XML custom schema (#4114)

Author: sfc-gh-mayliu

CHANGELOG.md

@@ -11,6 +11,7 @@
 - Fixed a bug in `Session.client_telemetry` that trace does not have snowflake style trace id.
 - Fixed a bug when saving a fdn table into an iceberg table in overwrite mode, error is raised because `StringType` is saved in wrong length.
 - Fixed a bug in `ai_complete` where `model_parameters` and `response_format` values containing single quotes would generate malformed SQL.
+- Fixed a bug in `DataFrameReader.xml()` where reading XML with a custom schema whose field names contain colons (e.g., `px:name`) raised a `SnowparkColumnException`.
 
 #### Improvements
 

src/snowflake/snowpark/_internal/type_utils.py

@@ -1250,15 +1250,21 @@ def find_top_level_colon(field_def: str) -> int:
     """
     Returns the index of the first top-level colon in 'field_def',
     or -1 if there is no top-level colon. A colon is considered top-level
-    if it is not enclosed in <...> or (...).
+    if it is not enclosed in <...>, (...), or "...".
 
     Example:
-      'a struct<i: integer>' => returns -1 (colon is nested).
+      'a struct<i: integer>'  => returns -1 (colon is nested).
       'x: struct<i: integer>' => returns index of the colon after 'x'.
+      '"px:name": string'     => returns index of the colon after the closing '"'.
     """
     bracket_depth = 0
+    in_quotes = False
     for i, ch in enumerate(field_def):
-        if ch in ("<", "("):
+        if ch == '"':
+            in_quotes = not in_quotes
+        elif in_quotes:
+            continue
+        elif ch in ("<", "("):
             bracket_depth += 1
         elif ch in (">", ")"):
             bracket_depth -= 1

src/snowflake/snowpark/_internal/xml_reader.py

@@ -57,12 +57,58 @@ def replace_entity(match: re.Match) -> str:
         return match.group(0)
 
 
+# TODO SNOW-3217320: The escape/restore below works around a bug in
+# type_utils.find_top_level_colon which mis-splits colons inside double-quoted
+# identifiers. The fix for find_top_level_colon will ship in this same release
+# (type_utils.py), but the UDTF resolves snowflake-snowpark-python from
+# Snowflake's Anaconda channel at runtime, so the server may still run an older
+# version. Once the Anaconda-channel version includes the fix, this
+# escape/restore logic becomes a harmless no-op and can be removed.
+_COLON_PLACEHOLDER = "\x00COLON\x00"
+
+
+def _escape_colons_in_quotes(schema_str: str) -> str:
+    """Replace colons inside double-quoted identifiers with a placeholder.
+
+    The server-side type_utils.find_top_level_colon does not skip colons inside
+    double-quoted identifiers, so a field like ``"px:name": string`` gets split at
+    the wrong colon.  By replacing those colons with a placeholder, we let the existing
+    parser find the correct top-level colon that separates the field name from its type.
+    """
+    result = []
+    in_quotes = False
+    for ch in schema_str:
+        if ch == '"':
+            in_quotes = not in_quotes
+            result.append(ch)
+        elif ch == ":" and in_quotes:
+            result.append(_COLON_PLACEHOLDER)
+        else:
+            result.append(ch)
+    return "".join(result)
+
+
+def _restore_colons_in_template(template: Optional[dict]) -> Optional[dict]:
+    """Recursively restore colon placeholders back to real colons in template keys."""
+    if template is None:
+        return None
+    restored: Dict[str, Any] = {}
+    for key, value in template.items():
+        real_key = key.replace(_COLON_PLACEHOLDER, ":")
+        if isinstance(value, dict):
+            restored[real_key] = _restore_colons_in_template(value)
+        else:
+            restored[real_key] = value
+    return restored
+
+
 def schema_string_to_result_dict_and_struct_type(schema_string: str) -> Optional[dict]:
     if schema_string == "":
         return None
-    schema = type_string_to_type_object(schema_string)
-
-    return struct_type_to_result_template(schema)
+    safe_string = _escape_colons_in_quotes(schema_string)
+    schema = type_string_to_type_object(safe_string)
+    template = struct_type_to_result_template(schema)
+    return _restore_colons_in_template(template)
 
 
 def struct_type_to_result_template(dt: DataType) -> Optional[dict]:

tests/integ/test_xml_reader_row_tag.py

@@ -693,6 +693,28 @@ def test_user_schema_without_rowtag(session):
         session.read.schema(user_schema).xml(f"@{tmp_stage_name}/{test_file_books_xml}")
 
 
+@pytest.mark.parametrize("ignore_namespace", [True, False])
+def test_read_xml_custom_schema_with_colon_tags(session, ignore_namespace):
+    schema = StructType(
+        [
+            StructField("px:name", StringType(), True),
+            StructField("px:value", StringType(), True),
+        ]
+    )
+    df = (
+        session.read.option("rowTag", "px:item")
+        .schema(schema)
+        .option("ignoreNamespace", ignore_namespace)
+        .xml(f"@{tmp_stage_name}/{test_file_xml_undeclared_namespace}")
+    )
+    result = df.collect()
+    assert len(result) == 2
+    names = {r[0] for r in result}
+    values = {r[1] for r in result}
+    assert names == {"Item One", "Item Two"}
+    assert values == {"100", "200"}
+
+
 def test_value_tag_custom_schema(session):
     test_schema1 = StructType(
         [

tests/unit/test_types.py

@@ -51,6 +51,7 @@
     snow_type_to_dtype_str,
     split_top_level_comma_fields,
     type_string_to_type_object,
+    find_top_level_colon,
 )
 from snowflake.snowpark.types import (
     ArrayType,
@@ -2539,3 +2540,12 @@ def process(self, input_a, input_b):
         "", "process", class_name="MyOtherClass", _source=source
     )
     assert arg_names is None
+
+
+def test_find_top_level_colon():
+    assert find_top_level_colon("x: int") == 1
+    assert find_top_level_colon("a struct<i: integer>") == -1
+    assert find_top_level_colon("x: struct<i: integer>") == 1
+    assert find_top_level_colon('"px:name": string') == 9
+    assert find_top_level_colon('"a:b:c": string') == 7
+    assert find_top_level_colon("plain_field") == -1

tests/unit/test_xml_reader.py

@@ -25,6 +25,9 @@
     DEFAULT_CHUNK_SIZE,
     struct_type_to_result_template,
     schema_string_to_result_dict_and_struct_type,
+    _escape_colons_in_quotes,
+    _restore_colons_in_template,
+    _COLON_PLACEHOLDER,
 )
 from snowflake.snowpark.types import (
     StructType,
@@ -991,3 +994,85 @@ def test_user_schema_value_tag():
     res = element_to_dict_or_str(element, result_template=result_template)
     assert result_template == {"NUM": None, "STR1": None, "STR2": None, "STR3": None}
     assert res == {"NUM": "1", "STR1": "NULL", "STR2": None, "STR3": "xxx"}
+
+
+def test_escape_colons_in_quotes():
+    s = '"px:name": string'
+    escaped = _escape_colons_in_quotes(s)
+    assert escaped.count(":") == 1
+    assert _COLON_PLACEHOLDER in escaped
+
+    s = '"px:name": string, "px:value": string'
+    escaped = _escape_colons_in_quotes(s)
+    assert escaped.count(":") == 2
+    assert escaped.count(_COLON_PLACEHOLDER) == 2
+
+    s = 'struct<"px:name": string, "detail": struct<"px:sub": string>>'
+    escaped = _escape_colons_in_quotes(s)
+    assert escaped.count(_COLON_PLACEHOLDER) == 2
+    assert escaped.count(":") == 3
+
+    s = '"Author": string, "Title": string'
+    assert _escape_colons_in_quotes(s) == s
+    assert _escape_colons_in_quotes("") == ""
+    assert _escape_colons_in_quotes("a: int, b: string") == "a: int, b: string"
+
+
+def test_restore_colons_in_template():
+    assert _restore_colons_in_template(None) is None
+    assert _restore_colons_in_template({"Author": None}) == {"Author": None}
+
+    # Flat
+    template = {
+        f"px{_COLON_PLACEHOLDER}name": None,
+        f"px{_COLON_PLACEHOLDER}value": None,
+    }
+    assert _restore_colons_in_template(template) == {
+        "px:name": None,
+        "px:value": None,
+    }
+
+    # Nested
+    template = {
+        f"eq{_COLON_PLACEHOLDER}event": {
+            f"eq{_COLON_PLACEHOLDER}sub-id": None,
+        },
+        "plain": None,
+    }
+    assert _restore_colons_in_template(template) == {
+        "eq:event": {"eq:sub-id": None},
+        "plain": None,
+    }
+
+
+def test_schema_string_round_trip_with_colons():
+    # Flat
+    schema_str = 'struct<"px:name": string, "px:value": string>'
+    assert schema_string_to_result_dict_and_struct_type(schema_str) == {
+        "px:name": None,
+        "px:value": None,
+    }
+
+    # Nested
+    schema_str = (
+        'struct<"eq:event-id": string,' '"eq:detail": struct<"eq:sub-id": string>>'
+    )
+    assert schema_string_to_result_dict_and_struct_type(schema_str) == {
+        "eq:event-id": None,
+        "eq:detail": {"eq:sub-id": None},
+    }
+
+    # Mixed
+    schema_str = 'struct<"px:name": string, "Title": string, price: double>'
+    assert schema_string_to_result_dict_and_struct_type(schema_str) == {
+        "px:name": None,
+        "Title": None,
+        "PRICE": None,
+    }
+
+    # No colon fields
+    schema_str = 'struct<"Author": string, "TITLE": string>'
+    assert schema_string_to_result_dict_and_struct_type(schema_str) == {
+        "Author": None,
+        "TITLE": None,
+    }