github.com/segmentio/parquet-go@v0.0.0-20230712180008-5d42db8f0d47/compress/zstd/zstd.go (about)

     1  // Package zstd implements the ZSTD parquet compression codec.
     2  package zstd
     3  
     4  import (
     5  	"sync"
     6  
     7  	"github.com/klauspost/compress/zstd"
     8  	"github.com/segmentio/parquet-go/format"
     9  )
    10  
    11  type Level = zstd.EncoderLevel
    12  
    13  const (
    14  	// SpeedFastest will choose the fastest reasonable compression.
    15  	// This is roughly equivalent to the fastest Zstandard mode.
    16  	SpeedFastest = zstd.SpeedFastest
    17  
    18  	// SpeedDefault is the default "pretty fast" compression option.
    19  	// This is roughly equivalent to the default Zstandard mode (level 3).
    20  	SpeedDefault = zstd.SpeedDefault
    21  
    22  	// SpeedBetterCompression will yield better compression than the default.
    23  	// Currently it is about zstd level 7-8 with ~ 2x-3x the default CPU usage.
    24  	// By using this, notice that CPU usage may go up in the future.
    25  	SpeedBetterCompression = zstd.SpeedBetterCompression
    26  
    27  	// SpeedBestCompression will choose the best available compression option.
    28  	// This will offer the best compression no matter the CPU cost.
    29  	SpeedBestCompression = zstd.SpeedBestCompression
    30  )
    31  
    32  const (
    33  	DefaultLevel = SpeedDefault
    34  )
    35  
    36  type Codec struct {
    37  	Level Level
    38  
    39  	encoders sync.Pool // *zstd.Encoder
    40  	decoders sync.Pool // *zstd.Decoder
    41  }
    42  
    43  func (c *Codec) String() string {
    44  	return "ZSTD"
    45  }
    46  
    47  func (c *Codec) CompressionCodec() format.CompressionCodec {
    48  	return format.Zstd
    49  }
    50  
    51  func (c *Codec) Encode(dst, src []byte) ([]byte, error) {
    52  	e, _ := c.encoders.Get().(*zstd.Encoder)
    53  	if e == nil {
    54  		var err error
    55  		e, err = zstd.NewWriter(nil,
    56  			zstd.WithEncoderConcurrency(1),
    57  			zstd.WithEncoderLevel(c.level()),
    58  			zstd.WithZeroFrames(true),
    59  			zstd.WithEncoderCRC(false),
    60  		)
    61  		if err != nil {
    62  			return dst[:0], err
    63  		}
    64  	}
    65  	defer c.encoders.Put(e)
    66  	return e.EncodeAll(src, dst[:0]), nil
    67  }
    68  
    69  func (c *Codec) Decode(dst, src []byte) ([]byte, error) {
    70  	d, _ := c.decoders.Get().(*zstd.Decoder)
    71  	if d == nil {
    72  		var err error
    73  		d, err = zstd.NewReader(nil,
    74  			zstd.WithDecoderConcurrency(1),
    75  		)
    76  		if err != nil {
    77  			return dst[:0], err
    78  		}
    79  	}
    80  	defer c.decoders.Put(d)
    81  	return d.DecodeAll(src, dst[:0])
    82  }
    83  
    84  func (c *Codec) level() Level {
    85  	if c.Level != 0 {
    86  		return c.Level
    87  	}
    88  	return DefaultLevel
    89  }