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

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

@@ -156,6 +156,14 @@ 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/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,60 @@ 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/execution/CacheManager.scala

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.execution
 
+import java.util.IdentityHashMap
+
 import scala.util.control.NonFatal
 
 import org.apache.hadoop.fs.{FileSystem, Path}
@@ -79,13 +81,47 @@ class CacheManager extends Logging with AdaptiveSparkPlanHelper {
   @transient @volatile
   private var cachedData = IndexedSeq[CachedData]()
 
+  /**
+   * Tracks the sessions that explicitly cached each entry. A cache entry can be shared by
+   * multiple sessions because cache lookup uses plan semantics rather than session identity.
+   */
+  @transient
+  private val cacheOwners = new IdentityHashMap[CachedData, Set[String]]
+
   /** Clears all cached tables. */
   def clearCache(): Unit = this.synchronized {
     cachedData.foreach(_.cachedRepresentation.cacheBuilder.clearCache())
     cachedData = IndexedSeq[CachedData]()
+    cacheOwners.clear()
     CacheManager.logCacheOperation(log"Cleared all Dataframe cache entries")
   }
 
+  /** Clears cached data that is owned only by the given session. */
+  private[sql] def clearCache(session: SparkSession): Unit = {
+    val sessionUUID = session.sessionUUID
+    val plansToUncache = this.synchronized {
+      val plans = cachedData.filter { cd =>
+        Option(cacheOwners.get(cd)).contains(Set(sessionUUID))
+      }
+      cachedData = cachedData.filterNot(cd => plans.exists(_ eq cd))
+      val ownerEntries = cacheOwners.entrySet().iterator()
+      while (ownerEntries.hasNext) {
+        val entry = ownerEntries.next()
+        val remainingOwners = entry.getValue - sessionUUID
+        if (remainingOwners.nonEmpty) {
+          entry.setValue(remainingOwners)
+        } else {
+          ownerEntries.remove()
+        }
+      }
+      plans
+    }
+    plansToUncache.foreach(_.cachedRepresentation.cacheBuilder.clearCache())
+    CacheManager.logCacheOperation(
+      log"Cleared ${MDC(SIZE, plansToUncache.size)} Dataframe cache entries for session " +
+        log"${MDC(SESSION_ID, sessionUUID)}")
+  }
+
   /** Checks if the cache is empty. */
   def isEmpty: Boolean = {
     cachedData.isEmpty
@@ -144,7 +180,7 @@ class CacheManager extends Logging with AdaptiveSparkPlanHelper {
         log"Asked to cache a plan that is inapplicable for caching: " +
         log"${MDC(LOGICAL_PLAN, unnormalizedPlan)}"
       )
-    } else if (lookupCachedDataInternal(normalizedPlan).nonEmpty) {
+    } else if (registerCacheOwner(normalizedPlan, spark.sessionUUID)) {
       logWarning("Asked to cache already cached data.")
     } else {
       val sessionWithConfigsOff = getOrCloneSessionWithConfigsOff(spark)
@@ -158,19 +194,31 @@ class CacheManager extends Logging with AdaptiveSparkPlanHelper {
       }
 
       this.synchronized {
-        if (lookupCachedDataInternal(normalizedPlan).nonEmpty) {
+        if (registerCacheOwner(normalizedPlan, spark.sessionUUID)) {
           logWarning("Data has already been cached.")
         } else {
           // the cache key is the normalized plan
           val cd = CachedData(normalizedPlan, inMemoryRelation)
           cachedData = cd +: cachedData
+          cacheOwners.put(cd, Set(spark.sessionUUID))
           CacheManager.logCacheOperation(log"Added Dataframe cache entry:" +
             log"${MDC(DATAFRAME_CACHE_ENTRY, cd)}")
         }
       }
     }
   }
 
+  private def registerCacheOwner(plan: LogicalPlan, sessionUUID: String): Boolean =
+    this.synchronized {
+      lookupCachedDataInternal(plan) match {
+        case Some(cd) =>
+          val owners = Option(cacheOwners.get(cd)).getOrElse(Set.empty)
+          cacheOwners.put(cd, owners + sessionUUID)
+          true
+        case None => false
+      }
+    }
+
   /**
    * Un-cache the given plan or all the cache entries that refer to the given plan.
    *
@@ -301,6 +349,7 @@ class CacheManager extends Logging with AdaptiveSparkPlanHelper {
     val plansToUncache = cachedData.filter(cd => shouldRemove(cd.plan))
     this.synchronized {
       cachedData = cachedData.filterNot(cd => plansToUncache.exists(_ eq cd))
+      plansToUncache.foreach(cd => cacheOwners.remove(cd))
     }
     plansToUncache.foreach { _.cachedRepresentation.cacheBuilder.clearCache(blocking) }
     CacheManager.logCacheOperation(log"Removed ${MDC(SIZE, plansToUncache.size)} Dataframe " +
@@ -374,14 +423,22 @@ class CacheManager extends Logging with AdaptiveSparkPlanHelper {
     }
     needToRecache.foreach { cd =>
       cd.cachedRepresentation.cacheBuilder.clearCache()
-      tryRebuildCacheEntry(spark, cd).foreach { entry =>
-        this.synchronized {
-          if (lookupCachedDataInternal(entry.plan).nonEmpty) {
-            logWarning("While recaching, data was already added to cache.")
-          } else {
-            cachedData = entry +: cachedData
-            CacheManager.logCacheOperation(log"Re-cached Dataframe cache entry:" +
-              log"${MDC(DATAFRAME_CACHE_ENTRY, entry)}")
+      val rebuiltEntry = tryRebuildCacheEntry(spark, cd)
+      this.synchronized {
+        val previousOwners = Option(cacheOwners.remove(cd)).getOrElse(Set.empty)
+        rebuiltEntry.foreach { entry =>
+          if (previousOwners.nonEmpty) {
+            lookupCachedDataInternal(entry.plan) match {
+              case Some(existing) =>
+                val existingOwners = Option(cacheOwners.get(existing)).getOrElse(Set.empty)
+                cacheOwners.put(existing, existingOwners ++ previousOwners)
+                logWarning("While recaching, data was already added to cache.")
+              case None =>
+                cachedData = entry +: cachedData
+                cacheOwners.put(entry, previousOwners)
+                CacheManager.logCacheOperation(log"Re-cached Dataframe cache entry:" +
+                  log"${MDC(DATAFRAME_CACHE_ENTRY, entry)}")
+            }
           }
         }
       }