diff --git a/src/snowflake/snowpark/dataframe_reader.py b/src/snowflake/snowpark/dataframe_reader.py index 09ca05a30a..26f6ba753f 100644 --- a/src/snowflake/snowpark/dataframe_reader.py +++ b/src/snowflake/snowpark/dataframe_reader.py @@ -1703,7 +1703,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 diff --git a/tests/unit/test_xml_schema_inference.py b/tests/unit/test_xml_schema_inference.py index b047be37d0..411927812b 100644 --- a/tests/unit/test_xml_schema_inference.py +++ b/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()