block: generalize CompressionStats

Generalize `CompressionStats` so that it can be used with an arbitrary
number of settings.

Author: RaduBerinde

sstable/block/compression_stats.go

@@ -5,63 +5,83 @@
 package block
 
 import (
+	"cmp"
 	"fmt"
 	"iter"
+	"slices"
 	"strings"
 
 	"github.com/cockroachdb/pebble/internal/compression"
+	"github.com/cockroachdb/pebble/internal/invariants"
 )
 
-// CompressionStats collects compression statistics for a single file - the
-// total compressed and uncompressed sizes for each distinct compression.Setting
-// used.
+// CompressionStats collects compression statistics (either for a single file or
+// for a collection of files).
+//
+// Compression statistics consist of the total compressed and uncompressed sizes for
+// each distinct compression.Setting used.
 type CompressionStats struct {
-	n int
-	// Compression profiles have three settings (data, value, other) and
-	// NoCompression can also be used for data that didn't compress.
-	buf [4]CompressionStatsForSetting
+	// We inline common values to avoid allocating the map in most cases.
+
+	// Total number of bytes that are not compressed.
+	noCompressionBytes uint64
+	// Compression stats for fastestCompression.
+	fastest CompressionStatsForSetting
+
+	others map[compression.Setting]CompressionStatsForSetting
 }
 
 type CompressionStatsForSetting struct {
-	Setting           compression.Setting
-	UncompressedBytes uint64
 	CompressedBytes   uint64
+	UncompressedBytes uint64
+}
+
+func (cs *CompressionStatsForSetting) Add(other CompressionStatsForSetting) {
+	cs.CompressedBytes += other.CompressedBytes
+	cs.UncompressedBytes += other.UncompressedBytes
 }
 
 // add updates the stats to reflect a block that was compressed with the given setting.
