Add support for zlib compression (#14)

* Add support for zlib compression

* Address comments in PR. Use constant for Zlib default chunk size and use proper fmt printing verbs

* Fix maximum chunkLength check to ensure that the value to be used in the header remains withing the available field bounds

Author: mgilbir

buffererwriter.go

@@ -1,66 +1,173 @@
 package orc
 
 import (
-	"bufio"
 	"bytes"
+	"fmt"
+	"io"
+	"sync"
 )
 
-// BufferedWriter wraps a *bufio.Writer and records the current
-// position of the writer prior to flushing to the underlying
-// writer.
 type BufferedWriter struct {
-	*bufio.Writer
-	*bytes.Buffer
-	codec      CompressionCodec
-	checkpoint uint64
-	written    uint64
+	uncompressedBuffer *bytes.Buffer
+	encodedBuffer      *bytes.Buffer
+	codec              CompressionCodec
+	chunkSize          int
+	checkpoint         uint64
+	written            uint64
+	sync.Mutex
 }
 
 // NewBufferedWriter returns a new BufferedWriter using the provided
 // CompressionCodec.
 func NewBufferedWriter(codec CompressionCodec) *BufferedWriter {
-	buf := &bytes.Buffer{}
+	chunkSize := 1024
+	switch codec.(type) {
+	case CompressionNone:
+		chunkSize = 1
+	case CompressionZlib:
+		chunkSize = int(DefaultCompressionChunkSize)
+	}
 	return &BufferedWriter{
-		codec: codec,
-		Writer: bufio.NewWriterSize(
-			codec.Encoder(buf),
-			int(DefaultCompressionChunkSize),
-		),
-		Buffer: buf,
+		codec:              codec,
+		uncompressedBuffer: &bytes.Buffer{},
+		encodedBuffer:      &bytes.Buffer{},
+		chunkSize:          chunkSize,
 	}
 }
 
-// WriteByte writes a byte to the underlying buffer an increments the total
-// number of bytes written.
+// WriteByte writes a byte to the underlying buffer.
+// If the desired chunk size is reached, the buffer is compressed
 func (b *BufferedWriter) WriteByte(c byte) error {
-	b.written++
-	return b.Writer.WriteByte(c)
+	b.Lock()
+	defer b.Unlock()
+
+	if b.uncompressedBuffer.Len() == b.chunkSize {
+		err := b.spill()
+		if err != nil {
+			return err
+		}
+	}
+	_, err := b.uncompressedBuffer.Write([]byte{c})
+	return err
 }
 
-// Write writes the provided byte slice to the underlying buffer an increments
-// the total number of bytes written.
+// Write writes the provided byte slice to the underlying buffer.
+// If the desired chunk size is reached, the buffer is compressed
 func (b *BufferedWriter) Write(p []byte) (int, error) {
-	b.written += uint64(len(p))
-	return b.Writer.Write(p)
+	b.Lock()
+	defer b.Unlock()
+	pos := 0
+
+	var remaining int
+	l := len(p)
+	c := b.chunkSize - b.uncompressedBuffer.Len()
+	if c > l {
+		remaining = l
+	} else {
+		remaining = c
+	}
+
+	n, err := b.uncompressedBuffer.Write(p[pos : pos+remaining])
+	if err != nil {
+		return 0, err
+	}
+	pos += n
+	l -= n
+
+	for l != 0 {
+		if err := b.spill(); err != nil {
+			return 0, err
+		}
+
+		c = b.chunkSize - b.uncompressedBuffer.Len()
+		if c > l {
+			remaining = l
+		} else {
+			remaining = c
+		}
+
+		n, err = b.uncompressedBuffer.Write(p[pos : pos+remaining])
+		if err != nil {
+			return 0, err
+		}
+		pos += n
+		l -= n
+	}
+
+	return pos, nil
+}
+
+// spill to the encoder to handle the compression and update the number of
+// written bytes to the encoded buffer
+func (b *BufferedWriter) spill() error {
+	encoder := b.codec.Encoder(b.encodedBuffer)
+	l := b.uncompressedBuffer.Len()
+	n, err := io.Copy(encoder, b.uncompressedBuffer)
+	if err != nil {
+		return err
+	}
+	if int(n) != l {
+		return fmt.Errorf("Expected to write %d bytes, wrote %d", l, n)
+	}
+
+	err = encoder.Close()
+	if err != nil {
+		return err
+	}
+	b.written += uint64(n)
+	return nil
 }
 
 func (b *BufferedWriter) Positions() []uint64 {
+	b.Lock()
+	defer b.Unlock()
+
+	//TODO: Do we still need the checkpoint?
 	switch b.codec.(type) {
 	case CompressionNone:
 		checkpoint := b.checkpoint
 		b.checkpoint = b.written
 		return []uint64{checkpoint}
 	default:
-		return nil
+		//TODO: check if this is correct
+		checkpoint := b.checkpoint
+		b.checkpoint = b.written
+		return []uint64{checkpoint}
+
+		// return nil
 	}
 }
 
+func (b *BufferedWriter) Flush() error {
+	b.Lock()
+	defer b.Unlock()
+
+	return b.spill()
+}
+
+func (b *BufferedWriter) Read(p []byte) (int, error) {
+	b.Lock()
+	defer b.Unlock()
+
+	return b.encodedBuffer.Read(p)
+}
+
+func (b *BufferedWriter) Len() int {
+	b.Lock()
+	defer b.Unlock()
+
+	return b.encodedBuffer.Len()
+}
+
 // Close flushes any buffered bytes to the underlying writer.
 func (b *BufferedWriter) Close() error {
-	return b.Writer.Flush()
+	return b.spill()
 }
 
-// Reset resets the underlying bytes.Buffer.
+// Reset resets the underlying encoded buffer
 func (b *BufferedWriter) Reset() {
-	b.Buffer.Reset()
+	b.Lock()
+	defer b.Unlock()
+
+	b.encodedBuffer.Reset()
 }