[GitHub] attilapiros commented on a change in pull request #23679: [SPARK-23516][CORE] Avoid repeated release/acquire on storage memory

[GitBox <gi...@apache.org>]

attilapiros commented on a change in pull request #23679: [SPARK-23516][CORE] Avoid repeated release/acquire on storage memory
URL: https://github.com/apache/spark/pull/23679#discussion_r256930530
 
 

 ##########
 File path: core/src/main/scala/org/apache/spark/storage/memory/MemoryStore.scala
 ##########
 @@ -586,6 +587,29 @@ private[spark] class MemoryStore(
     }
   }
 
+  /**
+   * Only release memory used by this task from unroll memory
+   * map(`onHeapUnrollMemoryMap` or `offHeapUnrollMemoryMap`).
+   */
+  def releaseMemoryFromUnrollMemoryMap(
+      memoryMode: MemoryMode,
+      memory: Long = Long.MaxValue): Unit = {
+    val taskAttemptId = currentTaskAttemptId()
+    memoryManager.synchronized {
+      val unrollMemoryMap = memoryMode match {
+        case MemoryMode.ON_HEAP => onHeapUnrollMemoryMap
+        case MemoryMode.OFF_HEAP => offHeapUnrollMemoryMap
+      }
+      if (unrollMemoryMap.contains(taskAttemptId)) {
 
 Review comment:
   Possible micro-optimisation: you can use `getOrElse(taskAttemptId, -1)` instead of `contains` then you can decide whether the key is contained and get the value for the key with one lookup. 

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services

---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org