zstd: Big speedup on small dictionary encodes

zstd/dict_test.go

@@ -218,7 +218,7 @@ func TestEncoder_SmallDict(t *testing.T) {
 	}
 }
 
-func BenchmarkEncodeAllDict(b *testing.B) {
+func benchmarkEncodeAllLimitedBySize(b *testing.B, lowerLimit int, upperLimit int) {
 	fn := "testdata/dict-tests-small.zip"
 	data, err := ioutil.ReadFile(fn)
 	t := testing.TB(b)
@@ -232,7 +232,6 @@ func BenchmarkEncodeAllDict(b *testing.B) {
 	}
 	var dicts [][]byte
 	var encs []*Encoder
-	var noDictEncs []*Encoder
 	var encNames []string
 
 	for _, tt := range zr.File {
@@ -251,32 +250,25 @@ func BenchmarkEncodeAllDict(b *testing.B) {
 			}
 			dicts = append(dicts, in)
 			for level := SpeedFastest; level < speedLast; level++ {
-				enc, err := NewWriter(nil, WithEncoderConcurrency(1), WithEncoderDict(in), WithEncoderLevel(level), WithWindowSize(1<<17))
+				enc, err := NewWriter(nil, WithEncoderDict(in), WithEncoderLevel(level))
 				if err != nil {
 					t.Fatal(err)
 				}
 				encs = append(encs, enc)
 				encNames = append(encNames, fmt.Sprint("level-", level.String(), "-dict-", len(dicts)))
-
-				enc, err = NewWriter(nil, WithEncoderConcurrency(1), WithEncoderLevel(level), WithWindowSize(1<<17))
-				if err != nil {
-					t.Fatal(err)
-				}
-				noDictEncs = append(noDictEncs, enc)
 			}
 		}()
 	}
+	const nPer = int(speedLast - SpeedFastest)
 	dec, err := NewReader(nil, WithDecoderConcurrency(1), WithDecoderDicts(dicts...))
 	if err != nil {
 		t.Fatal(err)
 		return
 	}
 	defer dec.Close()
 
-	for i, tt := range zr.File {
-		if i == 5 {
-			break
-		}
+	tested := make(map[int]struct{})
+	for j, tt := range zr.File {
 		if !strings.HasSuffix(tt.Name, ".zst") {
 			continue
 		}
@@ -293,26 +285,64 @@ func BenchmarkEncodeAllDict(b *testing.B) {
 		if err != nil {
 			t.Fatal(err)
 		}
+
+		// Only test each size once
+		if _, ok := tested[len(decoded)]; ok {
+			continue
+		}
+		tested[len(decoded)] = struct{}{}
+
+		if len(decoded) < lowerLimit {
+			continue
+		}
+
+		if upperLimit > 0 && len(decoded) > upperLimit {
+			continue
+		}
+
 		for i := range encs {
-			// Only do 1 dict (3 encoders) for now.
-			if i == 3 {
+			// Only do 1 dict (4 encoders) for now.
+			if i == nPer-1 {
 				break
 			}
 			// Attempt to compress with all dicts
-			var dst []byte
-			enc := encs[i]
-			b.Run(fmt.Sprintf("length-%d-%s", len(decoded), encNames[i]), func(b *testing.B) {
-				b.SetBytes(int64(len(decoded)))
-				b.ResetTimer()
-				b.ReportAllocs()
-				for i := 0; i < b.N; i++ {
-					dst = enc.EncodeAll(decoded, dst[:0])
-				}
+			encIdx := (i + j*nPer) % len(encs)
+			enc := encs[encIdx]
+			b.Run(fmt.Sprintf("length-%d-%s", len(decoded), encNames[encIdx]), func(b *testing.B) {
+				b.RunParallel(func(pb *testing.PB) {
+					dst := make([]byte, 0, len(decoded)+10)
+					b.SetBytes(int64(len(decoded)))
+					b.ResetTimer()
+					b.ReportAllocs()
+					for pb.Next() {
+						dst = enc.EncodeAll(decoded, dst[:0])
+					}
+				})
 			})
 		}
 	}
 }
 
+func BenchmarkEncodeAllDict0_1024(b *testing.B) {
+	benchmarkEncodeAllLimitedBySize(b, 0, 1024)
+}
+
+func BenchmarkEncodeAllDict1024_8192(b *testing.B) {
+	benchmarkEncodeAllLimitedBySize(b, 1024, 8192)
+}
+
+func BenchmarkEncodeAllDict8192_16384(b *testing.B) {
+	benchmarkEncodeAllLimitedBySize(b, 8192, 16384)
+}
+
+func BenchmarkEncodeAllDict16384_65536(b *testing.B) {
+	benchmarkEncodeAllLimitedBySize(b, 16384, 65536)
+}
+
+func BenchmarkEncodeAllDict65536_0(b *testing.B) {
+	benchmarkEncodeAllLimitedBySize(b, 65536, 0)
+}
+
 func TestDecoder_MoreDicts(t *testing.T) {
 	// All files have CRC
 	// https://files.klauspost.com/compress/zstd-dict-tests.zip

zstd/enc_base.go

@@ -7,6 +7,10 @@ import (
 	"github.com/klauspost/compress/zstd/internal/xxhash"
 )
 
+const (
+	dictShardBits = 6
+)
+
 type fastBase struct {
 	// cur is the offset at the start of hist
 	cur int32