From f63d181c5da66bf39ad9b76d41df54269a710047 Mon Sep 17 00:00:00 2001 From: JanakiRamChittineni <341781521+raams95@users.noreply.github.com> Date: Tue, 13 Oct 2026 02:00:39 -0500 Subject: [PATCH] fix(spark): drop DataFrame temp views after inserting from them The Spark and Databricks adapters register a temp view for every DataFrame insert but never drop it, so long sessions accumulate views and keep the DataFrames alive on the driver. Drop the view in the source query cleanup, as the other adapters do for their temp tables. Fixes #6125 Signed-off-by: JanakiRamChittineni <341781521+raams95@users.noreply.github.com> --- sqlmesh/core/engine_adapter/databricks.py | 8 +++-- sqlmesh/core/engine_adapter/spark.py | 12 ++++--- tests/core/engine_adapter/test_databricks.py | 38 ++++++++++++++++++++ tests/core/engine_adapter/test_spark.py | 25 +++++++++++++ 4 files changed, 77 insertions(+), 6 deletions(-) diff --git a/sqlmesh/core/engine_adapter/databricks.py b/sqlmesh/core/engine_adapter/databricks.py index 8d47cdfd61..7ee351950b 100644 --- a/sqlmesh/core/engine_adapter/databricks.py +++ b/sqlmesh/core/engine_adapter/databricks.py @@ -258,13 +258,17 @@ def _df_to_source_queries( df, target_columns_to_types, source_columns=source_columns ) + temp_table = self._get_temp_table(target_table or "spark", table_only=True) + def query_factory() -> Query: - temp_table = self._get_temp_table(target_table or "spark", table_only=True) pyspark_df.createOrReplaceTempView(temp_table.sql(dialect=self.dialect)) self._connection_pool.set_attribute("use_spark_engine_adapter", True) return exp.select(*self._select_columns(target_columns_to_types)).from_(temp_table) - return [SourceQuery(query_factory=query_factory)] + def cleanup() -> None: + self.spark.catalog.dropTempView(temp_table.name) + + return [SourceQuery(query_factory=query_factory, cleanup_func=cleanup)] def _fetch_native_df( self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False diff --git a/sqlmesh/core/engine_adapter/spark.py b/sqlmesh/core/engine_adapter/spark.py index 9199aa3bcd..72672bafd7 100644 --- a/sqlmesh/core/engine_adapter/spark.py +++ b/sqlmesh/core/engine_adapter/spark.py @@ -282,14 +282,18 @@ def _df_to_source_queries( source_columns: t.Optional[t.List[str]] = None, ) -> t.List[SourceQuery]: df = self._ensure_pyspark_df(df, target_columns_to_types, source_columns=source_columns) + temp_table = self._get_temp_table(target_table or "spark", table_only=True) + temp_view_name = temp_table.sql(dialect=self.dialect) + temp_table.set("db", "global_temp") def query_factory() -> Query: - temp_table = self._get_temp_table(target_table or "spark", table_only=True) - df.createOrReplaceGlobalTempView(temp_table.sql(dialect=self.dialect)) # type: ignore - temp_table.set("db", "global_temp") + df.createOrReplaceGlobalTempView(temp_view_name) # type: ignore return exp.select(*self._select_columns(target_columns_to_types)).from_(temp_table) - return [SourceQuery(query_factory=query_factory)] + def cleanup() -> None: + self.spark.catalog.dropGlobalTempView(temp_table.name) + + return [SourceQuery(query_factory=query_factory, cleanup_func=cleanup)] def _ensure_pyspark_df( self, diff --git a/tests/core/engine_adapter/test_databricks.py b/tests/core/engine_adapter/test_databricks.py index 65d8e0388f..7e03a1f11e 100644 --- a/tests/core/engine_adapter/test_databricks.py +++ b/tests/core/engine_adapter/test_databricks.py @@ -837,3 +837,41 @@ def test_databricks_connect_cluster_id_not_overridden_by_falsy_serverless( token="mytoken", cluster_id="0123-456789-mycluster", ) + + +def test_df_source_query_drops_temp_view( + mocker: MockFixture, make_mocked_engine_adapter: t.Callable +): + mocker.patch( + "sqlmesh.core.engine_adapter.databricks.DatabricksEngineAdapter.set_current_catalog" + ) + adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog") + mocker.patch.object( + DatabricksEngineAdapter, + "_use_spark_session", + new_callable=mocker.PropertyMock, + return_value=True, + ) + spark = mocker.Mock() + mocker.patch.object( + DatabricksEngineAdapter, "spark", new_callable=mocker.PropertyMock, return_value=spark + ) + pyspark_df = mocker.Mock() + mocker.patch.object(adapter, "_ensure_pyspark_df", return_value=pyspark_df) + mocker.patch.object( + adapter, + "_get_temp_table", + return_value=exp.to_table("`temp_test_abcdefgh`", dialect="databricks"), + ) + + source_queries = adapter._df_to_source_queries( + pyspark_df, {"a": exp.DataType.build("int")}, batch_size=0, target_table="db.test" + ) + assert len(source_queries) == 1 + + with source_queries[0] as query: + pyspark_df.createOrReplaceTempView.assert_called_once_with("`temp_test_abcdefgh`") + assert query.sql(dialect="databricks") == "SELECT `a` FROM `temp_test_abcdefgh`" + spark.catalog.dropTempView.assert_not_called() + + spark.catalog.dropTempView.assert_called_once_with("temp_test_abcdefgh") diff --git a/tests/core/engine_adapter/test_spark.py b/tests/core/engine_adapter/test_spark.py index d7c3127f05..fa257161d3 100644 --- a/tests/core/engine_adapter/test_spark.py +++ b/tests/core/engine_adapter/test_spark.py @@ -1106,3 +1106,28 @@ def test_get_data_object_wap_branch(make_mocked_engine_adapter: t.Callable, mock d.schema_("sqlmesh__test", "catalog"), {"test__test_view__630979748"}, ) + + +def test_df_source_query_drops_global_temp_view( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(SparkEngineAdapter) + pyspark_df = mocker.Mock() + mocker.patch.object(adapter, "_ensure_pyspark_df", return_value=pyspark_df) + mocker.patch.object( + adapter, + "_get_temp_table", + return_value=exp.to_table("`temp_test_abcdefgh`", dialect="spark"), + ) + + source_queries = adapter._df_to_source_queries( + pyspark_df, {"a": exp.DataType.build("int")}, batch_size=0, target_table="db.test" + ) + assert len(source_queries) == 1 + + with source_queries[0] as query: + pyspark_df.createOrReplaceGlobalTempView.assert_called_once_with("`temp_test_abcdefgh`") + assert query.sql(dialect="spark") == "SELECT `a` FROM global_temp.`temp_test_abcdefgh`" + adapter.spark.catalog.dropGlobalTempView.assert_not_called() + + adapter.spark.catalog.dropGlobalTempView.assert_called_once_with("temp_test_abcdefgh")