[O11Y-224] add support for merging parquet files.

buffer.go

@@ -1,126 +1,259 @@
 package parquet
 
 import (
-	"bytes"
-	"encoding/binary"
-	"io"
-	"os"
-	"path/filepath"
-	"sync"
+	"sort"
 )
 
-const (
-	defaultBufferSize      = 4096
-	defaultLevelBufferSize = 1024
-)
-
-var (
-	defaultBufferPool bufferPool
-)
-
-type Buffer interface {
-	io.Reader
-	io.Writer
+// Buffer represents an in-memory group of parquet rows.
+//
+// The main purpose of the Buffer type is to provide a way to sort rows before
+// writing them to a parquet file. Buffer implements sort.Interface as a way
+// to support reordering the rows that have been written to it.
+type Buffer struct {
+	config  *RowGroupConfig
+	schema  *Schema
+	rowbuf  []Value
+	colbuf  [][]Value
+	columns []ColumnBuffer
+	sorted  []ColumnBuffer
 }
 
-type BufferPool interface {
-	GetBuffer() Buffer
-	PutBuffer(Buffer)
+// NewBuffer constructs a new buffer, using the given list of buffer options
+// to configure the buffer returned by the function.
+//
+// The function panics if the buffer configuration is invalid. Programs that
+// cannot guarantee the validity of the options passed to NewBuffer should
+// construct the buffer configuration independently prior to calling this
+// function:
+//
+//	config, err := parquet.NewRowGroupConfig(options...)
+//	if err != nil {
+//		// handle the configuration error
+//		...
+//	} else {
+//		// this call to create a buffer is guaranteed not to panic
+//		buffer := parquet.NewBuffer(config)
+//		...
+//	}
+//
+func NewBuffer(options ...RowGroupOption) *Buffer {
+	config, err := NewRowGroupConfig(options...)
+	if err != nil {
+		panic(err)
+	}
+	buf := &Buffer{
+		config: config,
+	}
+	if config.Schema != nil {
+		buf.configure(config.Schema)
+	}
+	return buf
 }
 
-func NewBufferPool() BufferPool { return new(bufferPool) }
+func (buf *Buffer) configure(schema *Schema) {
+	if schema == nil {
+		return
+	}
+	sortingColumns := buf.config.SortingColumns
+	buf.sorted = make([]ColumnBuffer, len(sortingColumns))
 
-type bufferPool struct{ sync.Pool }
+	forEachLeafColumnOf(schema, func(leaf leafColumn) {
+		nullOrdering := nullsGoLast
+		columnType := leaf.node.Type()
+		bufferSize := buf.config.ColumnBufferSize
+		dictionary := (Dictionary)(nil)
+		encoding, _ := encodingAndCompressionOf(leaf.node)
 
-func (pool *bufferPool) GetBuffer() Buffer {
-	b, _ := pool.Get().(*buffer)
-	if b == nil {
-		b = new(buffer)
-	} else {
-		b.Reset()
+		if isDictionaryEncoding(encoding) {
+			bufferSize /= 2
+			dictionary = columnType.NewDictionary(bufferSize)
+			columnType = dictionary.Type()
+		}
+
+		column := columnType.NewColumnBuffer(leaf.columnIndex, bufferSize)
+		switch {
+		case leaf.maxRepetitionLevel > 0:
+			column = newRepeatedColumnBuffer(column, leaf.maxRepetitionLevel, leaf.maxDefinitionLevel, nullOrdering)
+		case leaf.maxDefinitionLevel > 0:
+			column = newOptionalColumnBuffer(column, leaf.maxDefinitionLevel, nullOrdering)
+		}
+		buf.columns = append(buf.columns, column)
+
+		if sortingIndex := searchSortingColumn(sortingColumns, leaf.path); sortingIndex < len(sortingColumns) {
+			if sortingColumns[sortingIndex].Descending() {
+				column = &reversedColumnBuffer{column}
+			}
+			if sortingColumns[sortingIndex].NullsFirst() {
+				nullOrdering = nullsGoFirst
+			}
+			buf.sorted[sortingIndex] = column
+		}
+	})
+
+	buf.schema = schema
+	buf.rowbuf = make([]Value, 0, 10)
+	buf.colbuf = make([][]Value, len(buf.columns))
+}
+
+// Size returns the estimated size of the buffer in memory (in bytes).
+func (buf *Buffer) Size() int64 {
+	size := int64(0)
+	for _, col := range buf.columns {
+		size += col.Size()
 	}
-	return b
+	return size
 }
 
-func (pool *bufferPool) PutBuffer(buf Buffer) {
-	if b, _ := buf.(*buffer); b != nil {
-		pool.Put(b)
+// NumRows returns the number of rows written to the buffer.
+func (buf *Buffer) NumRows() int {
+	if len(buf.columns) == 0 {
+		return 0
+	} else {
+		// All columns have the same number of rows.
+		return buf.columns[0].Len()
 	}
 }
 
-type buffer struct{ bytes.Buffer }
+// NumColumns returns the number of columns in the buffer.
+//
+// The count will be zero until a schema is configured on buf.
+func (buf *Buffer) NumColumns() int { return len(buf.columns) }
 
-func (b *buffer) Close() error {
-	b.Reset()
-	return nil
-}
+// Column returns the buffer column at index i.
+//
+// The method panics if i is negative or beyond the last column index in buf.
+func (buf *Buffer) Column(i int) ColumnChunk { return buf.columns[i] }
 
-type fileBufferPool struct {
-	err     error
-	tempdir string
-	pattern string
-}
+// Schema returns the schema of the buffer.
+//
+// The schema is either configured by passing a Schema in the option list when
+// constructing the buffer, or lazily discovered when the first row is written.
+func (buf *Buffer) Schema() *Schema { return buf.schema }
+
+// SortingColumns returns the list of columns by which the buffer will be
+// sorted.
+//
+// The sorting order is configured by passing a SortingColumns option when
+// constructing the buffer.
+func (buf *Buffer) SortingColumns() []SortingColumn { return buf.config.SortingColumns }
+
+// Len returns the number of rows written to the buffer.
+func (buf *Buffer) Len() int { return buf.NumRows() }
 
-func NewFileBufferPool(tempdir, pattern string) BufferPool {
-	pool := &fileBufferPool{
-		tempdir: tempdir,
-		pattern: pattern,
+// Less returns true if row[i] < row[j] in the buffer.
+func (buf *Buffer) Less(i, j int) bool {
+	for _, col := range buf.sorted {
+		switch {
+		case col.Less(i, j):
+			return true
+		case col.Less(j, i):
+			return false
+		}
 	}
-	pool.tempdir, pool.err = filepath.Abs(pool.tempdir)
-	return pool
+	return false
 }
 
-func (pool *fileBufferPool) GetBuffer() Buffer {
-	if pool.err != nil {
-		return &errorBuffer{err: pool.err}
+// Swap exchanges the rows at indexes i and j.
+func (buf *Buffer) Swap(i, j int) {
+	for _, col := range buf.columns {
+		col.Swap(i, j)
 	}
-	f, err := os.CreateTemp(pool.tempdir, pool.pattern)
-	if err != nil {
-		return &errorBuffer{err: err}
+}
+
+// Reset clears the content of the buffer, allowing it to be reused.
+func (buf *Buffer) Reset() {
+	for _, col := range buf.columns {
+		col.Reset()
 	}
-	return f
 }
 
-func (pool *fileBufferPool) PutBuffer(buf Buffer) {
-	if f, _ := buf.(*os.File); f != nil {
-		defer f.Close()
-		os.Remove(f.Name())
+// Write writes a row held in a Go value to the buffer.
+func (buf *Buffer) Write(row interface{}) error {
+	if buf.schema == nil {
+		buf.configure(SchemaOf(row))
 	}
+	defer func() {
+		clearValues(buf.rowbuf)
+	}()
+	buf.rowbuf = buf.schema.Deconstruct(buf.rowbuf[:0], row)
+	return buf.WriteRow(buf.rowbuf)
 }
 
-type errorBuffer struct{ err error }
+// WriteRow writes a parquet row to the buffer.
+func (buf *Buffer) WriteRow(row Row) error {
+	defer func() {
+		for i, colbuf := range buf.colbuf {
+			clearValues(colbuf)
+			buf.colbuf[i] = colbuf[:0]
+		}
+	}()
 
-func (errbuf *errorBuffer) Read([]byte) (int, error)          { return 0, errbuf.err }
-func (errbuf *errorBuffer) Write([]byte) (int, error)         { return 0, errbuf.err }
-func (errbuf *errorBuffer) ReadFrom(io.Reader) (int64, error) { return 0, errbuf.err }
-func (errbuf *errorBuffer) WriteTo(io.Writer) (int64, error)  { return 0, errbuf.err }
+	if buf.schema == nil {
+		return ErrRowGroupSchemaMissing
+	}
 
-var (
-	_ io.ReaderFrom = (*errorBuffer)(nil)
-	_ io.WriterTo   = (*errorBuffer)(nil)
-)
+	for _, value := range row {
+		columnIndex := value.Column()
+		buf.colbuf[columnIndex] = append(buf.colbuf[columnIndex], value)
+	}
 
-type lengthPrefixedWriter struct {
-	writer io.Writer
-	buffer []byte
-}
+	for columnIndex, values := range buf.colbuf {
+		if err := buf.columns[columnIndex].WriteRow(values); err != nil {
+			return err
+		}
+	}
 
-func (w *lengthPrefixedWriter) Reset(ww io.Writer) {
-	w.writer = ww
-	w.buffer = append(w.buffer[:0], 0, 0, 0, 0)
+	return nil
 }
 
-func (w *lengthPrefixedWriter) Close() error {
-	if len(w.buffer) > 0 {
-		defer func() { w.buffer = w.buffer[:0] }()
-		binary.LittleEndian.PutUint32(w.buffer, uint32(len(w.buffer))-4)
-		_, err := w.writer.Write(w.buffer)
-		return err
+// WriteRowGroup satisfies the RowGroupWriter interface.
+func (buf *Buffer) WriteRowGroup(rowGroup RowGroup) (int64, error) {
+	rowGroupSchema := rowGroup.Schema()
+	switch {
+	case rowGroupSchema == nil:
+		return 0, ErrRowGroupSchemaMissing
+	case buf.schema == nil:
+		buf.configure(rowGroupSchema)
+	case !nodesAreEqual(buf.schema, rowGroupSchema):
+		return 0, ErrRowGroupSchemaMismatch
 	}
-	return nil
+	if !sortingColumnsHavePrefix(rowGroup.SortingColumns(), buf.SortingColumns()) {
+		return 0, ErrRowGroupSortingColumnsMismatch
+	}
+	n := buf.NumRows()
+	_, err := CopyRows(bufferWriter{buf}, rowGroup.Rows())
+	return int64(buf.NumRows() - n), err
+}
+
+// Rows returns a reader exposing the current content of the buffer.
+//
+// The buffer and the returned reader share memory. Mutating the buffer
+// concurrently to reading rows may result in non-deterministic behavior.
+func (buf *Buffer) Rows() RowReader { return &rowGroupRowReader{rowGroup: buf} }
+
+// bufferWriter is an adapter for Buffer which implements both RowWriter and
+// PageWriter to enable optimizations in CopyRows for types that support writing
+// rows by copying whole pages instead of calling WriteRow repeatedly.
+type bufferWriter struct{ buf *Buffer }
+
+func (w bufferWriter) WriteRow(row Row) error {
+	return w.buf.WriteRow(row)
 }
 
-func (w *lengthPrefixedWriter) Write(b []byte) (int, error) {
-	w.buffer = append(w.buffer, b...)
-	return len(b), nil
+func (w bufferWriter) WriteValues(values []Value) (int, error) {
+	return w.buf.columns[values[0].Column()].WriteValues(values)
 }
+
+func (w bufferWriter) WritePage(page Page) (int64, error) {
+	return CopyValues(w.buf.columns[page.Column()], page.Values())
+}
+
+var (
+	_ RowGroup       = (*Buffer)(nil)
+	_ RowGroupWriter = (*Buffer)(nil)
+	_ sort.Interface = (*Buffer)(nil)
+
+	_ RowWriter   = (*bufferWriter)(nil)
+	_ PageWriter  = (*bufferWriter)(nil)
+	_ ValueWriter = (*bufferWriter)(nil)
+)