Merge pull request #1014 from segmentio/snappy-configurable-level

compress/snappy: configurable compression level

Author: chriso

compress/snappy/snappy.go

@@ -4,6 +4,7 @@ import (
 	"io"
 	"sync"
 
+	"github.com/klauspost/compress/s2"
 	"github.com/klauspost/compress/snappy"
 )
 
@@ -16,6 +17,16 @@ const (
 	Unframed
 )
 
+// Compression level.
+type Compression int
+
+const (
+	DefaultCompression Compression = iota
+	FasterCompression
+	BetterCompression
+	BestCompression
+)
+
 var (
 	readerPool sync.Pool
 	writerPool sync.Pool
@@ -28,6 +39,9 @@ type Codec struct {
 	//
 	// Default to Framed.
 	Framing Framing
+
+	// Compression level.
+	Compression Compression
 }
 
 // Code implements the compress.Codec interface.
@@ -56,12 +70,19 @@ func (c *Codec) NewWriter(w io.Writer) io.WriteCloser {
 	if x != nil {
 		x.Reset(w)
 	} else {
-		x = &xerialWriter{
-			writer: w,
-			encode: snappy.Encode,
-		}
+		x = &xerialWriter{writer: w}
 	}
 	x.framed = c.Framing == Framed
+	switch c.Compression {
+	case FasterCompression:
+		x.encode = s2.EncodeSnappy
+	case BetterCompression:
+		x.encode = s2.EncodeSnappyBetter
+	case BestCompression:
+		x.encode = s2.EncodeSnappyBest
+	default:
+		x.encode = snappy.Encode // aka. s2.EncodeSnappyBetter
+	}
 	return &writer{xerialWriter: x}
 }