add s2 and snappy and gzip algorithm for compress (#2)

* add s2 and snappy algorithm for compress

* add zip algorithm for compress

* update go mod

Author: hengfeiyang

compress/compress.go

@@ -0,0 +1,59 @@
+/* Copyright 2022 Zinc Labs Inc. and Contributors
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+*     http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+ */
+
+package compress
+
+import (
+	"errors"
+)
+
+const (
+	SNAPPY = iota
+	S2
+	ZSTD
+	GZIP
+)
+
+var Algorithm = S2
+
+func Decompress(dst, src []byte) ([]byte, error) {
+	switch Algorithm {
+	case SNAPPY:
+		return SnappyDecompress(dst, src)
+	case S2:
+		return S2Decompress(dst, src)
+	case ZSTD:
+		return ZSTDDecompress(dst, src)
+	case GZIP:
+		return GzipDecompress(dst, src)
+	default:
+		return nil, errors.New("unknown compress algorithm")
+	}
+}
+
+func Compress(dst, src []byte) ([]byte, error) {
+	switch Algorithm {
+	case SNAPPY:
+		return SnappyCompress(dst, src)
+	case S2:
+		return S2Compress(dst, src)
+	case ZSTD:
+		return ZSTDCompress(dst, src, ZSTDCompressionLevel)
+	case GZIP:
+		return GzipCompress(dst, src)
+	default:
+		return nil, errors.New("unknown compress algorithm")
+	}
+}

compress/gzip.go

@@ -0,0 +1,47 @@
+/* Copyright 2022 Zinc Labs Inc. and Contributors
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+*     http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+ */
+
+package compress
+
+import (
+	"bytes"
+	"compress/gzip"
+)
+
+// GzipDecompress decompresses a block using gzip algorithm.
+func GzipDecompress(dst, src []byte) ([]byte, error) {
+	buf := bytes.NewBuffer(src)
+	r, err := gzip.NewReader(buf)
+	if err != nil {
+		return nil, err
+	}
+	out := bytes.NewBuffer(dst[:0])
+	out.ReadFrom(r)
+	return out.Bytes(), nil
+}
+
+// GzipCompress compresses a block using gzip algorithm.
+func GzipCompress(dst, src []byte) ([]byte, error) {
+	buf := bytes.NewBuffer(dst[:0])
+	w := gzip.NewWriter(buf)
+	_, err := w.Write(src)
+	if err != nil {
+		return nil, err
+	}
+	if err = w.Close(); err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}

compress/s2.go

@@ -0,0 +1,30 @@
+/* Copyright 2022 Zinc Labs Inc. and Contributors
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+*     http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+ */
+
+package compress
+
+import (
+	s2c "github.com/klauspost/compress/s2"
+)
+
+// S2Decompress decompresses a block using s2 algorithm.
+func S2Decompress(dst, src []byte) ([]byte, error) {
+	return s2c.Decode(dst, src)
+}
+
+// S2Compress compresses a block using s2 algorithm.
+func S2Compress(dst, src []byte) ([]byte, error) {
+	return s2c.EncodeBetter(dst, src), nil
+}

compress/snappy.go

@@ -0,0 +1,28 @@
+/* Copyright 2022 Zinc Labs Inc. and Contributors
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+*     http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+ */
+
+package compress
+
+import "github.com/klauspost/compress/snappy"
+
+// SnappyDecompress decompresses a block using snappy algorithm.
+func SnappyDecompress(dst, src []byte) ([]byte, error) {
+	return snappy.Decode(dst, src)
+}
+
+// SnappyCompress compresses a block using snappy algorithm.
+func SnappyCompress(dst, src []byte) ([]byte, error) {
+	return snappy.Encode(dst, src), nil
+}

zstd.go renamed to compress/zstd.go

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package ice
+package compress
 
 import (
 	"log"
@@ -23,7 +23,7 @@ import (
 	"github.com/klauspost/compress/zstd"
 )
 
-const ZSTDCompressionLevel = 3 // 1, 3, 9
+var ZSTDCompressionLevel = 3 // 1, 3, 9
 
 var (
 	decoder *zstd.Decoder

contentcoder.go

@@ -18,6 +18,8 @@ import (
 	"bytes"
 	"encoding/binary"
 	"io"
+
+	"github.com/blugelabs/ice/compress"
 )
 
 var termSeparator byte = 0xff
@@ -120,7 +122,7 @@ func (c *chunkedContentCoder) flushContents() error {
 	metaData := c.chunkMetaBuf.Bytes()
 	c.final = append(c.final, c.chunkMetaBuf.Bytes()...)
 	// write the compressed data to the final data
-	c.compressed, err = ZSTDCompress(c.compressed[:cap(c.compressed)], c.chunkBuf.Bytes(), ZSTDCompressionLevel)
+	c.compressed, err = compress.Compress(c.compressed[:cap(c.compressed)], c.chunkBuf.Bytes())
 	if err != nil {
 		return err
 	}

documentcoder.go

@@ -4,6 +4,8 @@ import (
 	"bytes"
 	"encoding/binary"
 	"io"
+
+	"github.com/blugelabs/ice/compress"
 )
 
 const defaultDocumentChunkSize uint32 = 128
@@ -70,7 +72,7 @@ func (c *chunkedDocumentCoder) newLine() error {
 func (c *chunkedDocumentCoder) flush() error {
 	if c.buf.Len() > 0 {
 		var err error
-		c.compressed, err = ZSTDCompress(c.compressed[:cap(c.compressed)], c.buf.Bytes(), ZSTDCompressionLevel)
+		c.compressed, err = compress.Compress(c.compressed[:cap(c.compressed)], c.buf.Bytes())
 		if err != nil {
 			return err
 		}

docvalues.go

@@ -22,6 +22,7 @@ import (
 	"sort"
 
 	segment "github.com/blugelabs/bluge_segment_api"
+	"github.com/blugelabs/ice/compress"
 )
 
 type docNumTermsVisitor func(docNum uint64, terms []byte) error
@@ -209,7 +210,7 @@ func (di *docValueReader) iterateAllDocValues(s *Segment, visitor docNumTermsVis
 		}
 
 		// uncompress the already loaded data
-		uncompressed, err := ZSTDDecompress(di.uncompressed[:cap(di.uncompressed)], di.curChunkData)
+		uncompressed, err := compress.Decompress(di.uncompressed[:cap(di.uncompressed)], di.curChunkData)
 		if err != nil {
 			return err
 		}
@@ -244,7 +245,7 @@ func (di *docValueReader) visitDocValues(docNum uint64,
 		uncompressed = di.uncompressed
 	} else {
 		// uncompress the already loaded data
-		uncompressed, err = ZSTDDecompress(di.uncompressed[:cap(di.uncompressed)], di.curChunkData)
+		uncompressed, err = compress.Decompress(di.uncompressed[:cap(di.uncompressed)], di.curChunkData)
 		if err != nil {
 			return err
 		}

intcoder.go

@@ -18,6 +18,8 @@ import (
 	"bytes"
 	"encoding/binary"
 	"io"
+
+	"github.com/blugelabs/ice/compress"
 )
 
 // We can safely use 0 to represent termNotEncoded since 0
@@ -104,7 +106,7 @@ func (c *chunkedIntCoder) Add(docNum uint64, vals ...uint64) error {
 // to be encoded.
 func (c *chunkedIntCoder) Close() error {
 	var err error
-	c.compressed, err = ZSTDCompress(c.compressed[:cap(c.compressed)], c.chunkBuf.Bytes(), ZSTDCompressionLevel)
+	c.compressed, err = compress.Compress(c.compressed[:cap(c.compressed)], c.chunkBuf.Bytes())
 	if err != nil {
 		return err
 	}

intdecoder.go

@@ -19,6 +19,7 @@ import (
 	"fmt"
 
 	segment "github.com/blugelabs/bluge_segment_api"
+	"github.com/blugelabs/ice/compress"
 )
 
 type chunkedIntDecoder struct {
@@ -87,7 +88,10 @@ func (d *chunkedIntDecoder) loadChunk(chunk int) error {
 	if err != nil {
 		return err
 	}
-	d.uncompressed, err = ZSTDDecompress(d.uncompressed[:cap(d.uncompressed)], curChunkBytesData)
+	if len(curChunkBytesData) == 0 {
+		return nil
+	}
+	d.uncompressed, err = compress.Decompress(d.uncompressed[:cap(d.uncompressed)], curChunkBytesData)
 	if err != nil {
 		return err
 	}