Fix: Use SUPPORTS_INSERT_OVERWRITE and cleanup (#893)

eakmanrq · web-flow · commit f24d075e008d · 2023-05-24T16:11:35.000-07:00
* use SUPPORTS_INSERT_OVERWRITE and cleanup

* feedback

* remove where clause
diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py
@@ -596,13 +596,29 @@ def _insert_overwrite_by_condition(
         where: t.Optional[exp.Condition] = None,
         columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
     ) -> None:
-        if where is None:
-            raise SQLMeshError(
-                "Where condition is required when doing a delete/insert for insert/overwrite"
+        table = exp.to_table(table_name)
+        if self.SUPPORTS_INSERT_OVERWRITE:
+            df = self.try_get_pandas_df(query_or_df)
+            if df is not None:
+                query_or_df = next(
+                    pandas_to_sql(
+                        df,
+                        alias=table.alias_or_name,
+                        columns_to_types=columns_to_types,
+                    )
+                )
+            query = t.cast("Query", query_or_df)
+            self.execute(
+                exp.insert(query, table, columns=list(columns_to_types or []), overwrite=True)
             )
-        with self.transaction():
-            self.delete_from(table_name, where=where)
-            self.insert_append(table_name, query_or_df, columns_to_types=columns_to_types)
+        else:
+            if where is None:
+                raise SQLMeshError(
+                    "Where condition is required when doing a delete/insert for insert/overwrite"
+                )
+            with self.transaction():
+                self.delete_from(table_name, where=where)
+                self.insert_append(table_name, query_or_df, columns_to_types=columns_to_types)
 
     def update_table(
         self,
diff --git a/sqlmesh/core/engine_adapter/spark.py b/sqlmesh/core/engine_adapter/spark.py
@@ -5,7 +5,6 @@
 import pandas as pd
 from sqlglot import exp
 
-from sqlmesh.core.dialect import pandas_to_sql
 from sqlmesh.core.engine_adapter.base import EngineAdapter
 from sqlmesh.core.engine_adapter.shared import (
     DataObject,
@@ -21,7 +20,6 @@
         DF,
         PySparkDataFrame,
         PySparkSession,
-        Query,
         QueryOrDF,
     )
     from sqlmesh.core.model.meta import IntervalUnit
@@ -62,29 +60,11 @@ def _insert_overwrite_by_condition(
         where: t.Optional[exp.Condition] = None,
         columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
     ) -> None:
-        table = exp.to_table(table_name)
-        df = self.try_get_pandas_df(query_or_df)
-        pyspark_df = self.try_get_pyspark_df(query_or_df)
-        if self._use_spark_session and (df is not None or pyspark_df):
-            if df is not None:
-                pyspark_df = self._ensure_pyspark_df(df)
-            assert pyspark_df
-            self._insert_pyspark_df(table_name, pyspark_df, overwrite=True)
+        df = self.try_get_df(query_or_df)
+        if self._use_spark_session and df is not None:
+            self._insert_pyspark_df(table_name, self._ensure_pyspark_df(df), overwrite=True)
         else:
-            if df is not None:
-                query_or_df = next(
-                    pandas_to_sql(
-                        df,
-                        alias=table.alias_or_name,
-                        columns_to_types=columns_to_types,
-                    )
-                )
-            column_names = list(columns_to_types or [])
-            self.execute(
-                exp.insert(
-                    t.cast("Query", query_or_df), table, columns=column_names, overwrite=True
-                )
-            )
+            super()._insert_overwrite_by_condition(table_name, query_or_df, where, columns_to_types)
 
     def insert_append(
         self,
@@ -94,10 +74,10 @@ def insert_append(
         contains_json: bool = False,
     ) -> None:
         df = self.try_get_df(query_or_df)
-        if df is None or not self._use_spark_session:
-            super().insert_append(table_name, query_or_df, columns_to_types, contains_json)
-        else:
+        if self._use_spark_session and df is not None:
             self._insert_append_pyspark_df(table_name, self._ensure_pyspark_df(df))
+        else:
+            super().insert_append(table_name, query_or_df, columns_to_types, contains_json)
 
     def merge(
         self,
@@ -108,16 +88,16 @@ def merge(
     ) -> None:
         column_names = columns_to_types.keys()
         df = self.try_get_df(source_table)
-        if df is None or not self._use_spark_session:
-            super().merge(target_table, source_table, columns_to_types, unique_key)
-        else:
-            df = self._ensure_pyspark_df(df)
+        if self._use_spark_session and df is not None:
+            pyspark_df = self._ensure_pyspark_df(df)
             temp_view_name = self._get_temp_table(target_table, table_only=True).sql(
                 dialect=self.dialect
             )
-            df.createOrReplaceTempView(temp_view_name)
+            pyspark_df.createOrReplaceTempView(temp_view_name)
             query = exp.select(*column_names).from_(temp_view_name)
             super().merge(target_table, query, columns_to_types, unique_key)
+        else:
+            super().merge(target_table, source_table, columns_to_types, unique_key)
 
     def _insert_append_pandas_df(
         self,
@@ -126,10 +106,10 @@ def _insert_append_pandas_df(
         columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
         contains_json: bool = False,
     ) -> None:
-        if not self._use_spark_session:
-            super()._insert_append_pandas_df(table_name, df, columns_to_types, contains_json)
-        else:
+        if self._use_spark_session:
             self._insert_pyspark_df(table_name, self._ensure_pyspark_df(df), overwrite=False)
+        else:
+            super()._insert_append_pandas_df(table_name, df, columns_to_types, contains_json)
 
     def _insert_append_pyspark_df(
         self,
@@ -160,13 +140,13 @@ def _create_table_from_df(
         replace: bool = True,
         **kwargs: t.Any,
     ) -> None:
-        if not self._use_spark_session:
-            super()._create_table_from_df(table_name, df, columns_to_types, exists, replace)
-        else:
+        if self._use_spark_session:
             df = self._ensure_pyspark_df(df)
             if isinstance(table_name, exp.Table):
                 table_name = table_name.sql(dialect=self.dialect)
             df.write.saveAsTable(table_name, mode="overwrite")
+        else:
+            super()._create_table_from_df(table_name, df, columns_to_types, exists, replace)
 
     def _get_data_objects(
         self, schema_name: str, catalog_name: t.Optional[str] = None
diff --git a/tests/core/engine_adapter/test_base.py b/tests/core/engine_adapter/test_base.py
@@ -118,6 +118,45 @@ def test_insert_overwrite_by_time_partition(mocker: MockerFixture):
     )
 
 
+def test_insert_overwrite_by_time_partition_supports_insert_overwrite(mocker: MockerFixture):
+    connection_mock = mocker.NonCallableMock()
+    cursor_mock = mocker.Mock()
+    connection_mock.cursor.return_value = cursor_mock
+
+    adapter = EngineAdapter(lambda: connection_mock, "")  # type: ignore
+    adapter.SUPPORTS_INSERT_OVERWRITE = True
+    adapter._insert_overwrite_by_condition(
+        "test_table",
+        parse_one("SELECT a, b FROM tbl"),
+        where=parse_one("b BETWEEN '2022-01-01' and '2022-01-02'"),
+        columns_to_types={"a": exp.DataType.build("INT"), "b": exp.DataType.build("STRING")},
+    )
+
+    cursor_mock.execute.assert_called_once_with(
+        "INSERT OVERWRITE TABLE test_table (a, b) SELECT a, b FROM tbl"
+    )
+
+
+def test_insert_overwrite_by_time_partition_supports_insert_overwrite_pandas(mocker: MockerFixture):
+    connection_mock = mocker.NonCallableMock()
+    cursor_mock = mocker.Mock()
+    connection_mock.cursor.return_value = cursor_mock
+
+    adapter = EngineAdapter(lambda: connection_mock, "")  # type: ignore
+    adapter.SUPPORTS_INSERT_OVERWRITE = True
+    df = pd.DataFrame({"a": [1, 2], "ds": ["2022-01-01", "2022-01-02"]})
+    adapter._insert_overwrite_by_condition(
+        "test_table",
+        df,
+        where=parse_one("ds BETWEEN '2022-01-01' and '2022-01-02'"),
+        columns_to_types={"a": exp.DataType.build("INT"), "ds": exp.DataType.build("STRING")},
+    )
+
+    cursor_mock.execute.assert_called_once_with(
+        "INSERT OVERWRITE TABLE test_table (a, ds) SELECT CAST(a AS INT) AS a, CAST(ds AS TEXT) AS ds FROM (VALUES (1, '2022-01-01'), (2, '2022-01-02')) AS test_table(a, ds)"
+    )
+
+
 def test_insert_append_query(mocker: MockerFixture):
     connection_mock = mocker.NonCallableMock()
     cursor_mock = mocker.Mock()
