From 370382affdc98fa3e9c96031eda5c49e0cdc0efd Mon Sep 17 00:00:00 2001 From: May Liu Date: Fri, 17 Apr 2026 14:34:05 -0700 Subject: [PATCH 1/2] Preserve 1-pass read for SCOS XML user schema performance regardless of inferSchema value --- src/snowflake/snowpark/dataframe_reader.py | 9 +++++++- tests/unit/test_xml_schema_inference.py | 27 ++++++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) 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..31dfced54c 100644 --- a/tests/unit/test_xml_schema_inference.py +++ b/tests/unit/test_xml_schema_inference.py @@ -1937,3 +1937,30 @@ 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 +# =========================================================================== + + +def test_xml_skip_inference_option_prevents_inference_call(): + s = mock.MagicMock() + reader = DataFrameReader(s, _emit_ast=False) + + # _XML_SKIP_INFERENCE=False, should call _infer_schema_for_xml + reader._cur_options[_dr_mod.XML_ROW_TAG_STRING] = "row" + skip = reader._cur_options.get("_XML_SKIP_INFERENCE", False) + assert skip is False + assert not reader._user_schema + should_infer = (not reader._user_schema) and (not skip) + assert should_infer is True + + # set _XML_SKIP_INFERENCE=True, should not call _infer_schema_for_xml + reader._cur_options["_XML_SKIP_INFERENCE"] = True + with mock.patch.object(reader, "_infer_schema_for_xml") as mock_infer: + skip = reader._cur_options.get("_XML_SKIP_INFERENCE", False) + should_infer = (not reader._user_schema) and (not skip) + assert skip is True + assert should_infer is False + mock_infer.assert_not_called() From f36a7d4e558983aa7f800a653d8ea68de5e73d9f Mon Sep 17 00:00:00 2001 From: May Liu Date: Mon, 20 Apr 2026 11:03:02 -0700 Subject: [PATCH 2/2] mock _read_semi_structured_file --- tests/unit/test_xml_schema_inference.py | 39 +++++++++++++------------ 1 file changed, 20 insertions(+), 19 deletions(-) diff --git a/tests/unit/test_xml_schema_inference.py b/tests/unit/test_xml_schema_inference.py index 31dfced54c..411927812b 100644 --- a/tests/unit/test_xml_schema_inference.py +++ b/tests/unit/test_xml_schema_inference.py @@ -1944,23 +1944,24 @@ def test_udtf_process_string_types_only(): # =========================================================================== -def test_xml_skip_inference_option_prevents_inference_call(): - s = mock.MagicMock() - reader = DataFrameReader(s, _emit_ast=False) - - # _XML_SKIP_INFERENCE=False, should call _infer_schema_for_xml +@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" - skip = reader._cur_options.get("_XML_SKIP_INFERENCE", False) - assert skip is False - assert not reader._user_schema - should_infer = (not reader._user_schema) and (not skip) - assert should_infer is True - - # set _XML_SKIP_INFERENCE=True, should not call _infer_schema_for_xml - reader._cur_options["_XML_SKIP_INFERENCE"] = True - with mock.patch.object(reader, "_infer_schema_for_xml") as mock_infer: - skip = reader._cur_options.get("_XML_SKIP_INFERENCE", False) - should_infer = (not reader._user_schema) and (not skip) - assert skip is True - assert should_infer is False - mock_infer.assert_not_called() + 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()