-func (c *CompressionStats) add(
-	setting compression.Setting, sizeUncompressed, sizeCompressed uint64,
-) {
-	for i := 0; i < c.n; i++ {
-		if c.buf[i].Setting == setting {
-			c.buf[i].UncompressedBytes += sizeUncompressed
-			c.buf[i].CompressedBytes += sizeCompressed
-			return
+func (c *CompressionStats) add(setting compression.Setting, stats CompressionStatsForSetting) {
+	switch setting {
+	case compression.None:
+		c.noCompressionBytes += stats.UncompressedBytes
+		if invariants.Enabled && stats.UncompressedBytes != stats.CompressedBytes {
+			panic("invalid stats for no-compression")
 		}
+	case fastestCompression:
+		c.fastest.Add(stats)
+	default:
+		if c.others == nil {
+			c.others = make(map[compression.Setting]CompressionStatsForSetting, 2)
+		}
+		prev := c.others[setting]
+		prev.Add(stats)
+		c.others[setting] = prev
 	}
-	if c.n >= len(c.buf)-1 {
-		panic("too many compression settings")
-	}
-	c.buf[c.n] = CompressionStatsForSetting{
-		Setting:           setting,
-		UncompressedBytes: sizeUncompressed,
-		CompressedBytes:   sizeCompressed,
-	}
-	c.n++
 }
 
 // MergeWith updates the receiver stats to include the other stats.
 func (c *CompressionStats) MergeWith(other *CompressionStats) {
-	for i := 0; i < other.n; i++ {
-		c.add(other.buf[i].Setting, other.buf[i].UncompressedBytes, other.buf[i].CompressedBytes)
+	for s, cs := range other.All() {
+		c.add(s, cs)
 	}
 }
 
 // All returns an iterator over the collected stats, in arbitrary order.
-func (c CompressionStats) All() iter.Seq[CompressionStatsForSetting] {
-	return func(yield func(cs CompressionStatsForSetting) bool) {
-		for i := 0; i < c.n; i++ {
-			if !yield(c.buf[i]) {
+func (c *CompressionStats) All() iter.Seq2[compression.Setting, CompressionStatsForSetting] {
+	return func(yield func(s compression.Setting, cs CompressionStatsForSetting) bool) {
+		if c.noCompressionBytes != 0 && !yield(compression.None, CompressionStatsForSetting{
+			UncompressedBytes: c.noCompressionBytes,
+			CompressedBytes:   c.noCompressionBytes,
+		}) {
+			return
+		}
+		if c.fastest.UncompressedBytes != 0 && !yield(fastestCompression, c.fastest) {
+			return
+		}
+		for s, cs := range c.others {
+			if !yield(s, cs) {
 				return
 			}
 		}
@@ -70,14 +90,39 @@ func (c CompressionStats) All() iter.Seq[CompressionStatsForSetting] {
 
 // String returns a string representation of the stats, in the format:
 // "<setting1>:<compressed1>/<uncompressed1>,<setting2>:<compressed2>/<uncompressed2>,..."
+//
+// The settings are ordered alphabetically.
 func (c CompressionStats) String() string {
+	n := len(c.others)
+	if c.noCompressionBytes != 0 {
+		n++
+	}
+	if c.fastest.UncompressedBytes != 0 {
+		n++
+	}
+
+	type entry struct {
+		s  compression.Setting
+		cs CompressionStatsForSetting
+	}
+	entries := make([]entry, 0, n)
+	for s, cs := range c.All() {
+		entries = append(entries, entry{s, cs})
+	}
+	slices.SortFunc(entries, func(x, y entry) int {
+		if x.s.Algorithm != y.s.Algorithm {
+			return cmp.Compare(x.s.Algorithm.String(), y.s.Algorithm.String())
+		}
+		return cmp.Compare(x.s.Level, y.s.Level)
+	})
+
 	var buf strings.Builder
-	buf.Grow(c.n * 64)
-	for i := 0; i < c.n; i++ {
-		if i > 0 {
+	buf.Grow(n * 64)
+	for _, e := range entries {
+		if buf.Len() > 0 {
 			buf.WriteString(",")
 		}
-		fmt.Fprintf(&buf, "%s:%d/%d", c.buf[i].Setting.String(), c.buf[i].CompressedBytes, c.buf[i].UncompressedBytes)
+		fmt.Fprintf(&buf, "%s:%d/%d", e.s.String(), e.cs.CompressedBytes, e.cs.UncompressedBytes)
 	}
 	return buf.String()
 }

sstable/block/compressor.go

@@ -89,18 +89,24 @@ func (c *Compressor) Compress(dst, src []byte, kind Kind) (CompressionIndicator,
 	//      before
 	if setting.Algorithm != compression.NoCompression &&
 		int64(len(out))*100 > int64(len(src))*int64(100-c.minReductionPercent) {
-		c.stats.add(compression.None, uint64(len(src)), uint64(len(src)))
-		return NoCompressionIndicator, append(out[:0], src...)
+		setting.Algorithm = compression.NoCompression
+		out = append(out[:0], src...)
 	}
-	c.stats.add(setting, uint64(len(src)), uint64(len(out)))
+	c.stats.add(setting, CompressionStatsForSetting{
+		UncompressedBytes: uint64(len(src)),
+		CompressedBytes:   uint64(len(out)),
+	})
 	return compressionIndicatorFromAlgorithm(setting.Algorithm), out
 }
 
 // UncompressedBlock informs the compressor that a block of the given size and
 // kind was written uncompressed. This is used so that the final statistics are
 // complete.
 func (c *Compressor) UncompressedBlock(size int, kind Kind) {
-	c.stats.add(compression.None, uint64(size), uint64(size))
+	c.stats.add(compression.None, CompressionStatsForSetting{
+		UncompressedBytes: uint64(size),
+		CompressedBytes:   uint64(size),
+	})
 }
 
 // Stats returns the compression stats. The result can only be used until the

sstable/testdata/writer_value_blocks

@@ -670,7 +670,7 @@ sstable
  │    ├── 00630    rocksdb.top-level.index.size (24)
  │    ├── restart points
  │    │    └── 00654 [restart 0]
- │    └── trailer [compression=none checksum=0x3508516e]
+ │    └── trailer [compression=none checksum=0x88688345]
  ├── meta-index  offset: 1360  length: 64
  │    ├── 0000    pebble.value_index block:680/8 value-blocks-index-lengths: 1(num), 2(offset), 1(length) [restart]
  │    ├── 0027    rocksdb.properties block:693/662 [restart]
@@ -1247,7 +1247,7 @@ sstable
  │    ├── 00698    rocksdb.property.collectors (41)
  │    ├── 00739    rocksdb.raw.key.size (21)
  │    ├── 00760    rocksdb.raw.value.size (23)
- │    └── trailer [compression=snappy checksum=0xf941a001]
+ │    └── trailer [compression=snappy checksum=0x4ed1f050]
  ├── meta-index  offset: 908  length: 72
  │    ├── 0000    pebble.value_index block:257/3 value-blocks-index-lengths: 1(num), 1(offset), 1(length)
  │    │