[SPARK-29434][Core] Improve the MapStatuses Serialization Performance #26085
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -17,10 +17,9 @@ | |
|
|
||
| package org.apache.spark | ||
|
|
||
| import java.io.{ByteArrayInputStream, ObjectInputStream, ObjectOutputStream} | ||
| import java.util.concurrent.{ConcurrentHashMap, LinkedBlockingQueue, ThreadPoolExecutor, TimeUnit} | ||
| import java.util.concurrent.locks.ReentrantReadWriteLock | ||
|
|
||
| import scala.collection.JavaConverters._ | ||
| import scala.collection.mutable.{HashMap, ListBuffer, Map} | ||
|
|
@@ -29,6 +28,10 @@ import scala.concurrent.duration.Duration | |
| import scala.reflect.ClassTag | ||
| import scala.util.control.NonFatal | ||
|
|
||
| import com.github.luben.zstd.ZstdInputStream | ||
| import com.github.luben.zstd.ZstdOutputStream | ||
| import org.apache.commons.io.output.{ByteArrayOutputStream => ApacheByteArrayOutputStream} | ||
|
|
||
| import org.apache.spark.broadcast.{Broadcast, BroadcastManager} | ||
| import org.apache.spark.internal.Logging | ||
| import org.apache.spark.internal.config._ | ||
|
|
@@ -885,13 +888,18 @@ private[spark] object MapOutputTracker extends Logging { | |
| private val BROADCAST = 1 | ||
|
|
||
| // Serialize an array of map output locations into an efficient byte format so that we can send | ||
| // it to reduce tasks. We do this by compressing the serialized bytes using Zstd. They will | ||
| // generally be pretty compressible because many map outputs will be on the same hostname. | ||
| def serializeMapStatuses(statuses: Array[MapStatus], broadcastManager: BroadcastManager, | ||
| isLocal: Boolean, minBroadcastSize: Int): (Array[Byte], Broadcast[Array[Byte]]) = { | ||
| // Using `org.apache.commons.io.output.ByteArrayOutputStream` instead of the standard one | ||
| // This implementation doesn't reallocate the whole memory block but allocates | ||
| // additional buffers. This way no buffers need to be garbage collected and | ||
| // the contents don't have to be copied to the new buffer. | ||
| val out = new ApacheByteArrayOutputStream() | ||
| val compressedOut = new ApacheByteArrayOutputStream() | ||
|
|
||
| val objOut = new ObjectOutputStream(out) | ||
| Utils.tryWithSafeFinally { | ||
| // Since statuses can be modified in parallel, sync on it | ||
| statuses.synchronized { | ||
|
|
@@ -900,18 +908,42 @@ private[spark] object MapOutputTracker extends Logging { | |
| } { | ||
| objOut.close() | ||
| } | ||
|
|
||
| val arr: Array[Byte] = { | ||
| val zos = new ZstdOutputStream(compressedOut) | ||
| Utils.tryWithSafeFinally { | ||
| compressedOut.write(DIRECT) | ||
| // `out.writeTo(zos)` will write the uncompressed data from `out` to `zos` | ||
| // without copying to avoid unnecessary allocation and copy of byte[]. | ||
| out.writeTo(zos) | ||
| } { | ||
| zos.close() | ||
| } | ||
| compressedOut.toByteArray | ||
| } | ||
| if (arr.length >= minBroadcastSize) { | ||
| // Use broadcast instead. | ||
| // Important arr(0) is the tag == DIRECT, ignore that while deserializing ! | ||
| val bcast = broadcastManager.newBroadcast(arr, isLocal) | ||
| // toByteArray creates copy, so we can reuse out | ||
dongjoon-hyun marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| out.reset() | ||
| val oos = new ObjectOutputStream(out) | ||
| Utils.tryWithSafeFinally { | ||
| oos.writeObject(bcast) | ||
| } { | ||
| oos.close() | ||
| } | ||
| val outArr = { | ||
| compressedOut.reset() | ||
| val zos = new ZstdOutputStream(compressedOut) | ||
|
Contributor
There was a problem hiding this comment. Hi, @dbtsai , I am back-porting this into our internal repo. Looks like this compression is unnecessary since
Member
Author
There was a problem hiding this comment. The actually value of the data (which is already compressed) will not be in the serialized form of
Contributor
There was a problem hiding this comment. Thanks for your clarification. It's indeed not including the compressed data. |
||
| Utils.tryWithSafeFinally { | ||
| compressedOut.write(BROADCAST) | ||
| out.writeTo(zos) | ||
| } { | ||
| zos.close() | ||
| } | ||
| compressedOut.toByteArray | ||
| } | ||
| logInfo("Broadcast mapstatuses size = " + outArr.length + ", actual size = " + arr.length) | ||
| (outArr, bcast) | ||
| } else { | ||
|
|
@@ -924,7 +956,7 @@ private[spark] object MapOutputTracker extends Logging { | |
| assert (bytes.length > 0) | ||
|
|
||
| def deserializeObject(arr: Array[Byte], off: Int, len: Int): AnyRef = { | ||
| val objIn = new ObjectInputStream(new ZstdInputStream( | ||
| new ByteArrayInputStream(arr, off, len))) | ||
| Utils.tryWithSafeFinally { | ||
| objIn.readObject() | ||
|
|
||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.