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

     1  package compress_test
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"testing"
     7  
     8  	"github.com/segmentio/parquet-go/compress"
     9  	"github.com/segmentio/parquet-go/compress/brotli"
    10  	"github.com/segmentio/parquet-go/compress/gzip"
    11  	"github.com/segmentio/parquet-go/compress/lz4"
    12  	"github.com/segmentio/parquet-go/compress/snappy"
    13  	"github.com/segmentio/parquet-go/compress/uncompressed"
    14  	"github.com/segmentio/parquet-go/compress/zstd"
    15  )
    16  
    17  var tests = [...]struct {
    18  	scenario string
    19  	codec    compress.Codec
    20  }{
    21  	{
    22  		scenario: "uncompressed",
    23  		codec:    new(uncompressed.Codec),
    24  	},
    25  
    26  	{
    27  		scenario: "snappy",
    28  		codec:    new(snappy.Codec),
    29  	},
    30  
    31  	{
    32  		scenario: "gzip",
    33  		codec:    new(gzip.Codec),
    34  	},
    35  
    36  	{
    37  		scenario: "brotli",
    38  		codec:    new(brotli.Codec),
    39  	},
    40  
    41  	{
    42  		scenario: "zstd",
    43  		codec:    new(zstd.Codec),
    44  	},
    45  
    46  	{
    47  		scenario: "lz4",
    48  		codec:    new(lz4.Codec),
    49  	},
    50  }
    51  
    52  var testdata = bytes.Repeat([]byte("1234567890qwertyuiopasdfghjklzxcvbnm"), 10e3)
    53  
    54  func TestCompressionCodec(t *testing.T) {
    55  	buffer := make([]byte, 0, len(testdata))
    56  	output := make([]byte, 0, len(testdata))
    57  
    58  	for _, test := range tests {
    59  		t.Run(test.scenario, func(t *testing.T) {
    60  			const N = 10
    61  			// Run the test multiple times to exercise codecs that maintain
    62  			// state across compression/decompression.
    63  			for i := 0; i < N; i++ {
    64  				var err error
    65  
    66  				buffer, err = test.codec.Encode(buffer[:0], testdata)
    67  				if err != nil {
    68  					t.Fatal(err)
    69  				}
    70  
    71  				output, err = test.codec.Decode(output[:0], buffer)
    72  				if err != nil {
    73  					t.Fatal(err)
    74  				}
    75  
    76  				if !bytes.Equal(testdata, output) {
    77  					t.Errorf("content mismatch after compressing and decompressing (attempt %d/%d)", i+1, N)
    78  				}
    79  			}
    80  		})
    81  	}
    82  }
    83  
    84  func BenchmarkEncode(b *testing.B) {
    85  	buffer := make([]byte, 0, len(testdata))
    86  
    87  	for _, test := range tests {
    88  		b.Run(test.scenario, func(b *testing.B) {
    89  			b.SetBytes(int64(len(testdata)))
    90  			benchmarkZeroAllocsPerRun(b, func() {
    91  				buffer, _ = test.codec.Encode(buffer[:0], testdata)
    92  			})
    93  		})
    94  	}
    95  }
    96  
    97  func BenchmarkDecode(b *testing.B) {
    98  	buffer := make([]byte, 0, len(testdata))
    99  	output := make([]byte, 0, len(testdata))
   100  
   101  	for _, test := range tests {
   102  		b.Run(test.scenario, func(b *testing.B) {
   103  			buffer, _ = test.codec.Encode(buffer[:0], testdata)
   104  			b.SetBytes(int64(len(testdata)))
   105  			benchmarkZeroAllocsPerRun(b, func() {
   106  				output, _ = test.codec.Encode(output[:0], buffer)
   107  			})
   108  		})
   109  	}
   110  }
   111  
   112  type simpleReader struct{ io.Reader }
   113  
   114  func (s *simpleReader) Close() error            { return nil }
   115  func (s *simpleReader) Reset(r io.Reader) error { s.Reader = r; return nil }
   116  
   117  type simpleWriter struct{ io.Writer }
   118  
   119  func (s *simpleWriter) Close() error      { return nil }
   120  func (s *simpleWriter) Reset(w io.Writer) { s.Writer = w }
   121  
   122  func BenchmarkCompressor(b *testing.B) {
   123  	compressor := compress.Compressor{}
   124  	src := make([]byte, 1000)
   125  	dst := make([]byte, 1000)
   126  
   127  	benchmarkZeroAllocsPerRun(b, func() {
   128  		dst, _ = compressor.Encode(dst, src, func(w io.Writer) (compress.Writer, error) {
   129  			return &simpleWriter{Writer: w}, nil
   130  		})
   131  	})
   132  }
   133  
   134  func BenchmarkDecompressor(b *testing.B) {
   135  	decompressor := compress.Decompressor{}
   136  	src := make([]byte, 1000)
   137  	dst := make([]byte, 1000)
   138  
   139  	benchmarkZeroAllocsPerRun(b, func() {
   140  		dst, _ = decompressor.Decode(dst, src, func(r io.Reader) (compress.Reader, error) {
   141  			return &simpleReader{Reader: r}, nil
   142  		})
   143  	})
   144  }
   145  
   146  func benchmarkZeroAllocsPerRun(b *testing.B, f func()) {
   147  	if allocs := testing.AllocsPerRun(b.N, f); allocs != 0 && !testing.Short() {
   148  		b.Errorf("too many memory allocations: %g > 0", allocs)
   149  	}
   150  }