github.com/grailbio/base@v0.0.11/recordio/recordioutil/compress.go (about)

     1  // Copyright 2017 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache-2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package recordioutil
     6  
     7  import (
     8  	"bytes"
     9  	"io"
    10  
    11  	"github.com/klauspost/compress/flate"
    12  )
    13  
    14  // FlateTransform represents a 'transform' that can be used with
    15  // recordio.PackedWriter and Scanner to compress/decompress items written to a
    16  // single record.
    17  type FlateTransform struct {
    18  	level       int
    19  	passthrough bool
    20  }
    21  
    22  func sizeof(bufs [][]byte) int {
    23  	size := 0
    24  	for _, b := range bufs {
    25  		size += len(b)
    26  	}
    27  	return size
    28  }
    29  
    30  // NewFlateTransform creates a new Flate instance for use with
    31  // recordio.PackedWriter and PackedScanner. Level indicates the compression
    32  // level to use as per the flate package's contstants. The chosen level
    33  // has no effect when decompressing.
    34  func NewFlateTransform(level int) *FlateTransform {
    35  	return &FlateTransform{level: level}
    36  }
    37  
    38  // CompressTransform is intended for use Recordio.PackedWriterOpts.Transform.
    39  func (f *FlateTransform) CompressTransform(bufs [][]byte) ([]byte, error) {
    40  	if f.passthrough {
    41  		return bytes.Join(bufs, nil), nil
    42  	}
    43  	size := sizeof(bufs)
    44  	out := bytes.NewBuffer(make([]byte, 0, size))
    45  	fwr, err := flate.NewWriter(out, f.level)
    46  	if err != nil {
    47  		return nil, err
    48  	}
    49  	for _, b := range bufs {
    50  		_, err := fwr.Write(b)
    51  		if err != nil {
    52  			return nil, err
    53  		}
    54  	}
    55  	if err := fwr.Close(); err != nil {
    56  		return nil, err
    57  	}
    58  	return out.Bytes(), nil
    59  }
    60  
    61  // DecompressTransform is intended for use Recordio.PackedScannerOpts.Transform.
    62  func (f *FlateTransform) DecompressTransform(buf []byte) ([]byte, error) {
    63  	if f.passthrough {
    64  		return buf, nil
    65  	}
    66  	// Guess the size of the decompressed buffer.
    67  	size := len(buf) * 2
    68  	out := bytes.NewBuffer(make([]byte, 0, size))
    69  	frd := flate.NewReader(bytes.NewBuffer(buf))
    70  	if _, err := io.Copy(out, frd); err != nil {
    71  		return nil, err
    72  	}
    73  	if err := frd.Close(); err != nil {
    74  		return nil, err
    75  	}
    76  	return out.Bytes(), nil
    77  }