[SPARK-53446][CORE] Optimize BlockManager remove operations with cached block mappings

### What changes were proposed in this pull request?
Continue #52210.
Introduced three concurrent hash maps to track block ID associations for optimize BlockManager remove operations by introducing cached mappings to eliminate O(n) linear scans.

### Why are the changes needed?
Previously, removeRdd(), removeBroadcast(), and removeCache() required scanning all blocks in blockInfoManager.entries to find matches. This approach becomes a serious bottleneck when:

1. Large block counts: In production deployments with millions or even tens of millions of cached blocks, linear scans can be prohibitively slow
2. High cleanup frequency: Workloads that repeatedly create and discard RDDs or broadcast variables accumulate overhead quickly

The original removeRdd() method already contained a TODO noting that an additional mapping would be needed to avoid linear scans. This PR implements that improvement.

### Does this PR introduce _any_ user-facing change?
No.

### How was this patch tested?
New UT.

### Was this patch authored or co-authored using generative AI tooling?
No.

Closes #52646 from zml1206/SPARK-53446.

Authored-by: zml1206 <zhuml1206@gmail.com>
Signed-off-by: Wenchen Fan <wenchen@databricks.com>

Author: zml1206

core/src/main/scala/org/apache/spark/storage/BlockInfoManager.scala

@@ -149,6 +149,14 @@ private[storage] class BlockInfoManager(trackingCacheVisibility: Boolean = false
    */
   private[this] val blockInfoWrappers = new ConcurrentHashMap[BlockId, BlockInfoWrapper]
 
+  // Cache mappings to avoid O(n) scans in remove operations.
+  private[this] val rddToBlockIds =
+    new ConcurrentHashMap[Int, ConcurrentHashMap.KeySetView[BlockId, java.lang.Boolean]]()
+  private[this] val broadcastToBlockIds =
+    new ConcurrentHashMap[Long, ConcurrentHashMap.KeySetView[BlockId, java.lang.Boolean]]()
+  private[this] val sessionToBlockIds =
+    new ConcurrentHashMap[String, ConcurrentHashMap.KeySetView[BlockId, java.lang.Boolean]]()
+
   /**
    * Record invisible rdd blocks stored in the block manager, entries will be removed when blocks
    * are marked as visible or blocks are removed by [[removeBlock()]].
@@ -445,6 +453,7 @@ private[storage] class BlockInfoManager(trackingCacheVisibility: Boolean = false
         }
 
         if (previous == null) {
+          addToMapping(blockId)
           // New block lock it for writing.
           val result = lockForWriting(blockId, blocking = false)
           assert(result.isDefined)
@@ -535,6 +544,23 @@ private[storage] class BlockInfoManager(trackingCacheVisibility: Boolean = false
     blockInfoWrappers.entrySet().iterator().asScala.map(kv => kv.getKey -> kv.getValue.info)
   }
 
+  /**
+   * Return all blocks belonging to the given RDD.
+   */
+  def rddBlockIds(rddId: Int): Seq[BlockId] = getBlockIdsFromMapping(rddToBlockIds, rddId)
+
+  /**
+   * Return all blocks belonging to the given broadcast.
+   */
+  def broadcastBlockIds(broadcastId: Long): Seq[BlockId] =
+    getBlockIdsFromMapping(broadcastToBlockIds, broadcastId)
+
+  /**
+   * Return cache blocks that might be related to cached local relations.
+   */
+  def sessionBlockIds(sessionUUID: String): Seq[BlockId] =
+    getBlockIdsFromMapping(sessionToBlockIds, sessionUUID)
+
   /**
    * Removes the given block and releases the write lock on it.
    *
@@ -551,6 +577,7 @@ private[storage] class BlockInfoManager(trackingCacheVisibility: Boolean = false
       } else {
         invisibleRDDBlocks.synchronized {
           blockInfoWrappers.remove(blockId)
+          removeFromMapping(blockId)
           blockId.asRDDId.foreach(invisibleRDDBlocks.remove)
         }
         info.readerCount = 0
@@ -573,11 +600,75 @@ private[storage] class BlockInfoManager(trackingCacheVisibility: Boolean = false
       }
     }
     blockInfoWrappers.clear()
+    rddToBlockIds.clear()
+    broadcastToBlockIds.clear()
+    sessionToBlockIds.clear()
     readLocksByTask.clear()
     writeLocksByTask.clear()
     invisibleRDDBlocks.synchronized {
       invisibleRDDBlocks.clear()
     }
   }
 
+  /**
+   * Return all blocks in the cache mapping for a given key.
+   */
+  private def getBlockIdsFromMapping[K](
+      map: ConcurrentHashMap[K, ConcurrentHashMap.KeySetView[BlockId, java.lang.Boolean]],
+      key: K): Seq[BlockId] = {
+    Option(map.get(key)).map(_.asScala.toSeq).getOrElse(Seq.empty)
+  }
+
+  /**
+   * Add a block ID to the corresponding cache mapping based on its type.
+   */
+  private def addToMapping(blockId: BlockId): Unit = {
+    blockId match {
+      case rddBlockId: RDDBlockId =>
+        rddToBlockIds
+          .computeIfAbsent(rddBlockId.rddId, _ => ConcurrentHashMap.newKeySet())
+          .add(blockId)
+      case broadcastBlockId: BroadcastBlockId =>
+        broadcastToBlockIds
+          .computeIfAbsent(broadcastBlockId.broadcastId, _ => ConcurrentHashMap.newKeySet())
+          .add(blockId)
+      case cacheId: CacheId =>
+        sessionToBlockIds
+          .computeIfAbsent(cacheId.sessionUUID, _ => ConcurrentHashMap.newKeySet())
+          .add(blockId)
+      case _ => // Do nothing for other block types
+    }
+  }
+
+  /**
+   * Remove a block ID from the corresponding cache mapping based on its type.
+   */
+  private def removeFromMapping(blockId: BlockId): Unit = {
+    def doRemove[K](
+        map: ConcurrentHashMap[K, ConcurrentHashMap.KeySetView[BlockId, java.lang.Boolean]],
+        key: K,
+        block: BlockId): Unit = {
+      map.compute(key,
+        (_, set) => {
+          if (null != set) {
+            set.remove(block)
+            if (set.isEmpty) null else set
+          } else {
+            // missing
+            null
+          }
+        }
+      )
+    }
+
+    blockId match {
+      case rddBlockId: RDDBlockId =>
+        doRemove(rddToBlockIds, rddBlockId.rddId, rddBlockId)
+      case broadcastBlockId: BroadcastBlockId =>
+        doRemove(broadcastToBlockIds, broadcastBlockId.broadcastId, broadcastBlockId)
+      case cacheId: CacheId =>
+        doRemove(sessionToBlockIds, cacheId.sessionUUID, cacheId)
+      case _ => // Do nothing for other block types
+    }
+  }
 }

