github.com/grailbio/base@v0.0.11/recordio/recordioutil/rioutil.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  	"io"
     9  
    10  	"github.com/grailbio/base/recordio"
    11  	"github.com/grailbio/base/recordio/deprecated"
    12  	"github.com/klauspost/compress/flate"
    13  )
    14  
    15  // WriterOpts represents the options accepted by NewWriter.
    16  type WriterOpts struct {
    17  	// Marshal is called to marshal an object to a byte slice.
    18  	Marshal recordio.MarshalFunc
    19  
    20  	// Index is called whenever a new record is written.
    21  	Index deprecated.RecordIndex
    22  
    23  	// MaxItems sets recordio.LegacyPackedWriterOpts when calling NewLegacyPackedWriter
    24  	MaxItems uint32
    25  
    26  	// MaxBytes sets recordio.LegacyPackedWriterOpts when calling NewLegacyPackedWriter.
    27  	MaxBytes uint32
    28  
    29  	// Flushed is passed to recordio.PackedWriterOpts.
    30  	Flushed func() error
    31  
    32  	// FlateLevel indicates the compression level to use, it accepts
    33  	// the values supported by the flate package. The default value (0)
    34  	// is flate.NoCompression.
    35  	// Compression is always performed before encryption.
    36  	FlateLevel int
    37  }
    38  
    39  type writer struct {
    40  	deprecated.LegacyPackedWriter
    41  	compressor *FlateTransform
    42  	opts       WriterOpts
    43  }
    44  
    45  // NewWriter returns a recordio.LegacyPackedWriter that can optionally compress
    46  // and encrypt the items it writes.
    47  func NewWriter(w io.Writer, opts WriterOpts) (deprecated.LegacyPackedWriter, error) {
    48  	wr := &writer{opts: opts}
    49  	subopts := deprecated.LegacyPackedWriterOpts{
    50  		Marshal:  deprecated.MarshalFunc(opts.Marshal),
    51  		Index:    opts.Index,
    52  		MaxItems: opts.MaxItems,
    53  		MaxBytes: opts.MaxBytes,
    54  		Flushed:  opts.Flushed,
    55  	}
    56  
    57  	compress := false
    58  	switch opts.FlateLevel {
    59  	case flate.BestSpeed,
    60  		flate.BestCompression,
    61  		flate.DefaultCompression,
    62  		flate.HuffmanOnly:
    63  		compress = true
    64  	}
    65  
    66  	if compress {
    67  		wr.compressor = NewFlateTransform(opts.FlateLevel)
    68  		subopts.Transform = wr.compressor.CompressTransform
    69  	}
    70  	wr.LegacyPackedWriter = deprecated.NewLegacyPackedWriter(w, subopts)
    71  	return wr, nil
    72  }