[SPARK-50569][CONNECT]: clean up data cached by individual sessions

python/pyspark/sql/catalog.py

@@ -1409,15 +1409,24 @@ def uncacheTable(self, tableName: str) -> None:
         """
         self._jcatalog.uncacheTable(tableName)
 
-    def clearCache(self) -> None:
-        """Removes all cached tables from the in-memory cache.
+    def clearCache(self, allSessions: bool = True) -> None:
+        """Removes cached tables from the in-memory cache.
 
         .. versionadded:: 2.0.0
 
+        Parameters
+        ----------
+        allSessions : bool, optional
+            Whether to clear cached data across all sessions. If ``False``, only data cached by
+            the current session is cleared, while data also cached by another session is preserved.
+            Defaults to ``True``.
+
+            .. versionadded:: 4.4.0
+
         Notes
         -----
         Cached data is shared across all Spark sessions on the cluster, so clearing
-        the cache affects all sessions.
+        the cache affects all sessions by default.
 
         Examples
         --------
@@ -1428,7 +1437,7 @@ def clearCache(self) -> None:
         False
         >>> _ = spark.sql("DROP TABLE tbl1")
         """
-        self._jcatalog.clearCache()
+        self._jcatalog.clearCache(allSessions)
 
     def refreshTable(self, tableName: str) -> None:
         """Invalidates and refreshes all the cached data and metadata of the given table.

python/pyspark/sql/connect/catalog.py

@@ -302,8 +302,8 @@ def uncacheTable(self, tableName: str) -> None:
 
     uncacheTable.__doc__ = PySparkCatalog.uncacheTable.__doc__
 
-    def clearCache(self) -> None:
-        self._execute_and_fetch(plan.ClearCache())
+    def clearCache(self, allSessions: bool = True) -> None:
+        self._execute_and_fetch(plan.ClearCache(all_sessions=allSessions))
 
     clearCache.__doc__ = PySparkCatalog.clearCache.__doc__
 

python/pyspark/sql/connect/plan.py

@@ -2622,12 +2622,13 @@ def plan(self, session: "SparkConnectClient") -> proto.Relation:
 
 
 class ClearCache(LogicalPlan):
-    def __init__(self) -> None:
+    def __init__(self, all_sessions: bool = True) -> None:
         super().__init__(None)
+        self._all_sessions = all_sessions
 
     def plan(self, session: "SparkConnectClient") -> proto.Relation:
         plan = self._create_proto_relation()
-        plan.catalog.clear_cache.SetInParent()
+        plan.catalog.clear_cache.all_sessions = self._all_sessions
         return plan
 
 

python/pyspark/sql/connect/proto/catalog_pb2.pyi

@@ -1126,9 +1126,29 @@ class ClearCache(google.protobuf.message.Message):
 
     DESCRIPTOR: google.protobuf.descriptor.Descriptor
 
+    ALL_SESSIONS_FIELD_NUMBER: builtins.int
+    all_sessions: builtins.bool
+    """(Optional) Whether to clear cached data across all sessions. Defaults to true when omitted."""
     def __init__(
         self,
+        *,
+        all_sessions: builtins.bool | None = ...,
+    ) -> None: ...
+    def HasField(
+        self,
+        field_name: typing_extensions.Literal[
+            "_all_sessions", b"_all_sessions", "all_sessions", b"all_sessions"
+        ],
+    ) -> builtins.bool: ...
+    def ClearField(
+        self,
+        field_name: typing_extensions.Literal[
+            "_all_sessions", b"_all_sessions", "all_sessions", b"all_sessions"
+        ],
     ) -> None: ...
+    def WhichOneof(
+        self, oneof_group: typing_extensions.Literal["_all_sessions", b"_all_sessions"]
+    ) -> typing_extensions.Literal["all_sessions"] | None: ...
 
 global___ClearCache = ClearCache
 

python/pyspark/sql/tests/test_catalog.py

@@ -18,6 +18,7 @@
 
 from pyspark import StorageLevel
 from pyspark.errors import AnalysisException, PySparkTypeError
+from pyspark.sql import is_remote
 from pyspark.sql.types import IntegerType, StructField, StructType
 from pyspark.testing.sqlutils import ReusedSQLTestCase
 
@@ -449,6 +450,31 @@ def assert_cached(c: bool):
                 spark.catalog.clearCache()
                 assert_cached(False)
 
+    def test_clear_cache_current_session(self):
+        # SPARK-50569: clearCache can preserve data cached by other sessions.
+        spark = self.spark
+        other_session = spark.newSession()
+        first_view = "clear_cache_first_view"
+        second_view = "clear_cache_second_view"
+        try:
+            spark.range(1).createTempView(first_view)
+            other_session.range(1, 2).createTempView(second_view)
+            spark.catalog.cacheTable(first_view)
+            other_session.catalog.cacheTable(second_view)
+
+            spark.catalog.clearCache(allSessions=False)
+
+            self.assertFalse(spark.catalog.isCached(first_view))
+            self.assertTrue(other_session.catalog.isCached(second_view))
+            other_session.catalog.clearCache(allSessions=False)
+            self.assertFalse(other_session.catalog.isCached(second_view))
+        finally:
+            spark.catalog.clearCache()
+            spark.catalog.dropTempView(first_view)
+            other_session.catalog.dropTempView(second_view)
+            if is_remote():
+                other_session.client.close()
+
     def test_table_exists(self):
         # SPARK-36176: testing that table_exists returns correct boolean
         spark = self.spark

sql/api/src/main/scala/org/apache/spark/sql/catalog/Catalog.scala

@@ -645,6 +645,22 @@ abstract class Catalog {
    */
   def clearCache(): Unit
 
+  /**
+   * Removes cached tables from the in-memory cache.
+   *
+   * @param allSessions
+   *   when true, removes cached data across all Spark sessions. When false, removes cached data
+   *   owned by the current session while preserving data that is also cached by another session.
+   * @since 4.4.0
+   */
+  def clearCache(allSessions: Boolean): Unit = {
+    if (allSessions) {
+      clearCache()
+    } else {
+      catalogUnsupported("clearCache")
+    }
+  }
+
   /**
    * Invalidates and refreshes all the cached data and metadata of the given table. For
    * performance reasons, Spark SQL or the external data source library it uses might cache

sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/CatalogSuite.scala

@@ -209,6 +209,30 @@ class CatalogSuite extends ConnectFunSuite with RemoteSparkSession with SQLHelpe
     }
   }
 
+  test("SPARK-50569: clearCache can be scoped to the current session") {
+    val otherSession = spark.newSession()
+    val firstView = "clear_cache_first_view"
+    val secondView = "clear_cache_second_view"
+    try {
+      spark.range(1).createTempView(firstView)
+      otherSession.range(1, 2).createTempView(secondView)
+      spark.catalog.cacheTable(firstView)
+      otherSession.catalog.cacheTable(secondView)
+
+      spark.catalog.clearCache(allSessions = false)
+
+      assert(!spark.catalog.isCached(firstView))
+      assert(otherSession.catalog.isCached(secondView))
+      otherSession.catalog.clearCache(allSessions = false)
+      assert(!otherSession.catalog.isCached(secondView))
+    } finally {
+      spark.catalog.clearCache()
+      spark.catalog.dropTempView(firstView)
+      otherSession.catalog.dropTempView(secondView)
+      otherSession.close()
+    }
+  }
+
   test("TempView APIs") {
     val viewName = "view1"
     val globalViewName = "g_view1"

sql/connect/common/src/main/protobuf/spark/connect/catalog.proto

@@ -223,7 +223,10 @@ message UncacheTable {
 }
 
 // See `spark.catalog.clearCache`
-message ClearCache { }
+message ClearCache {
+  // (Optional) Whether to clear cached data across all sessions. Defaults to true when omitted.
+  optional bool all_sessions = 1;
+}
 
 // See `spark.catalog.refreshTable`
 message RefreshTable {

sql/connect/common/src/main/scala/org/apache/spark/sql/connect/Catalog.scala

@@ -643,8 +643,17 @@ class Catalog(sparkSession: SparkSession) extends catalog.Catalog {
    * @since 3.5.0
    */
   override def clearCache(): Unit = {
+    clearCache(allSessions = true)
+  }
+
+  /**
+   * Removes cached tables from the in-memory cache.
+   *
+   * @since 4.4.0
+   */
+  override def clearCache(allSessions: Boolean): Unit = {
     sparkSession.execute { builder =>
-      builder.getCatalogBuilder.getClearCacheBuilder
+      builder.getCatalogBuilder.getClearCacheBuilder.setAllSessions(allSessions)
     }
   }
 

sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala

@@ -156,6 +156,15 @@ object Connect {
       .timeConf(TimeUnit.MILLISECONDS)
       .createWithDefaultString("30s")
 
+  val CONNECT_SESSION_MANAGER_CLEANUP_CACHED_DATA_ENABLED =
+    buildStaticConf("spark.connect.session.manager.cleanupCachedData.enabled")
+      .doc(
+        "When true, cached data persisted by an isolated session is removed when the session " +
+          "is closed. Cached data that is also persisted by another session is preserved.")
+      .version("4.4.0")
+      .booleanConf
+      .createWithDefault(false)
+
   val CONNECT_EXECUTE_MANAGER_DETACHED_TIMEOUT =
     buildStaticConf("spark.connect.execute.manager.detachedTimeout")
       .internal()

sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala

@@ -306,7 +306,8 @@ class SparkConnectPlanner(
       case proto.Catalog.CatTypeCase.CACHE_TABLE => transformCacheTable(catalog.getCacheTable)
       case proto.Catalog.CatTypeCase.UNCACHE_TABLE =>
         transformUncacheTable(catalog.getUncacheTable)
-      case proto.Catalog.CatTypeCase.CLEAR_CACHE => transformClearCache()
+      case proto.Catalog.CatTypeCase.CLEAR_CACHE =>
+        transformClearCache(catalog.getClearCache)
       case proto.Catalog.CatTypeCase.REFRESH_TABLE =>
         transformRefreshTable(catalog.getRefreshTable)
       case proto.Catalog.CatTypeCase.REFRESH_BY_PATH =>
@@ -4303,8 +4304,9 @@ class SparkConnectPlanner(
     emptyLocalRelation
   }
 
-  private def transformClearCache(): LogicalPlan = {
-    session.catalog.clearCache()
+  private def transformClearCache(clearCache: proto.ClearCache): LogicalPlan = {
+    val allSessions = !clearCache.hasAllSessions || clearCache.getAllSessions
+    session.catalog.clearCache(allSessions)
     emptyLocalRelation
   }
 

sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala

@@ -434,6 +434,10 @@ case class SessionHolder(userId: String, sessionId: String, session: SparkSessio
     // Clean up ML cache (only if ML models were created)
     mlCache.close()
 
+    if (SparkEnv.get.conf.get(Connect.CONNECT_SESSION_MANAGER_CLEANUP_CACHED_DATA_ENABLED)) {
+      session.sharedState.cacheManager.clearCache(session)
+    }
+
     session.cleanupPythonWorkerLogs()
 
     eventManager.postClosed()

sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManagerSuite.scala

@@ -21,14 +21,28 @@ import java.util.UUID
 
 import org.scalatest.time.SpanSugar._
 
-import org.apache.spark.SparkSQLException
+import org.apache.spark.{SparkEnv, SparkSQLException}
 import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.connect.config.Connect
 import org.apache.spark.sql.pipelines.graph.{DataflowGraph, PipelineUpdateContextImpl}
 import org.apache.spark.sql.pipelines.logging.PipelineEvent
 import org.apache.spark.sql.test.SharedSparkSession
 
 class SparkConnectSessionManagerSuite extends SharedSparkSession {
 
+  private def withSparkConf(pairs: (String, String)*)(f: => Unit): Unit = {
+    val conf = SparkEnv.get.conf
+    val previousValues = pairs.map { case (key, _) => key -> conf.getOption(key) }
+    pairs.foreach { case (key, value) => conf.set(key, value) }
+    try f
+    finally {
+      previousValues.foreach {
+        case (key, Some(value)) => conf.set(key, value)
+        case (key, None) => conf.remove(key)
+      }
+    }
+  }
+
   override def beforeEach(): Unit = {
     super.beforeEach()
     SparkConnectService.sessionManager.invalidateAllSessions()
@@ -177,6 +191,66 @@ class SparkConnectSessionManagerSuite extends SharedSparkSession {
       "pipeline execution was not removed")
   }
 
+  test("SPARK-50569: cached data cleanup on session close is configurable and isolated") {
+    Seq(false, true).foreach { cleanupCachedData =>
+      withClue(s"cleanupCachedData=$cleanupCachedData") {
+        withSparkConf(
+          Connect.CONNECT_SESSION_MANAGER_CLEANUP_CACHED_DATA_ENABLED.key ->
+            cleanupCachedData.toString) {
+          val first = SparkConnectService.sessionManager.getOrCreateIsolatedSession(
+            SessionKey("user", UUID.randomUUID().toString),
+            None)
+          val second = SparkConnectService.sessionManager.getOrCreateIsolatedSession(
+            SessionKey("user", UUID.randomUUID().toString),
+            None)
+          val firstDataFrame = first.session.range(1)
+          second.session.range(1, 2).createTempView("second_view")
+          second.session.catalog.cacheTable("second_view")
+          val secondDataFrame = second.session.table("second_view")
+
+          firstDataFrame.persist()
+          SparkConnectService.sessionManager.closeSession(first.key)
+
+          assert(
+            first.session.sharedState.cacheManager.lookupCachedData(firstDataFrame).isDefined ===
+              !cleanupCachedData)
+          assert(
+            second.session.sharedState.cacheManager.lookupCachedData(secondDataFrame).isDefined)
+
+          SparkConnectService.sessionManager.closeSession(second.key)
+          assert(
+            second.session.sharedState.cacheManager
+              .lookupCachedData(secondDataFrame)
+              .isDefined ===
+              !cleanupCachedData)
+          spark.catalog.clearCache()
+        }
+      }
+    }
+  }
+
+  test("SPARK-50569: cached data cleanup preserves entries persisted by another session") {
+    withSparkConf(Connect.CONNECT_SESSION_MANAGER_CLEANUP_CACHED_DATA_ENABLED.key -> "true") {
+      val first = SparkConnectService.sessionManager.getOrCreateIsolatedSession(
+        SessionKey("user", UUID.randomUUID().toString),
+        None)
+      val second = SparkConnectService.sessionManager.getOrCreateIsolatedSession(
+        SessionKey("user", UUID.randomUUID().toString),
+        None)
+      val firstDataFrame = first.session.range(1)
+      val secondDataFrame = second.session.range(1)
+
+      firstDataFrame.persist()
+      secondDataFrame.persist()
+      SparkConnectService.sessionManager.closeSession(first.key)
+
+      assert(second.session.sharedState.cacheManager.lookupCachedData(secondDataFrame).isDefined)
+
+      SparkConnectService.sessionManager.closeSession(second.key)
+      assert(second.session.sharedState.cacheManager.lookupCachedData(secondDataFrame).isEmpty)
+    }
+  }
+
   test("baseSession allows creating sessions after default session is cleared") {
     // Create a new session manager to test initialization
     val sessionManager = new SparkConnectSessionManager()

sql/core/src/main/scala/org/apache/spark/sql/classic/Catalog.scala

@@ -875,7 +875,21 @@ class Catalog(sparkSession: SparkSession) extends catalog.Catalog with Logging {
    * @since 2.0.0
    */
   override def clearCache(): Unit = {
-    sparkSession.sharedState.cacheManager.clearCache()
+    clearCache(allSessions = true)
+  }
+
+  /**
+   * Removes cached tables or views from the in-memory cache.
+   *
+   * @group cachemgmt
+   * @since 4.4.0
+   */
+  override def clearCache(allSessions: Boolean): Unit = {
+    if (allSessions) {
+      sparkSession.sharedState.cacheManager.clearCache()
+    } else {
+      sparkSession.sharedState.cacheManager.clearCache(sparkSession)
+    }
   }
 
   /**