SNOW-3375781: Preserve 1-pass read for SCOS XML user schema performance (#4185)

Author: sfc-gh-mayliu

src/snowflake/snowpark/dataframe_reader.py

@@ -1741,7 +1741,14 @@ def _read_semi_structured_file(self, path: str, format: str) -> DataFrame:
 
         xml_inferred_schema = None
         if format == "XML" and XML_ROW_TAG_STRING in self._cur_options:
-            if context._is_snowpark_connect_compatible_mode and not self._user_schema:
+            # Internal flag set by SCOS to skip the inference pass when a user schema
+            # is already present, maintaining 1-pass reading when user schema is provided.
+            skip_inference = self._cur_options.get("_XML_SKIP_INFERENCE", False)
+            if (
+                context._is_snowpark_connect_compatible_mode
+                and not self._user_schema
+                and not skip_inference
+            ):
                 string_types_only = not self._cur_options.get("INFER_SCHEMA", True)
                 xml_inferred_schema = self._infer_schema_for_xml(
                     path, string_types_only

tests/unit/test_xml_schema_inference.py

@@ -1937,3 +1937,31 @@ def test_udtf_process_string_types_only():
         )[0][0]
     assert "string" in schema_str
     assert "bigint" not in schema_str and "date" not in schema_str
+
+
+# ===========================================================================
+# _XML_SKIP_INFERENCE option
+# ===========================================================================
+
+
+@pytest.mark.parametrize("skip_inference", [True, False])
+def test_xml_skip_inference_option(skip_inference):
+    reader = DataFrameReader(mock.MagicMock(), _emit_ast=False)
+    reader._cur_options[_dr_mod.XML_ROW_TAG_STRING] = "row"
+    if skip_inference:
+        reader._cur_options["_XML_SKIP_INFERENCE"] = True
+
+    with mock.patch.object(
+        reader, "_infer_schema_for_xml", return_value=None
+    ) as mock_infer, mock.patch.object(
+        _dr_mod.context, "_is_snowpark_connect_compatible_mode", True
+    ):
+        try:
+            reader._read_semi_structured_file("@s/f.xml", "XML")
+        except Exception:
+            pass
+
+        if skip_inference:
+            mock_infer.assert_not_called()
+        else:
+            mock_infer.assert_called_once()