otterc commented on a change in pull request #32140:
URL: https://github.com/apache/spark/pull/32140#discussion_r648537705



##########
File path: 
core/src/main/scala/org/apache/spark/storage/ShuffleBlockFetcherIterator.scala
##########
@@ -347,20 +361,48 @@ final class ShuffleBlockFetcherIterator(
     }
   }
 
-  private[this] def partitionBlocksByFetchMode(): ArrayBuffer[FetchRequest] = {
+  /**
+   * This is called from initialize and also from the fallback which is 
triggered from
+   * [[PushBasedFetchHelper]].
+   */
+  private[this] def partitionBlocksByFetchMode(
+      blocksByAddress: Iterator[(BlockManagerId, Seq[(BlockId, Long, Int)])],
+      localBlocks: scala.collection.mutable.LinkedHashSet[(BlockId, Int)],
+      hostLocalBlocksByExecutor: mutable.LinkedHashMap[BlockManagerId, 
Seq[(BlockId, Long, Int)]],
+      mergedLocalBlocks: mutable.LinkedHashSet[BlockId]): 
ArrayBuffer[FetchRequest] = {
     logDebug(s"maxBytesInFlight: $maxBytesInFlight, targetRemoteRequestSize: "
       + s"$targetRemoteRequestSize, maxBlocksInFlightPerAddress: 
$maxBlocksInFlightPerAddress")
 
-    // Partition to local, host-local and remote blocks. Remote blocks are 
further split into
-    // FetchRequests of size at most maxBytesInFlight in order to limit the 
amount of data in flight
+    // Partition to local, host-local, merged-local, remote (includes 
merged-remote) blocks.
+    // Remote blocks are further split into FetchRequests of size at most 
maxBytesInFlight in order
+    // to limit the amount of data in flight
     val collectedRemoteRequests = new ArrayBuffer[FetchRequest]
+    val hostLocalBlocksCurrentIteration = mutable.LinkedHashSet[(BlockId, 
Int)]()
     var localBlockBytes = 0L
     var hostLocalBlockBytes = 0L
+    var mergedLocalBlockBytes = 0L
     var remoteBlockBytes = 0L
+    val prevNumBlocksToFetch = numBlocksToFetch
 
     val fallback = FallbackStorage.FALLBACK_BLOCK_MANAGER_ID.executorId
     for ((address, blockInfos) <- blocksByAddress) {
-      if (Seq(blockManager.blockManagerId.executorId, 
fallback).contains(address.executorId)) {
+      if (pushBasedFetchHelper.isMergedShuffleBlockAddress(address)) {
+        // These are push-based merged blocks or chunks of these merged blocks.
+        if (address.host == blockManager.blockManagerId.host) {
+          checkBlockSizes(blockInfos)

Review comment:
       For a merged blocks that are remote, this check is performed in 
`collectFetchRequests`. So, if we do it before the condition ` if (address.host 
== blockManager.blockManagerId.host)`, then it will be done twice for remote 
merged blocks.
   I think this is also a reason that currently `checkBlockSizes()`  is being 
called explicitly for each block type because for remote blocks (the last else 
block), the size is validated in `collectFetchRequests`.




-- 
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.

For queries about this service, please contact Infrastructure at:
[email protected]



---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to