SNOW-2866374: Enforce deterministic column ordering in infer_schema (#4021)

Author: sfc-gh-mayliu

CHANGELOG.md

@@ -10,6 +10,8 @@
 
 #### Improvements
 
+- Improved `DataFrameReader` to return columns in deterministic order when using `INFER_SCHEMA`.
+
 ### Snowpark pandas API Updates
 
 #### New Features

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

@@ -1299,8 +1299,13 @@ 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,
 ) -> str:
+    """
+    Note: Results are ordered by ORDER_ID and COLUMN_NAME for deterministic column ordering.
+    """
     return (
         SELECT
         + STAR
@@ -1328,6 +1333,8 @@ def infer_schema_statement(
         )
         + RIGHT_PARENTHESIS
         + RIGHT_PARENTHESIS
+        + ORDER_BY
+        + "ORDER_ID, COLUMN_NAME"
     )
 
 

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,22 @@ def test_read_json_with_infer_schema(session, mode):
     ]
 
 
+@pytest.mark.skipif(
+    "config.getoption('local_testing_mode', default=False)",
+    reason="Local Testing does not support JSONL format or INFER_SCHEMA.",
+)
+@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/unit/scala/test_utils_suite.py

@@ -345,6 +345,7 @@ def check_zip_files_and_close_stream(input_stream, expected_files):
                 "resources/testCSVspecialFormat.csv",
                 "resources/testJSONspecialFormat.json.gz",
                 "resources/testJson.json",
+                "resources/testJsonDupeColumnID.jsonl",
                 "resources/testJsonNewSchema.json",
                 "resources/testJsonSameSchema.json",
                 "resources/test_all_data_types.parquet",

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")