[GitHub] [spark] wankunde commented on a diff in pull request #37922: [SPARK-40480][SHUFFLE] Remove push-based shuffle data after query finished

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

wankunde commented on code in PR #37922:
URL: https://github.com/apache/spark/pull/37922#discussion_r1066901294


##########
core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala:
##########
@@ -321,6 +321,12 @@ class BlockManagerMasterEndpoint(
   }
 
   private def removeShuffle(shuffleId: Int): Future[Seq[Boolean]] = {
+    val mergerLocations =
+      if (Utils.isPushBasedShuffleEnabled(conf, isDriver)) {
+        mapOutputTracker.getShufflePushMergerLocations(shuffleId)
+      } else {
+        Seq.empty[BlockManagerId]
+      }

Review Comment:
   For `def removeShuffle(shuffleId: Int): Future[Seq[Boolean]]`  in `BlockManagerMasterEndpoint`, I prefer to put this code before 
   ```
       val removeMsg = RemoveShuffle(shuffleId)
       val removeShuffleFromExecutorsFutures = blockManagerInfo.values.map { bm =>
         bm.storageEndpoint.ask[Boolean](removeMsg).recover {
           // use false as default value means no shuffle data were removed
           handleBlockRemovalFailure("shuffle", shuffleId.toString, bm.blockManagerId, false)
         }
       }.toSeq
   ```
   because mapOutputTracker may unregisterShuffle this shuffle in `BlockManagerStorageEndpoint` before `mapOutputTracker.getShufflePushMergerLocations(shuffleId)`.
   ```
       case RemoveShuffle(shuffleId) =>
         doAsync[Boolean]("removing shuffle " + shuffleId, context) {
           if (mapOutputTracker != null) {
             mapOutputTracker.unregisterShuffle(shuffleId)
           }
           SparkEnv.get.shuffleManager.unregisterShuffle(shuffleId)
         }
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


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