Enforce deterministic column ordering from infer_schema

Author: sfc-gh-mayliu

src/snowflake/snowpark/_internal/analyzer/analyzer_utils.py

@@ -1299,9 +1299,16 @@ def create_file_format_statement(
 
 
 def infer_schema_statement(
-    path: str, file_format_name: str, options: Optional[Dict[str, str]] = None
+    path: str,
+    file_format_name: str,
+    options: Optional[Dict[str, str]] = None,
+    ordered: bool = True,
 ) -> str:
-    return (
+    """
+    Note: If ordered is set to True (default), results are ordered by ORDER_ID and
+    COLUMN_NAME for deterministic column ordering. Set to False to disable ordering.
+    """
+    base_query = (
         SELECT
         + STAR
         + FROM
@@ -1330,6 +1337,11 @@ def infer_schema_statement(
         + RIGHT_PARENTHESIS
     )
 
+    if ordered:
+        base_query += ORDER_BY + "ORDER_ID, COLUMN_NAME"
+
+    return base_query
+
 
 def file_operation_statement(
     command: str, file_name: str, stage_location: str, options: Dict[str, str]

tests/integ/conftest.py

@@ -237,6 +237,7 @@ def connection(db_parameters, local_testing_mode):
             "user",
             "password",
             "private_key_file",
+            "private_key_file_pwd",
             "host",
             "port",
             "database",

tests/integ/scala/test_dataframe_reader_suite.py

@@ -67,6 +67,7 @@
 test_file_json = "testJson.json"
 test_file_json_same_schema = "testJsonSameSchema.json"
 test_file_json_new_schema = "testJsonNewSchema.json"
+test_file_json_dupe_column_id = "testJsonDupeColumnID.jsonl"
 test_file_avro = "test.avro"
 test_file_parquet = "test.parquet"
 test_file_all_data_types_parquet = "test_all_data_types.parquet"
@@ -194,6 +195,12 @@ def setup(session, resources_path, local_testing_mode):
         test_files.test_file_json,
         compress=False,
     )
+    Utils.upload_to_stage(
+        session,
+        "@" + tmp_stage_name1,
+        test_files.test_file_json_dupe_column_id,
+        compress=False,
+    )
     Utils.upload_to_stage(
         session,
         "@" + tmp_stage_only_json_file,
@@ -1255,6 +1262,18 @@ def test_read_json_with_infer_schema(session, mode):
     ]
 
 
+@pytest.mark.parametrize("mode", ["select", "copy"])
+def test_read_json_with_infer_schema_deterministic_column_order(session, mode):
+    """Test that INFER_SCHEMA returns columns in deterministic order."""
+    json_path = f"@{tmp_stage_name1}/{test_file_json_dupe_column_id}"
+
+    # Run multiple times to verify deterministic ordering. Previously, column order would vary between runs:
+    # (a, b, c) vs (a, c, b) due to order_id of column b and c being the same.
+    for _ in range(10):
+        df = get_reader(session, mode).option("INFER_SCHEMA", True).json(json_path)
+        assert df.columns == ['"a"', '"b"', '"c"']
+
+
 @pytest.mark.skipif(
     "config.getoption('local_testing_mode', default=False)",
     reason="Local Testing does not support loading json with user specified schema.",

tests/resources/testJsonDupeColumnID.jsonl

@@ -0,0 +1,3 @@
+{"a":3.0, "b":3.5}
+{"a":4.0, "c":3.0}
+{"a":6, "c":7}

tests/utils.py

@@ -1541,6 +1541,10 @@ def test_file_json_same_schema(self):
     def test_file_json_new_schema(self):
         return os.path.join(self.resources_path, "testJsonNewSchema.json")
 
+    @property
+    def test_file_json_dupe_column_id(self):
+        return os.path.join(self.resources_path, "testJsonDupeColumnID.jsonl")
+
     @property
     def test_file_avro(self):
         return os.path.join(self.resources_path, "test.avro")