[SPARK-7081] Faster sort-based shuffle path using binary processing cache-aware sort

[JoshRosen]

This patch introduces a new shuffle manager that enhances the existing sort-based shuffle with a new cache-friendly sort algorithm that operates directly on binary data. The goals of this patch are to lower memory usage and Java object overheads during shuffle and to speed up sorting. It also lays groundwork for follow-up patches that will enable end-to-end processing of serialized records.

The new shuffle manager, UnsafeShuffleManager, can be enabled by setting spark.shuffle.manager=tungsten-sort in SparkConf.

The new shuffle manager uses directly-managed memory to implement several performance optimizations for certain types of shuffles. In cases where the new performance optimizations cannot be applied, the new shuffle manager delegates to SortShuffleManager to handle those shuffles.

UnsafeShuffleManager's optimizations will apply when all of the following conditions hold:

• The shuffle dependency specifies no aggregation or output ordering.
• The shuffle serializer supports relocation of serialized values (this is currently supported
  by KryoSerializer and Spark SQL's custom serializers).
• The shuffle produces fewer than 16777216 output partitions.
• No individual record is larger than 128 MB when serialized.

In addition, extra spill-merging optimizations are automatically applied when the shuffle compression codec supports concatenation of serialized streams. This is currently supported by Spark's LZF serializer.

At a high-level, UnsafeShuffleManager's design is similar to Spark's existing SortShuffleManager. In sort-based shuffle, incoming records are sorted according to their target partition ids, then written to a single map output file. Reducers fetch contiguous regions of this file in order to read their portion of the map output. In cases where the map output data is too large to fit in memory, sorted subsets of the output can are spilled to disk and those on-disk files are merged to produce the final output file.

UnsafeShuffleManager optimizes this process in several ways:

• Its sort operates on serialized binary data rather than Java objects, which reduces memory consumption and GC overheads. This optimization requires the record serializer to have certain properties to allow serialized records to be re-ordered without requiring deserialization. See SPARK-4550, where this optimization was first proposed and implemented, for more details.
• It uses a specialized cache-efficient sorter (UnsafeShuffleExternalSorter) that sorts arrays of compressed record pointers and partition ids. By using only 8 bytes of space per record in the sorting array, this fits more of the array into cache.
• The spill merging procedure operates on blocks of serialized records that belong to the same partition and does not need to deserialize records during the merge.
• When the spill compression codec supports concatenation of compressed data, the spill merge simply concatenates the serialized and compressed spill partitions to produce the final output partition. This allows efficient data copying methods, like NIO's transferTo, to be used and avoids the need to allocate decompression or copying buffers during the merge.

The shuffle read path is unchanged.

This patch is similar to SPARK-4550 / #4450 but uses a slightly different implementation. The unsafe-based implementation featured in this patch lays the groundwork for followup patches that will enable sorting to operate on serialized data pages that will be prepared by Spark SQL's new unsafe operators (such as the new aggregation operator introduced in #5725).

Future work

There are several tasks that build upon this patch, which will be left to future work:

• SPARK-7271 Redesign / extend the shuffle interfaces to accept binary data as input. The goal here is to let us bypass serialization steps in cases where the sort input is produced by an operator that operates directly on binary data.
• Extension / redesign of the Serializer API. We can add new methods which allow serializers to determine the size requirements for serializing objects and for serializing objects directly to a specified memory address (similar to how UnsafeRowConverter works in Spark SQL).

[AmplabJenkins]

Merged build triggered.

[AmplabJenkins]

Merged build started.

[SparkQA]

Test build #31671 has started for PR 5868 at commit 240864c.

[SparkQA]

Test build #31671 has finished for PR 5868 at commit 240864c.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds the following public classes (experimental):
  • public class UnsafeShuffleWriter<K, V> implements ShuffleWriter<K, V>
  • final class UnsafeSortDataFormat extends SortDataFormat<RecordPointerAndKeyPrefix, long[]>
  • public final class UnsafeSorter
  • public static final class RecordPointerAndKeyPrefix
  • public static abstract class RecordComparator
  • public static abstract class PrefixComparator

[AmplabJenkins]

Merged build finished. Test PASSed.

[AmplabJenkins]

Test PASSed.
Refer to this link for build results (access rights to CI server needed):
https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/31671/
Test PASSed.

[JoshRosen]

Whoops, didn't mean to include this logging level change. I'll roll it back.

[AmplabJenkins]

Merged build triggered.

[AmplabJenkins]

Merged build started.

[AmplabJenkins]

Merged build started.

[SparkQA]

Test build #32646 has started for PR 5868 at commit 7610f2f.

[SparkQA]

Test build #32646 has finished for PR 5868 at commit 7610f2f.

• This patch fails Scala style tests.
• This patch merges cleanly.
• This patch adds the following public classes (experimental):
  • final class DummySerializerInstance extends SerializerInstance
  • final class PackedRecordPointer
  • final class SpillInfo
  • final class UnsafeShuffleExternalSorter
  • final class UnsafeShuffleInMemorySorter
  • public static final class UnsafeShuffleSorterIterator
  • final class UnsafeShuffleSortDataFormat extends SortDataFormat<PackedRecordPointer, long[]>
  • public class UnsafeShuffleWriter<K, V> extends ShuffleWriter<K, V>
  • public final class TimeTrackingOutputStream extends OutputStream

[AmplabJenkins]

Merged build finished. Test FAILed.

[AmplabJenkins]

Test FAILed.
Refer to this link for build results (access rights to CI server needed):
https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/32646/
Test FAILed.

[AmplabJenkins]

Merged build triggered.

[AmplabJenkins]

Merged build started.

[SparkQA]

Test build #32648 has started for PR 5868 at commit ef0a86e.

[JoshRosen]

Here are some performance benchmark results based on a microbenchmark that I ran locally on my laptop. This benchmark measures the time for a single map task to write its shuffle output. The data type is (Int, Int) key-value pairs and I used 1000 shuffle output partitions. The source for this benchmark can be found at https://gist.github.com/JoshRosen/15ab99014fb21cbd1e65

In this benchmark, sort_serialize is the SortShuffleManager with #4450 optimizations enabled. I also tried benchmarking against Spark 1.2's sort-based shuffle but gave up because it quickly ran into GC thrashing problems (even with spilling enabled).

Each result is the average of 5 runs, preceded by 3 warmup runs.

	mean				std
	timeMs				timeMs
numKeys	10000000	20000000	30000000	40000000	10000000	20000000	30000000	40000000
mode
sort_serialize	5057.6	17871.2	27368.0	38202	174.925127	270.966234	2209.958371	1820.343237
tungsten-sort	4207.6	10720.6	22414.4	23966	643.662800	896.888678	5501.487917	4143.371514

To help interpret these results, here's the corresponding table that tracks the number of disk bytes spilled:

	mean
	diskBytesSpilled
numKeys	10000000	20000000	30000000	40000000
mode
sort_serialize	0	113607985	113607985	229053192
tungsten-sort	0	113599278	113599284	229043276

As I would expect, both shuffle implementations perform similarly when data fits in memory. I think that UnsafeShuffleManager's optimized shuffle merge path explains the large performance gains for shuffles that trigger disk spilling. Most of the variability here probably comes from my laptop's SSD.

I plan to run more comprehensive benchmarks on EC2, where we'll be able to see the benefits of using multiple disks and larger JVM heaps.

[zsxwing]

A bit confused. So the size of the last file won't be counted as DiskBytesSpilled nor ShuffleRecordsWritten?

[JoshRosen]

If control reaches this point, then we have written multiple files to disk in our external sorter. The last file should not be counted as bytes spilled in order to try to remain consistent with other code. In the existing sort-merge code, the final in-memory partition is merged against the on-disk partitions and thus does not get counted as extra disk write. Here, we end up dumping the in-memory partition to disk, incurring a bit of extra write but massively simplifying the merge path code. It's true that we're performing some disk writes that aren't accounted for anywhere in the final metrics, but I'm not sure how this should be reported / handled. I don't think that it's correct to treat these bytes as spilled bytes, but treating them as shuffle write is also confusing because then we'll end up reporting more shuffle bytes written than are transferred over the network.

The reason that I decrement the count here is so that the final count of shuffle bytes written is equal to the output file's size. Alternatively, I could have avoided incrementing shuffleBytesWritten when writing the last file, but that would cause complications when updating record counts.

I'm happy to work on fixes for the metrics, but I'd rather defer that to a separate followup patch that does a broader refactoring of our metrics code.

[SparkQA]

Test build #32648 has finished for PR 5868 at commit ef0a86e.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds the following public classes (experimental):
  • final class DummySerializerInstance extends SerializerInstance
  • final class PackedRecordPointer
  • final class SpillInfo
  • final class UnsafeShuffleExternalSorter
  • final class UnsafeShuffleInMemorySorter
  • public static final class UnsafeShuffleSorterIterator
  • final class UnsafeShuffleSortDataFormat extends SortDataFormat<PackedRecordPointer, long[]>
  • public class UnsafeShuffleWriter<K, V> extends ShuffleWriter<K, V>
  • public final class TimeTrackingOutputStream extends OutputStream

[AmplabJenkins]

Merged build finished. Test PASSed.

[AmplabJenkins]

Test PASSed.
Refer to this link for build results (access rights to CI server needed):
https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/32648/
Test PASSed.

[zsxwing]

LGTM now.

[rxin]

Thanks. I'm merging this in. The 100% unit test coverage is impressive. Let's continue to iterate and improve it.

[JoshRosen]

If anyone's been watching the Spark master Maven build, you may have noticed that it's started failing since this patch. I spent a bunch of time investigating this and narrowed it down to an issue where JavaAPISuite seemed to fail a pair of its sort tests if JavaAPISuite was run immediately after UnsafeShuffleWriterSuite. It turns out that this failure was not caused by this patch but was merely exposed by its impact on timing / scheduling, so please don't revert this!

It looks like there's a bug in snappy-java that can lead to stream corruption issues and the timing changes in the tests were sufficient to expose it: xerial/snappy-java#107. The same failures can be reproduced by adding small Thread.sleep calls in SnappyCompressionCodec.createOutputStream:

diff --git a/core/src/main/scala/org/apache/spark/io/CompressionCodec.scala b/core/src/main/scala/org/apache/spark/io/CompressionCodec.scala
index 0756cdb..cc2a912 100644
--- a/core/src/main/scala/org/apache/spark/io/CompressionCodec.scala
+++ b/core/src/main/scala/org/apache/spark/io/CompressionCodec.scala
@@ -154,6 +154,7 @@ class SnappyCompressionCodec(conf: SparkConf) extends CompressionCodec {

   override def compressedOutputStream(s: OutputStream): OutputStream = {
     val blockSize = conf.getSizeAsBytes("spark.io.compression.snappy.blockSize", "32k").toInt
+    Thread.sleep(10)
     new SnappyOutputStream(s, blockSize)
   }

I have a patch to fix this issue upstream in snappy-java. In the meantime, if we want to unbreak the Maven tests we can simply modify SnappyCompressionCodec to use the DefaultBufferAllocator instead of the caching one:

diff --git a/core/src/main/scala/org/apache/spark/io/CompressionCodec.scala b/core/src/main/scala/org/apache/spark/io/CompressionCodec.scala
index 0756cdb..e26db56 100644
--- a/core/src/main/scala/org/apache/spark/io/CompressionCodec.scala
+++ b/core/src/main/scala/org/apache/spark/io/CompressionCodec.scala
@@ -21,6 +21,7 @@ import java.io.{InputStream, OutputStream}

 import com.ning.compress.lzf.{LZFInputStream, LZFOutputStream}
 import net.jpountz.lz4.{LZ4BlockInputStream, LZ4BlockOutputStream}
+import org.xerial.snappy.buffer.DefaultBufferAllocator
 import org.xerial.snappy.{Snappy, SnappyInputStream, SnappyOutputStream}

 import org.apache.spark.SparkConf
@@ -154,7 +155,7 @@ class SnappyCompressionCodec(conf: SparkConf) extends CompressionCodec {

   override def compressedOutputStream(s: OutputStream): OutputStream = {
     val blockSize = conf.getSizeAsBytes("spark.io.compression.snappy.blockSize", "32k").toInt
-    new SnappyOutputStream(s, blockSize)
+    new SnappyOutputStream(s, blockSize, DefaultBufferAllocator.factory)
   }

[JoshRosen]

Also, just to underscore that the JavaAPISuite bug is nondeterministic, note that it's now passing again on one of the Maven builds even though we didn't change any related code: https://amplab.cs.berkeley.edu/jenkins/job/Spark-Master-Maven-pre-YARN/2350/hadoop.version=2.0.0-mr1-cdh4.1.2,label=centos/