[SPARK-19112][CORE] Support for ZStandard codec

Author: Sital Kedia

core/pom.xml

@@ -193,6 +193,10 @@
       <groupId>net.jpountz.lz4</groupId>
       <artifactId>lz4</artifactId>
     </dependency>
+    <dependency>
+      <groupId>com.github.luben</groupId>
+      <artifactId>zstd-jni</artifactId>
+    </dependency>
     <dependency>
       <groupId>org.roaringbitmap</groupId>
       <artifactId>RoaringBitmap</artifactId>

core/src/main/scala/org/apache/spark/io/CompressionCodec.scala

@@ -20,6 +20,7 @@ package org.apache.spark.io
 import java.io._
 import java.util.Locale
 
+import com.github.luben.zstd.{ZstdInputStream, ZstdOutputStream}
 import com.ning.compress.lzf.{LZFInputStream, LZFOutputStream}
 import net.jpountz.lz4.LZ4BlockOutputStream
 import org.xerial.snappy.{Snappy, SnappyInputStream, SnappyOutputStream}
@@ -50,13 +51,14 @@ private[spark] object CompressionCodec {
 
   private[spark] def supportsConcatenationOfSerializedStreams(codec: CompressionCodec): Boolean = {
     (codec.isInstanceOf[SnappyCompressionCodec] || codec.isInstanceOf[LZFCompressionCodec]
-      || codec.isInstanceOf[LZ4CompressionCodec])
+      || codec.isInstanceOf[LZ4CompressionCodec] || codec.isInstanceOf[ZStandardCompressionCodec])
   }
 
   private val shortCompressionCodecNames = Map(
     "lz4" -> classOf[LZ4CompressionCodec].getName,
     "lzf" -> classOf[LZFCompressionCodec].getName,
-    "snappy" -> classOf[SnappyCompressionCodec].getName)
+    "snappy" -> classOf[SnappyCompressionCodec].getName,
+    "zstd" -> classOf[SnappyCompressionCodec].getName)
 
   def getCodecName(conf: SparkConf): String = {
     conf.get(configKey, DEFAULT_COMPRESSION_CODEC)
@@ -216,3 +218,30 @@ private final class SnappyOutputStreamWrapper(os: SnappyOutputStream) extends Ou
     }
   }
 }
+
+/**
+ * :: DeveloperApi ::
+ * ZStandard implementation of [[org.apache.spark.io.CompressionCodec]].
+ *
+ * @note The wire protocol for this codec is not guaranteed to be compatible across versions
+ * of Spark. This is intended for use as an internal compression utility within a single Spark
+ * application.
+ */
+@DeveloperApi
+class ZStandardCompressionCodec(conf: SparkConf) extends CompressionCodec {
+
+  override def compressedOutputStream(s: OutputStream): OutputStream = {
+    val level = conf.getSizeAsBytes("spark.io.compression.zstandard.level", "1").toInt
+    val compressionBuffer = conf.getSizeAsBytes("spark.io.compression.lz4.blockSize", "32k").toInt
+    // Wrap the zstd output stream in a buffered output stream, so that we can
+    // avoid overhead excessive of JNI call while trying to compress small amount of data.
+    new BufferedOutputStream(new ZstdOutputStream(s, level), compressionBuffer)
+  }
+
+  override def compressedInputStream(s: InputStream): InputStream = {
+    val compressionBuffer = conf.getSizeAsBytes("spark.io.compression.lz4.blockSize", "32k").toInt
+    // Wrap the zstd input stream in a buffered input stream so that we can
+    // avoid overhead excessive of JNI call while trying to uncompress small amount of data.
+    new BufferedInputStream(new ZstdInputStream(s), compressionBuffer)
+  }
+}

core/src/test/scala/org/apache/spark/io/CompressionCodecSuite.scala

@@ -104,6 +104,24 @@ class CompressionCodecSuite extends SparkFunSuite {
     testConcatenationOfSerializedStreams(codec)
   }
 
+  test("zstd compression codec") {
+    val codec = CompressionCodec.createCodec(conf, classOf[ZStandardCompressionCodec].getName)
+    assert(codec.getClass === classOf[ZStandardCompressionCodec])
+    testCodec(codec)
+  }
+
+  test("zstd compression codec short form") {
+    val codec = CompressionCodec.createCodec(conf, "zstd")
+    assert(codec.getClass === classOf[ZStandardCompressionCodec])
+    testCodec(codec)
+  }
+
+  test("zstd supports concatenation of serialized zstd") {
+    val codec = CompressionCodec.createCodec(conf, classOf[ZStandardCompressionCodec].getName)
+    assert(codec.getClass === classOf[ZStandardCompressionCodec])
+    testConcatenationOfSerializedStreams(codec)
+  }
+
   test("bad compression codec") {
     intercept[IllegalArgumentException] {
       CompressionCodec.createCodec(conf, "foobar")

pom.xml

@@ -534,6 +534,11 @@
         <artifactId>lz4</artifactId>
         <version>1.3.0</version>
       </dependency>
+      <dependency>
+        <groupId>com.github.luben</groupId>
+        <artifactId>zstd-jni</artifactId>
+        <version>1.3.0-1</version>
+      </dependency>
       <dependency>
         <groupId>com.clearspring.analytics</groupId>
         <artifactId>stream</artifactId>