Paddle/go/recordio/writer.go

package recordio

import (
	"fmt"
	"io"
)

const (
	defaultMaxChunkSize = 32 * 1024 * 1024
)

// Writer creates a RecordIO file.
type Writer struct {
	io.Writer    // Set to nil to mark a closed writer.
	chunk        *Chunk
	maxChunkSize int // total records size, excluding metadata, before compression.
	compressor   int
}

// NewWriter creates a RecordIO file writer.  Each chunk is compressed
// using the deflate algorithm given compression level.  Note that
// level 0 means no compression and -1 means default compression.
func NewWriter(w io.Writer, maxChunkSize, compressor int) *Writer {
	if maxChunkSize < 0 {
		maxChunkSize = defaultMaxChunkSize
	}

	if compressor < 0 {
		compressor = defaultCompressor
	}

	return &Writer{
		Writer:       w,
		chunk:        &Chunk{},
		maxChunkSize: maxChunkSize,
		compressor:   compressor}
}

// Writes a record.  It returns an error if Close has been called.
func (w *Writer) Write(record []byte) (int, error) {
	if w.Writer == nil {
		return 0, fmt.Errorf("Cannot write since writer had been closed")
	}

	if w.chunk.numBytes+len(record) > w.maxChunkSize {
		if e := w.chunk.dump(w.Writer, w.compressor); e != nil {
			return 0, e
		}
	}

	w.chunk.add(record)
	return len(record), nil
}

// Close flushes the current chunk and makes the writer invalid.
func (w *Writer) Close() error {
	e := w.chunk.dump(w.Writer, w.compressor)
	w.Writer = nil
	return e
}