fix tests

Author: erindru

sqlmesh/core/engine_adapter/spark.py

@@ -279,10 +279,16 @@ def _ensure_pyspark_df(
     ) -> PySparkDataFrame:
         pyspark_df = self.try_get_pyspark_df(generic_df)
         if pyspark_df:
+            if columns_to_types:
+                # ensure Spark dataframe column order matches columns_to_types
+                pyspark_df = pyspark_df.select(*list(columns_to_types.keys()))
             return pyspark_df
         df = self.try_get_pandas_df(generic_df)
         if df is None:
             raise SQLMeshError("Ensure PySpark DF can only be run on a PySpark or Pandas DataFrame")
+        if columns_to_types:
+            # ensure Pandas dataframe column order matches columns_to_types
+            df = df[list(columns_to_types.keys())]
         kwargs = (
             dict(schema=self.sqlglot_to_spark_types(columns_to_types)) if columns_to_types else {}
         )

tests/core/engine_adapter/integration/test_integration.py

@@ -2735,7 +2735,10 @@ def _use_warehouse_as_state_connection(gateway_name: str, config: Config):
 
 
 def test_python_model_column_order(ctx: TestContext, tmp_path_factory: pytest.TempPathFactory):
-    if ctx.test_type != "df":
+    if ctx.test_type == "pyspark" and ctx.dialect in ("spark", "databricks"):
+        # dont skip
+        pass
+    elif ctx.test_type != "df":
         pytest.skip("python model column order test only needs to be run once per db")
 
     tmp_path = tmp_path_factory.mktemp(f"column_order_{ctx.test_id}")
