NO-SNOW: order databricks test data to address flakiness (#3390)

Author: sfc-gh-aling

tests/integ/datasource/test_databricks.py

@@ -64,13 +64,15 @@ def test_basic_databricks(session, input_type, input_value):
     input_dict = {
         input_type: input_value,
     }
-    df = session.read.dbapi(create_databricks_connection, **input_dict)
+    df = session.read.dbapi(create_databricks_connection, **input_dict).order_by(
+        "COL_BYTE"
+    )
     ret = df.collect()
     assert ret == EXPECTED_TEST_DATA and df.schema == EXPECTED_TYPE
 
     table_name = random_name_for_temp_object(TempObjectType.TABLE)
     df.write.save_as_table(table_name, mode="overwrite", table_type="temp")
-    df2 = session.table(table_name)
+    df2 = session.table(table_name).order_by("COL_BYTE")
     assert df2.collect() == EXPECTED_TEST_DATA and df2.schema == EXPECTED_TYPE
 
 
@@ -144,7 +146,7 @@ def local_create_databricks_connection():
         udtf_configs={
             "external_access_integration": DATABRICKS_TEST_EXTERNAL_ACCESS_INTEGRATION
         },
-    )
+    ).order_by("COL_BYTE")
     ret = df.collect()
     assert ret == EXPECTED_TEST_DATA and df.schema == EXPECTED_TYPE
 
@@ -160,10 +162,17 @@ def test_unit_udtf_ingestion():
     udtf_ingestion_instance = udtf_ingestion_class()
 
     dsp = DataSourcePartitioner(
-        create_databricks_connection, TEST_TABLE_NAME, is_query=False
+        create_databricks_connection,
+        f"(select * from {TEST_TABLE_NAME}) SORT BY COL_BYTE NULLS FIRST",
+        is_query=True,
     )
-    yield_data = udtf_ingestion_instance.process(dsp.partitions[0])
-    for row, expected_row in zip(yield_data, EXPECTED_TEST_DATA):
+    yield_data = list(udtf_ingestion_instance.process(dsp.partitions[0]))
+    # databricks sort by returns the all None row as the last row regardless of NULLS FIRST/LAST
+    # while in snowflake test data after default sort None is the first row
+    assert yield_data[-1] == EXPECTED_TEST_DATA[0]
+    for row, expected_row in zip(
+        yield_data[:-1], EXPECTED_TEST_DATA[1:]
+    ):  # None data ordering is the same
         for index, (field, value) in enumerate(zip(EXPECTED_TYPE.fields, row)):
             if isinstance(field.datatype, VariantType):
                 # Convert ArrayType, MapType, and StructType to JSON

tests/resources/test_data_source_dir/test_databricks_data.py

@@ -25,28 +25,7 @@
 TEST_TABLE_NAME = "ALL_TYPE_TABLE_2"  # ALL_TYPE_TABLE_2 contains None data while ALL_TYPE_TABLE doesn't
 TZ_INFO = pytz.timezone("America/Los_Angeles")
 EXPECTED_TEST_DATA = [
-    (
-        -34,
-        25393,
-        35234,
-        5644171805,
-        18.264881134033203,
-        9187.446999674603,
-        Decimal("269.89"),
-        "str_8541",
-        True,
-        bytearray(b"\xad\xa9\xdd\xa2"),
-        datetime.date(2025, 6, 8),
-        TZ_INFO.localize(
-            datetime.datetime(2025, 4, 16, 10, 39, 39, 565000), is_dst=True
-        ),
-        datetime.datetime(2025, 4, 16, 17, 49, 8, 565000),
-        "[\n  82,\n  40\n]",
-        '{\n  "key1": 71,\n  "key2": 81\n}',
-        '{\n  "field1": "f_25",\n  "field2": 25\n}',
-        "3-10",
-        "18 14:29:08.000000000",
-    ),
+    tuple([None] * 18),
     (
         -113,
         -14623,
@@ -92,26 +71,26 @@
         "2 11:12:05.000000000",
     ),
     (
-        114,
-        11139,
-        75014,
-        1135763646,
-        14.668656349182129,
-        1378.8325065107654,
-        Decimal("7411.91"),
-        "str_9765",
-        False,
+        -34,
+        25393,
+        35234,
+        5644171805,
+        18.264881134033203,
+        9187.446999674603,
+        Decimal("269.89"),
+        "str_8541",
+        True,
         bytearray(b"\xad\xa9\xdd\xa2"),
-        datetime.date(2025, 6, 29),
+        datetime.date(2025, 6, 8),
         TZ_INFO.localize(
-            datetime.datetime(2025, 4, 16, 10, 48, 27, 565000), is_dst=True
+            datetime.datetime(2025, 4, 16, 10, 39, 39, 565000), is_dst=True
         ),
-        datetime.datetime(2025, 4, 16, 17, 50, 8, 565000),
-        "[\n  92,\n  27\n]",
-        '{\n  "key1": 52,\n  "key2": 65\n}',
-        '{\n  "field1": "f_85",\n  "field2": 50\n}',
-        "7-4",
-        "22 04:52:41.000000000",
+        datetime.datetime(2025, 4, 16, 17, 49, 8, 565000),
+        "[\n  82,\n  40\n]",
+        '{\n  "key1": 71,\n  "key2": 81\n}',
+        '{\n  "field1": "f_25",\n  "field2": 25\n}',
+        "3-10",
+        "18 14:29:08.000000000",
     ),
     (
         -31,
@@ -135,7 +114,28 @@
         "0-7",
         "19 06:25:08.000000000",
     ),
-    tuple([None] * 18),
+    (
+        114,
+        11139,
+        75014,
+        1135763646,
+        14.668656349182129,
+        1378.8325065107654,
+        Decimal("7411.91"),
+        "str_9765",
+        False,
+        bytearray(b"\xad\xa9\xdd\xa2"),
+        datetime.date(2025, 6, 29),
+        TZ_INFO.localize(
+            datetime.datetime(2025, 4, 16, 10, 48, 27, 565000), is_dst=True
+        ),
+        datetime.datetime(2025, 4, 16, 17, 50, 8, 565000),
+        "[\n  92,\n  27\n]",
+        '{\n  "key1": 52,\n  "key2": 65\n}',
+        '{\n  "field1": "f_85",\n  "field2": 50\n}',
+        "7-4",
+        "22 04:52:41.000000000",
+    ),
 ]
 EXPECTED_TYPE = StructType(
     [