otterc commented on a change in pull request #32140:
URL: https://github.com/apache/spark/pull/32140#discussion_r660275905
##########
File path:
core/src/main/scala/org/apache/spark/storage/ShuffleBlockFetcherIterator.scala
##########
@@ -347,35 +355,56 @@ 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: mutable.LinkedHashSet[(BlockId, Int)],
+ hostLocalBlocksByExecutor: mutable.LinkedHashMap[BlockManagerId,
Seq[(BlockId, Long, Int)]],
+ pushMergedLocalBlocks: 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, push-merged-local, remote (includes
push-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 pushMergedLocalBlockBytes = 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)) {
- checkBlockSizes(blockInfos)
+ checkBlockSizes(blockInfos)
+ if (pushBasedFetchHelper.isPushMergedShuffleBlockAddress(address)) {
+ // These are push-merged blocks or shuffle chunks of these blocks.
+ if (address.host == blockManager.blockManagerId.host) {
+ numBlocksToFetch += blockInfos.size
+ pushMergedLocalBlocks ++= blockInfos.map(_._1)
+ pushMergedLocalBlockBytes += blockInfos.map(_._3).sum
+ } else {
+ collectFetchRequests(address, blockInfos, collectedRemoteRequests)
+ }
+ } else if (mutable.HashSet(blockManager.blockManagerId.executorId,
fallback)
+ .contains(address.executorId)) {
Review comment:
Done
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