@@ -2746,8 +2749,35 @@ def test_python_model_column_order(ctx: TestContext, tmp_path_factory: pytest.Te
 
     # note: this model deliberately defines the columns in the @model definition to be in a different order than what
     # is returned by the DataFrame within the model
-    (tmp_path / "models" / "python_model.py").write_text(
-        """
+    model_path = tmp_path / "models" / "python_model.py"
+    if ctx.test_type == "pyspark":
+        # python model that emits a PySpark dataframe
+        model_path.write_text(
+            """
+from pyspark.sql import DataFrame, Row
+import typing as t
+from sqlmesh import ExecutionContext, model
+
+@model(
+    "TEST_SCHEMA.model",
+    columns={
+        "id": "int",
+        "name": "varchar"
+    }
+)
+def execute(
+    context: ExecutionContext,
+    **kwargs: t.Any,
+) -> DataFrame:
+    return context.spark.createDataFrame([
+        Row(name="foo", id=1)
+    ])
+    """.replace("TEST_SCHEMA", test_schema)
+        )
+    else:
+        # python model that emits a Pandas DataFrame
+        model_path.write_text(
+            """
 import pandas as pd
 import typing as t
 from sqlmesh import ExecutionContext, model
@@ -2766,8 +2796,8 @@ def execute(
     return pd.DataFrame([
         {"name": "foo", "id": 1}
     ])
-""".replace("TEST_SCHEMA", test_schema)
-    )
+    """.replace("TEST_SCHEMA", test_schema)
+        )
 
     sqlmesh_ctx = ctx.create_context(path=tmp_path)
 

tests/core/engine_adapter/test_base.py

@@ -966,7 +966,7 @@ def test_merge_upsert(make_mocked_engine_adapter: t.Callable, assert_exp_eq):
 def test_merge_upsert_pandas(make_mocked_engine_adapter: t.Callable):
     adapter = make_mocked_engine_adapter(EngineAdapter)
 
-    df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+    df = pd.DataFrame({"id": [1, 2, 3], "ts": [4, 5, 6], "val": [1, 2, 3]})
     adapter.merge(
         target_table="target",
         source_table=df,
@@ -978,7 +978,7 @@ def test_merge_upsert_pandas(make_mocked_engine_adapter: t.Callable):
         unique_key=[exp.to_identifier("id")],
     )
     adapter.cursor.execute.assert_called_once_with(
-        'MERGE INTO "target" AS "__MERGE_TARGET__" USING (SELECT CAST("id" AS INT) AS "id", CAST("ts" AS TIMESTAMP) AS "ts", CAST("val" AS INT) AS "val" FROM (VALUES (1, 4), (2, 5), (3, 6)) AS "t"("id", "ts", "val")) AS "__MERGE_SOURCE__" ON "__MERGE_TARGET__"."id" = "__MERGE_SOURCE__"."id" '
+        'MERGE INTO "target" AS "__MERGE_TARGET__" USING (SELECT CAST("id" AS INT) AS "id", CAST("ts" AS TIMESTAMP) AS "ts", CAST("val" AS INT) AS "val" FROM (VALUES (1, 4, 1), (2, 5, 2), (3, 6, 3)) AS "t"("id", "ts", "val")) AS "__MERGE_SOURCE__" ON "__MERGE_TARGET__"."id" = "__MERGE_SOURCE__"."id" '
         'WHEN MATCHED THEN UPDATE SET "__MERGE_TARGET__"."id" = "__MERGE_SOURCE__"."id", "__MERGE_TARGET__"."ts" = "__MERGE_SOURCE__"."ts", "__MERGE_TARGET__"."val" = "__MERGE_SOURCE__"."val" '
         'WHEN NOT MATCHED THEN INSERT ("id", "ts", "val") VALUES ("__MERGE_SOURCE__"."id", "__MERGE_SOURCE__"."ts", "__MERGE_SOURCE__"."val")'
     )
@@ -995,7 +995,7 @@ def test_merge_upsert_pandas(make_mocked_engine_adapter: t.Callable):
         unique_key=[exp.to_identifier("id"), exp.to_identifier("ts")],
     )
     adapter.cursor.execute.assert_called_once_with(
-        'MERGE INTO "target" AS "__MERGE_TARGET__" USING (SELECT CAST("id" AS INT) AS "id", CAST("ts" AS TIMESTAMP) AS "ts", CAST("val" AS INT) AS "val" FROM (VALUES (1, 4), (2, 5), (3, 6)) AS "t"("id", "ts", "val")) AS "__MERGE_SOURCE__" ON "__MERGE_TARGET__"."id" = "__MERGE_SOURCE__"."id" AND "__MERGE_TARGET__"."ts" = "__MERGE_SOURCE__"."ts" '
+        'MERGE INTO "target" AS "__MERGE_TARGET__" USING (SELECT CAST("id" AS INT) AS "id", CAST("ts" AS TIMESTAMP) AS "ts", CAST("val" AS INT) AS "val" FROM (VALUES (1, 4, 1), (2, 5, 2), (3, 6, 3)) AS "t"("id", "ts", "val")) AS "__MERGE_SOURCE__" ON "__MERGE_TARGET__"."id" = "__MERGE_SOURCE__"."id" AND "__MERGE_TARGET__"."ts" = "__MERGE_SOURCE__"."ts" '
         'WHEN MATCHED THEN UPDATE SET "__MERGE_TARGET__"."id" = "__MERGE_SOURCE__"."id", "__MERGE_TARGET__"."ts" = "__MERGE_SOURCE__"."ts", "__MERGE_TARGET__"."val" = "__MERGE_SOURCE__"."val" '
         'WHEN NOT MATCHED THEN INSERT ("id", "ts", "val") VALUES ("__MERGE_SOURCE__"."id", "__MERGE_SOURCE__"."ts", "__MERGE_SOURCE__"."val")'
     )
@@ -1175,23 +1175,23 @@ def test_merge_filter(make_mocked_engine_adapter: t.Callable, assert_exp_eq):
         """
 MERGE INTO "target" AS "__MERGE_TARGET__"
 USING (
-    SELECT "ID", "ts", "val" 
+    SELECT "ID", "ts", "val"
     FROM "source"
 ) AS "__MERGE_SOURCE__"
 ON (
-    "__MERGE_SOURCE__"."ID" > 0 
+    "__MERGE_SOURCE__"."ID" > 0
     AND "__MERGE_TARGET__"."ts" < TIMESTAMP("2020-02-05")
 )
 AND "__MERGE_TARGET__"."ID" = "__MERGE_SOURCE__"."ID"
-WHEN MATCHED THEN 
-    UPDATE SET 
+WHEN MATCHED THEN
+    UPDATE SET
         "__MERGE_TARGET__"."val" = "__MERGE_SOURCE__"."val",
         "__MERGE_TARGET__"."ts" = COALESCE("__MERGE_SOURCE__"."ts", "__MERGE_TARGET__"."ts")
-WHEN NOT MATCHED THEN 
-    INSERT ("ID", "ts", "val") 
+WHEN NOT MATCHED THEN
+    INSERT ("ID", "ts", "val")
     VALUES (
-        "__MERGE_SOURCE__"."ID", 
-        "__MERGE_SOURCE__"."ts", 
+        "__MERGE_SOURCE__"."ID",
+        "__MERGE_SOURCE__"."ts",
         "__MERGE_SOURCE__"."val"
     );
 """,
@@ -1585,7 +1585,11 @@ def test_merge_scd_type_2_pandas(make_mocked_engine_adapter: t.Callable):
             "id2": [4, 5, 6],
             "name": ["muffins", "chips", "soda"],
             "price": [4.0, 5.0, 6.0],
-            "updated_at": ["2020-01-01 10:00:00", "2020-01-02 15:00:00", "2020-01-03 12:00:00"],
+            "test_updated_at": [
+                "2020-01-01 10:00:00",
+                "2020-01-02 15:00:00",
+                "2020-01-03 12:00:00",
+            ],
         }
     )
     adapter.scd_type_2_by_time(