Test partial column order should stay the same

Author: treysp

sqlmesh/core/test/definition.py

@@ -648,12 +648,12 @@ def _create_df(
             referenced_columns = list(dict.fromkeys(col for row in rows for col in row))
             _raise_if_unexpected_columns(columns, referenced_columns)
 
+            columns = [str(c) for c in columns]
             if partial:
-                columns = referenced_columns
+                # retain order of columns
+                columns = [c for c in columns if c in referenced_columns]
 
-        return pd.DataFrame.from_records(
-            rows, columns=[str(c) for c in columns] if columns else None
-        )
+        return pd.DataFrame.from_records(rows, columns=columns if columns else None)
 
     def _add_missing_columns(
         self, query: exp.Query, all_columns: t.Optional[t.Collection[str]] = None

tests/core/test_test.py

@@ -770,6 +770,42 @@ def test_partial_data_column_order(sushi_context: Context) -> None:
         ).run()
     )
 
+    #  - output df must differ if sorted by (id, event_date) vs. (event_date, id)
+    #  - output partial must be true
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
+test_foo:
+  model: sushi.foo
+  inputs:
+    sushi.items:
+      - id: 9876
+        event_date: 2020-01-01
+      - id: 1234
+        name: hello
+        event_date: 2020-01-02
+  outputs:
+    partial: true
+    query:
+      - event_date: 2020-01-01
+        id: 9876
+      - event_date: 2020-01-02
+        id: 1234
+        name: hello
+                """
+            ),
+            test_name="test_foo",
+            model=sushi_context.upsert_model(
+                _create_model(
+                    "SELECT id, name, price, event_date FROM sushi.items",
+                    default_catalog=sushi_context.default_catalog,
+                )
+            ),
+            context=sushi_context,
+        ).run()
+    )
+
 
 def test_partial_data_missing_schemas(sushi_context: Context) -> None:
     _check_successful_or_raise(