Skip to content

Commit

Permalink
zstd: Add RLE detection+encoding
Browse files Browse the repository at this point in the history
Add full block RLE encoding.

Fastest+Default: The first block in RLE will be detected automatically, but following will likely be backreferences.

Better+Best: Explicit RLE block detection. Will always be encoded as RLE blocks.

Full block RLE are very uncommon and the gains are small, but can happen in very sparse data, so makes sense for "stronger" settings.

Add feature #937
  • Loading branch information
klauspost committed Mar 8, 2024
1 parent 165be36 commit 0b60023
Show file tree
Hide file tree
Showing 4 changed files with 90 additions and 1 deletion.
20 changes: 20 additions & 0 deletions zstd/blockenc.go
Original file line number Diff line number Diff line change
Expand Up @@ -427,6 +427,16 @@ func (b *blockEnc) encodeLits(lits []byte, raw bool) error {
return nil
}

// encodeRLE will encode an RLE block.
func (b *blockEnc) encodeRLE(val byte, length uint32) {
var bh blockHeader
bh.setLast(b.last)
bh.setSize(length)
bh.setType(blockTypeRLE)
b.output = bh.appendTo(b.output)
b.output = append(b.output, val)
}

// fuzzFseEncoder can be used to fuzz the FSE encoder.
func fuzzFseEncoder(data []byte) int {
if len(data) > maxSequences || len(data) < 2 {
Expand Down Expand Up @@ -479,6 +489,16 @@ func (b *blockEnc) encode(org []byte, raw, rawAllLits bool) error {
if len(b.sequences) == 0 {
return b.encodeLits(b.literals, rawAllLits)
}
if len(b.sequences) == 1 && len(org) > 0 && len(b.literals) <= 1 {
// Check common RLE cases.
seq := b.sequences[0]
if seq.litLen == uint32(len(b.literals)) && seq.offset-3 == 1 {
// Offset == 1 and 0 or 1 literals.
b.encodeRLE(org[0], b.sequences[0].matchLen+zstdMinMatch+seq.litLen)
return nil
}
}

// We want some difference to at least account for the headers.
saved := b.size - len(b.literals) - (b.size >> 6)
if saved < 16 {
Expand Down
12 changes: 12 additions & 0 deletions zstd/enc_best.go
Original file line number Diff line number Diff line change
Expand Up @@ -135,8 +135,20 @@ func (e *bestFastEncoder) Encode(blk *blockEnc, src []byte) {
break
}

// Add block to history
s := e.addBlock(src)
blk.size = len(src)

// Check RLE first
if len(src) > zstdMinMatch {
ml := matchLen(src[1:], src)
if ml == len(src)-1 {
blk.literals = append(blk.literals, src[0])
blk.sequences = append(blk.sequences, seq{litLen: 1, matchLen: uint32(len(src)-1) - zstdMinMatch, offset: 1 + 3})
return
}
}

if len(src) < minNonLiteralBlockSize {
blk.extraLits = len(src)
blk.literals = blk.literals[:len(src)]
Expand Down
13 changes: 12 additions & 1 deletion zstd/enc_better.go
Original file line number Diff line number Diff line change
Expand Up @@ -102,9 +102,20 @@ func (e *betterFastEncoder) Encode(blk *blockEnc, src []byte) {
e.cur = e.maxMatchOff
break
}

// Add block to history
s := e.addBlock(src)
blk.size = len(src)

// Check RLE first
if len(src) > zstdMinMatch {
ml := matchLen(src[1:], src)
if ml == len(src)-1 {
blk.literals = append(blk.literals, src[0])
blk.sequences = append(blk.sequences, seq{litLen: 1, matchLen: uint32(len(src)-1) - zstdMinMatch, offset: 1 + 3})
return
}
}

if len(src) < minNonLiteralBlockSize {
blk.extraLits = len(src)
blk.literals = blk.literals[:len(src)]
Expand Down
46 changes: 46 additions & 0 deletions zstd/encoder_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -342,6 +342,52 @@ func TestEncoder_EncodeAllTwain(t *testing.T) {
}
}

func TestEncoder_EncodeRLE(t *testing.T) {
in := make([]byte, 1<<20)
testWindowSizes := testWindowSizes
if testing.Short() {
testWindowSizes = []int{1 << 20}
}

dec, err := NewReader(nil)
if err != nil {
t.Fatal(err)
}
defer dec.Close()

for level := speedNotSet + 1; level < speedLast; level++ {
t.Run(level.String(), func(t *testing.T) {
if isRaceTest && level >= SpeedBestCompression {
t.SkipNow()
}
for _, windowSize := range testWindowSizes {
t.Run(fmt.Sprintf("window:%d", windowSize), func(t *testing.T) {
e, err := NewWriter(nil, WithEncoderLevel(level), WithWindowSize(windowSize))
if err != nil {
t.Fatal(err)
}
defer e.Close()
start := time.Now()
dst := e.EncodeAll(in, nil)
t.Log("Simple Encoder len", len(in), "-> zstd len", len(dst))
mbpersec := (float64(len(in)) / (1024 * 1024)) / (float64(time.Since(start)) / (float64(time.Second)))
t.Logf("Encoded %d bytes with %.2f MB/s", len(in), mbpersec)

decoded, err := dec.DecodeAll(dst, nil)
if err != nil {
t.Error(err, len(decoded))
}
if !bytes.Equal(decoded, in) {
os.WriteFile("testdata/"+t.Name()+"-Mark.Twain-Tom.Sawyer.txt.got", decoded, os.ModePerm)
t.Fatal("Decoded does not match")
}
t.Log("Encoded content matched")
})
}
})
}
}

func TestEncoder_EncodeAllPi(t *testing.T) {
in, err := os.ReadFile("../testdata/pi.txt")
if err != nil {
Expand Down

0 comments on commit 0b60023

Please sign in to comment.