core/src/main/scala/org/apache/spark/storage/BlockManager.scala

@@ -2054,9 +2054,8 @@ private[spark] class BlockManager(
    * @return The number of blocks removed.
    */
   def removeRdd(rddId: Int): Int = {
-    // TODO: Avoid a linear scan by creating another mapping of RDD.id to blocks.
     logInfo(log"Removing RDD ${MDC(RDD_ID, rddId)}")
-    val blocksToRemove = blockInfoManager.entries.flatMap(_._1.asRDDId).filter(_.rddId == rddId)
+    val blocksToRemove = blockInfoManager.rddBlockIds(rddId)
     blocksToRemove.foreach { blockId => removeBlock(blockId, tellMaster = false) }
     blocksToRemove.size
   }
@@ -2090,9 +2089,7 @@ private[spark] class BlockManager(
    */
   def removeBroadcast(broadcastId: Long, tellMaster: Boolean): Int = {
     logDebug(s"Removing broadcast $broadcastId")
-    val blocksToRemove = blockInfoManager.entries.map(_._1).collect {
-      case bid @ BroadcastBlockId(`broadcastId`, _) => bid
-    }
+    val blocksToRemove = blockInfoManager.broadcastBlockIds(broadcastId)
     blocksToRemove.foreach { blockId => removeBlock(blockId, tellMaster) }
     blocksToRemove.size
   }
@@ -2104,9 +2101,7 @@ private[spark] class BlockManager(
    */
   def removeCache(sessionUUID: String): Int = {
     logDebug(s"Removing cache of spark session with UUID: $sessionUUID")
-    val blocksToRemove = blockInfoManager.entries.map(_._1).collect {
-      case cid: CacheId if cid.sessionUUID == sessionUUID => cid
-    }
+    val blocksToRemove = blockInfoManager.sessionBlockIds(sessionUUID)
     blocksToRemove.foreach { blockId => removeBlock(blockId) }
     blocksToRemove.size
   }

core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala

@@ -2634,6 +2634,28 @@ class BlockManagerSuite extends SparkFunSuite with Matchers with PrivateMethodTe
     assert(logBlockIds.contains(logBlockId1) && logBlockIds.contains(logBlockId2))
   }
 
+  test("SPARK-53446: Optimize BlockManager remove operations with cached block mappings") {
+    val store = makeBlockManager(8000, "executor1")
+    val broadcastId = 0
+    val rddId = 1
+    val sessionId = UUID.randomUUID.toString
+    val data = new Array[Byte](100)
+
+    store.putSingle(BroadcastBlockId(broadcastId), data, StorageLevel.MEMORY_ONLY)
+    assert(store.blockInfoManager.broadcastBlockIds(broadcastId).nonEmpty)
+    store.putSingle(rdd(rddId, 3), data, StorageLevel.MEMORY_ONLY)
+    assert(store.blockInfoManager.rddBlockIds(rddId).nonEmpty)
+    store.putSingle(CacheId(sessionId, "abc"), data, StorageLevel.MEMORY_ONLY)
+    assert(store.blockInfoManager.sessionBlockIds(sessionId).nonEmpty)
+
+    store.removeBroadcast(broadcastId, false)
+    assert(store.blockInfoManager.broadcastBlockIds(broadcastId).isEmpty)
+    store.removeRdd(rddId)
+    assert(store.blockInfoManager.rddBlockIds(rddId).isEmpty)
+    store.removeCache(sessionId)
+    assert(store.blockInfoManager.sessionBlockIds(sessionId).isEmpty)
+  }
+
   private def createKryoSerializerWithDiskCorruptedInputStream(): KryoSerializer = {
     class TestDiskCorruptedInputStream extends InputStream {
       override def read(): Int = throw new IOException("Input/output error")