Repository: segmentio/parquet-go
Pull request #437 (2022-12-07 10:10:07)

support configuring the sorting buffer pool
This PR adds the ability to configure a buffer pool that `parquet.SortingWriter` instances use to allocate the buffer where their temporary files.

I have reused the `parquet.PageBufferPool` API for this purpose, but by default I kept the sorting buffers and column page buffers separated because they tend to be of very different sizes.

I'm open to naming suggestions, it seems `parquet.PageBufferPool` is not so representative now that we have a new use case for it.

achille-roussel: > I'm open to naming suggestions, it seems parquet.PageBufferPool is not so representative now that we have a new use case for it.

For lack of creativity I renamed `parquet.PageBufferPool` to `parquet.BufferPool`. I'm a bit concerned that it could get confusing since there is a `parquet.Buffer` type which has nothing to do with it, I'm considering moving `parquet.BufferPool` to a subpackage to avoid the confusion but haven't made up my mind on a name for the package either so for now this is the way.

Patch:
diff --git a/page_buffer.go b/buffer_pool.go
similarity index 57%
rename from page_buffer.go
rename to buffer_pool.go
index ba67245..85b5407 100644
--- a/page_buffer.go
+++ b/buffer_pool.go
@@ -8,7 +8,7 @@ import (
 	"sync"
 )
 
-// PageBufferPool is an interface abstracting the underlying implementation of
+// BufferPool is an interface abstracting the underlying implementation of
 // page buffer pools.
 //
 // The parquet-go package provides two implementations of this interface, one
@@ -19,38 +19,38 @@ import (
 // of page buffers may choose to provide their own implementation and install it
 // via the parquet.ColumnPageBuffers writer option.
 //
-// PageBufferPool implementations must be safe to use concurrently from multiple
+// BufferPool implementations must be safe to use concurrently from multiple
 // goroutines.
-type PageBufferPool interface {
-	// GetPageBuffer is called when a parquet writer needs to acquire a new
+type BufferPool interface {
+	// GetBuffer is called when a parquet writer needs to acquire a new
 	// page buffer from the pool.
-	GetPageBuffer() io.ReadWriteSeeker
+	GetBuffer() io.ReadWriteSeeker
 
-	// PutPageBuffer is called when a parquet writer releases a page buffer to
+	// PutBuffer is called when a parquet writer releases a page buffer to
 	// the pool.
 	//
 	// The parquet.Writer type guarantees that the buffers it calls this method
-	// with were previously acquired by a call to GetPageBuffer on the same
+	// with were previously acquired by a call to GetBuffer on the same
 	// pool, and that it will not use them anymore after the call.
-	PutPageBuffer(io.ReadWriteSeeker)
+	PutBuffer(io.ReadWriteSeeker)
 }
 
-// NewPageBufferPool creates a new in-memory page buffer pool.
+// NewBufferPool creates a new in-memory page buffer pool.
 //
 // The implementation is backed by sync.Pool and allocates memory buffers on the
 // Go heap.
-func NewPageBufferPool() PageBufferPool { return new(pageBufferPool) }
+func NewBufferPool() BufferPool { return new(memoryBufferPool) }
 
-type pageBuffer struct {
+type memoryBuffer struct {
 	data []byte
 	off  int
 }
 
-func (p *pageBuffer) Reset() {
+func (p *memoryBuffer) Reset() {
 	p.data, p.off = p.data[:0], 0
 }
 
-func (p *pageBuffer) Read(b []byte) (n int, err error) {
+func (p *memoryBuffer) Read(b []byte) (n int, err error) {
 	n = copy(b, p.data[p.off:])
 	p.off += n
 	if p.off == len(p.data) {
@@ -59,7 +59,7 @@ func (p *pageBuffer) Read(b []byte) (n int, err error) {
 	return n, err
 }
 
-func (p *pageBuffer) Write(b []byte) (int, error) {
+func (p *memoryBuffer) Write(b []byte) (int, error) {
 	n := copy(p.data[p.off:cap(p.data)], b)
 	p.data = p.data[:p.off+n]
 
@@ -71,13 +71,13 @@ func (p *pageBuffer) Write(b []byte) (int, error) {
 	return len(b), nil
 }
 
-func (p *pageBuffer) WriteTo(w io.Writer) (int64, error) {
+func (p *memoryBuffer) WriteTo(w io.Writer) (int64, error) {
 	n, err := w.Write(p.data[p.off:])
 	p.off += n
 	return int64(n), err
 }
 
-func (p *pageBuffer) Seek(offset int64, whence int) (int64, error) {
+func (p *memoryBuffer) Seek(offset int64, whence int) (int64, error) {
 	switch whence {
 	case io.SeekCurrent:
 		offset += int64(p.off)
@@ -94,20 +94,20 @@ func (p *pageBuffer) Seek(offset int64, whence int) (int64, error) {
 	return offset, nil
 }
 
-type pageBufferPool struct{ sync.Pool }
+type memoryBufferPool struct{ sync.Pool }
 
-func (pool *pageBufferPool) GetPageBuffer() io.ReadWriteSeeker {
-	b, _ := pool.Get().(*pageBuffer)
+func (pool *memoryBufferPool) GetBuffer() io.ReadWriteSeeker {
+	b, _ := pool.Get().(*memoryBuffer)
 	if b == nil {
-		b = new(pageBuffer)
+		b = new(memoryBuffer)
 	} else {
 		b.Reset()
 	}
 	return b
 }
 
-func (pool *pageBufferPool) PutPageBuffer(buf io.ReadWriteSeeker) {
-	if b, _ := buf.(*pageBuffer); b != nil {
+func (pool *memoryBufferPool) PutBuffer(buf io.ReadWriteSeeker) {
+	if b, _ := buf.(*memoryBuffer); b != nil {
 		pool.Put(b)
 	}
 }
@@ -119,7 +119,7 @@ type fileBufferPool struct {
 }
 
 // NewFileBufferPool creates a new on-disk page buffer pool.
-func NewFileBufferPool(tempdir, pattern string) PageBufferPool {
+func NewFileBufferPool(tempdir, pattern string) BufferPool {
 	pool := &fileBufferPool{
 		tempdir: tempdir,
 		pattern: pattern,
@@ -128,7 +128,7 @@ func NewFileBufferPool(tempdir, pattern string) PageBufferPool {
 	return pool
 }
 
-func (pool *fileBufferPool) GetPageBuffer() io.ReadWriteSeeker {
+func (pool *fileBufferPool) GetBuffer() io.ReadWriteSeeker {
 	if pool.err != nil {
 		return &errorBuffer{err: pool.err}
 	}
@@ -139,7 +139,7 @@ func (pool *fileBufferPool) GetPageBuffer() io.ReadWriteSeeker {
 	return f
 }
 
-func (pool *fileBufferPool) PutPageBuffer(buf io.ReadWriteSeeker) {
+func (pool *fileBufferPool) PutBuffer(buf io.ReadWriteSeeker) {
 	if f, _ := buf.(*os.File); f != nil {
 		defer f.Close()
 		os.Remove(f.Name())
@@ -155,9 +155,35 @@ func (buf *errorBuffer) WriteTo(io.Writer) (int64, error)  { return 0, buf.err }
 func (buf *errorBuffer) Seek(int64, int) (int64, error)    { return 0, buf.err }
 
 var (
-	defaultPageBufferPool pageBufferPool
+	defaultColumnBufferPool  memoryBufferPool
+	defaultSortingBufferPool memoryBufferPool
 
 	_ io.ReaderFrom = (*errorBuffer)(nil)
 	_ io.WriterTo   = (*errorBuffer)(nil)
-	_ io.WriterTo   = (*pageBuffer)(nil)
+	_ io.WriterTo   = (*memoryBuffer)(nil)
 )
+
+type readerAt struct {
+	reader io.ReadSeeker
+	offset int64
+}
+
+func (r *readerAt) ReadAt(b []byte, off int64) (int, error) {
+	if r.offset < 0 || off != r.offset {
+		off, err := r.reader.Seek(off, io.SeekStart)
+		if err != nil {
+			return 0, err
+		}
+		r.offset = off
+	}
+	n, err := r.reader.Read(b)
+	r.offset += int64(n)
+	return n, err
+}
+
+func newReaderAt(r io.ReadSeeker) io.ReaderAt {
+	if rr, ok := r.(io.ReaderAt); ok {
+		return rr
+	}
+	return &readerAt{reader: r, offset: -1}
+}
diff --git a/config.go b/config.go
index d84aac0..f26397d 100644
--- a/config.go
+++ b/config.go
@@ -197,7 +197,7 @@ func (c *ReaderConfig) Validate() error {
 //	})
 type WriterConfig struct {
 	CreatedBy            string
-	ColumnPageBuffers    PageBufferPool
+	ColumnPageBuffers    BufferPool
 	ColumnIndexSizeLimit int
 	PageBufferSize       int
 	WriteBufferSize      int
@@ -216,13 +216,16 @@ type WriterConfig struct {
 func DefaultWriterConfig() *WriterConfig {
 	return &WriterConfig{
 		CreatedBy:            defaultCreatedBy(),
-		ColumnPageBuffers:    &defaultPageBufferPool,
+		ColumnPageBuffers:    &defaultColumnBufferPool,
 		ColumnIndexSizeLimit: DefaultColumnIndexSizeLimit,
 		PageBufferSize:       DefaultPageBufferSize,
 		WriteBufferSize:      DefaultWriteBufferSize,
 		DataPageVersion:      DefaultDataPageVersion,
 		DataPageStatistics:   DefaultDataPageStatistics,
 		MaxRowsPerRowGroup:   DefaultMaxRowsPerRowGroup,
+		Sorting: SortingConfig{
+			SortingBuffers: &defaultSortingBufferPool,
+		},
 	}
 }
 
@@ -258,7 +261,7 @@ func (c *WriterConfig) ConfigureWriter(config *WriterConfig) {
 
 	*config = WriterConfig{
 		CreatedBy:            coalesceString(c.CreatedBy, config.CreatedBy),
-		ColumnPageBuffers:    coalescePageBufferPool(c.ColumnPageBuffers, config.ColumnPageBuffers),
+		ColumnPageBuffers:    coalesceBufferPool(c.ColumnPageBuffers, config.ColumnPageBuffers),
 		ColumnIndexSizeLimit: coalesceInt(c.ColumnIndexSizeLimit, config.ColumnIndexSizeLimit),
 		PageBufferSize:       coalesceInt(c.PageBufferSize, config.PageBufferSize),
 		WriteBufferSize:      coalesceInt(c.WriteBufferSize, config.WriteBufferSize),
@@ -281,6 +284,7 @@ func (c *WriterConfig) Validate() error {
 		validatePositiveInt(baseName+"ColumnIndexSizeLimit", c.ColumnIndexSizeLimit),
 		validatePositiveInt(baseName+"PageBufferSize", c.PageBufferSize),
 		validateOneOfInt(baseName+"DataPageVersion", c.DataPageVersion, 1, 2),
+		c.Sorting.Validate(),
 	)
 }
 
@@ -303,6 +307,9 @@ type RowGroupConfig struct {
 func DefaultRowGroupConfig() *RowGroupConfig {
 	return &RowGroupConfig{
 		ColumnBufferCapacity: DefaultColumnBufferCapacity,
+		Sorting: SortingConfig{
+			SortingBuffers: &defaultSortingBufferPool,
+		},
 	}
 }
 
@@ -322,6 +329,7 @@ func (c *RowGroupConfig) Validate() error {
 	const baseName = "parquet.(*RowGroupConfig)."
 	return errorInvalidConfiguration(
 		validatePositiveInt(baseName+"ColumnBufferCapacity", c.ColumnBufferCapacity),
+		c.Sorting.Validate(),
 	)
 }
 
@@ -351,6 +359,7 @@ func (c *RowGroupConfig) ConfigureRowGroup(config *RowGroupConfig) {
 //		),
 //	})
 type SortingConfig struct {
+	SortingBuffers     BufferPool
 	SortingColumns     []SortingColumn
 	DropDuplicatedRows bool
 }
@@ -359,7 +368,7 @@ type SortingConfig struct {
 // default row group configuration.
 func DefaultSortingConfig() *SortingConfig {
 	return &SortingConfig{
-		DropDuplicatedRows: false,
+		SortingBuffers: &defaultSortingBufferPool,
 	}
 }
 
@@ -375,7 +384,10 @@ func NewSortingConfig(options ...SortingOption) (*SortingConfig, error) {
 }
 
 func (c *SortingConfig) Validate() error {
-	return nil
+	const baseName = "parquet.(*SortingConfig)."
+	return errorInvalidConfiguration(
+		validateNotNil(baseName+"SortingBuffers", c.SortingBuffers),
+	)
 }
 
 func (c *SortingConfig) Apply(options ...SortingOption) {
@@ -517,7 +529,7 @@ func CreatedBy(application, version, build string) WriterOption {
 // on the amount of memory available.
 //
 // Defaults to using in-memory buffers.
-func ColumnPageBuffers(buffers PageBufferPool) WriterOption {
+func ColumnPageBuffers(buffers BufferPool) WriterOption {
 	return writerOption(func(config *WriterConfig) { config.ColumnPageBuffers = buffers })
 }
 
@@ -625,6 +637,14 @@ func SortingColumns(columns ...SortingColumn) SortingOption {
 	return sortingOption(func(config *SortingConfig) { config.SortingColumns = columns })
 }
 
+// SortingBuffers creates a configuration option which sets the pool of buffers
+// used to hold intermediary state when sorting parquet rows.
+//
+// Defaults to using in-memory buffers.
+func SortingBuffers(buffers BufferPool) SortingOption {
+	return sortingOption(func(config *SortingConfig) { config.SortingBuffers = buffers })
+}
+
 // DropDuplicatedRows configures whether a sorting writer will keep or remove
 // duplicated rows.
 //
@@ -684,7 +704,7 @@ func coalesceBytes(b1, b2 []byte) []byte {
 	return b2
 }
 
-func coalescePageBufferPool(p1, p2 PageBufferPool) PageBufferPool {
+func coalesceBufferPool(p1, p2 BufferPool) BufferPool {
 	if p1 != nil {
 		return p1
 	}
@@ -707,6 +727,7 @@ func coalesceSortingColumns(s1, s2 []SortingColumn) []SortingColumn {
 
 func coalesceSortingConfig(c1, c2 SortingConfig) SortingConfig {
 	return SortingConfig{
+		SortingBuffers:     coalesceBufferPool(c1.SortingBuffers, c2.SortingBuffers),
 		SortingColumns:     coalesceSortingColumns(c1.SortingColumns, c2.SortingColumns),
 		DropDuplicatedRows: c1.DropDuplicatedRows,
 	}
diff --git a/sorting.go b/sorting.go
index 2311a0b..5b77d80 100644
--- a/sorting.go
+++ b/sorting.go
@@ -3,7 +3,6 @@
 package parquet
 
 import (
-	"bytes"
 	"io"
 	"sort"
 )
@@ -24,9 +23,9 @@ import (
 // causes a lot of cache misses since the data set cannot be held in CPU caches.
 type SortingWriter[T any] struct {
 	rows    *RowBuffer[T]
-	buffer  *bytes.Buffer
 	writer  *GenericWriter[T]
 	output  *GenericWriter[T]
+	buffer  io.ReadWriteSeeker
 	maxRows int64
 	numRows int64
 	sorting SortingConfig
@@ -48,21 +47,12 @@ func NewSortingWriter[T any](output io.Writer, sortRowCount int64, options ...Wr
 	if err != nil {
 		panic(err)
 	}
-
-	// At this time the intermediary buffer where partial row groups are
-	// written is held in memory. This may prove impractical if the parquet
-	// file exceeds the amount of memory available to the application, we will
-	// revisit when we have a need to put this buffer on a different storage
-	// medium.
-	buffer := bytes.NewBuffer(nil)
-
 	return &SortingWriter[T]{
 		rows: NewRowBuffer[T](&RowGroupConfig{
 			Schema:  config.Schema,
 			Sorting: config.Sorting,
 		}),
-		buffer: buffer,
-		writer: NewGenericWriter[T](buffer, &WriterConfig{
+		writer: NewGenericWriter[T](io.Discard, &WriterConfig{
 			CreatedBy:            config.CreatedBy,
 			ColumnPageBuffers:    config.ColumnPageBuffers,
 			ColumnIndexSizeLimit: config.ColumnIndexSizeLimit,
@@ -87,6 +77,8 @@ func (w *SortingWriter[T]) Close() error {
 }
 
 func (w *SortingWriter[T]) Flush() error {
+	defer w.resetSortingBuffer()
+
 	if err := w.sortAndWriteBufferedRows(); err != nil {
 		return err
 	}
@@ -95,17 +87,16 @@ func (w *SortingWriter[T]) Flush() error {
 		return nil
 	}
 
-	defer func() {
-		w.buffer.Reset()
-		w.writer.Reset(w.buffer)
-		w.numRows = 0
-	}()
-
 	if err := w.writer.Close(); err != nil {
 		return err
 	}
 
-	f, err := OpenFile(bytes.NewReader(w.buffer.Bytes()), int64(w.buffer.Len()),
+	size, err := w.buffer.Seek(0, io.SeekCurrent)
+	if err != nil {
+		return err
+	}
+
+	f, err := OpenFile(newReaderAt(w.buffer), size,
 		&FileConfig{
 			SkipPageIndex:    true,
 			SkipBloomFilters: true,
@@ -142,11 +133,19 @@ func (w *SortingWriter[T]) Flush() error {
 }
 
 func (w *SortingWriter[T]) Reset(output io.Writer) {
-	w.rows.Reset()
-	w.buffer.Reset()
-	w.writer.Reset(w.buffer)
 	w.output.Reset(output)
+	w.rows.Reset()
+	w.resetSortingBuffer()
+}
+
+func (w *SortingWriter[T]) resetSortingBuffer() {
+	w.writer.Reset(io.Discard)
 	w.numRows = 0
+
+	if w.buffer != nil {
+		w.sorting.SortingBuffers.PutBuffer(w.buffer)
+		w.buffer = nil
+	}
 }
 
 func (w *SortingWriter[T]) Write(rows []T) (int, error) {
@@ -208,6 +207,11 @@ func (w *SortingWriter[T]) sortAndWriteBufferedRows() error {
 	rows := w.rows.Rows()
 	defer rows.Close()
 
+	if w.buffer == nil {
+		w.buffer = w.sorting.SortingBuffers.GetBuffer()
+		w.writer.Reset(w.buffer)
+	}
+
 	n, err := CopyRows(w.writer, rows)
 	if err != nil {
 		return err
diff --git a/writer.go b/writer.go
index ad89217..ce95b40 100644
--- a/writer.go
+++ b/writer.go
@@ -857,7 +857,7 @@ func (wb *writerBuffers) swapPageAndScratchBuffers() {
 }
 
 type writerColumn struct {
-	pool  PageBufferPool
+	pool  BufferPool
 	pages []io.ReadWriteSeeker
 
 	columnPath   columnPath
@@ -903,7 +903,7 @@ func (c *writerColumn) reset() {
 		c.dictionary.Reset()
 	}
 	for _, page := range c.pages {
-		c.pool.PutPageBuffer(page)
+		c.pool.PutBuffer(page)
 	}
 	for i := range c.pages {
 		c.pages[i] = nil
@@ -1257,10 +1257,10 @@ func (w *writerColumn) writePageToFilter(page Page) (err error) {
 }
 
 func (c *writerColumn) writePageTo(size int64, writeTo func(io.Writer) (int64, error)) error {
-	buffer := c.pool.GetPageBuffer()
+	buffer := c.pool.GetBuffer()
 	defer func() {
 		if buffer != nil {
-			c.pool.PutPageBuffer(buffer)
+			c.pool.PutBuffer(buffer)
 		}
 	}()
 	written, err := writeTo(buffer)


Test patch:
diff --git a/page_buffer_test.go b/buffer_pool_test.go
similarity index 64%
rename from page_buffer_test.go
rename to buffer_pool_test.go
index 6166640..04437f8 100644
--- a/page_buffer_test.go
+++ b/buffer_pool_test.go
@@ -10,37 +10,37 @@ import (
 	"github.com/segmentio/parquet-go"
 )
 
-func TestPageBufferPool(t *testing.T) {
-	testPageBufferPool(t, parquet.NewPageBufferPool())
+func TestBufferPool(t *testing.T) {
+	testBufferPool(t, parquet.NewBufferPool())
 }
 
 func TestFileBufferPool(t *testing.T) {
-	testPageBufferPool(t, parquet.NewFileBufferPool("/tmp", "buffers.*"))
+	testBufferPool(t, parquet.NewFileBufferPool("/tmp", "buffers.*"))
 }
 
-func testPageBufferPool(t *testing.T, pool parquet.PageBufferPool) {
+func testBufferPool(t *testing.T, pool parquet.BufferPool) {
 	tests := []struct {
 		scenario string
-		function func(*testing.T, parquet.PageBufferPool)
+		function func(*testing.T, parquet.BufferPool)
 	}{
 		{
 			scenario: "write bytes",
-			function: testPageBufferPoolWriteBytes,
+			function: testBufferPoolWriteBytes,
 		},
 
 		{
 			scenario: "write string",
-			function: testPageBufferPoolWriteString,
+			function: testBufferPoolWriteString,
 		},
 
 		{
 			scenario: "copy to buffer",
-			function: testPageBufferPoolCopyToBuffer,
+			function: testBufferPoolCopyToBuffer,
 		},
 
 		{
 			scenario: "copy from buffer",
-			function: testPageBufferPoolCopyFromBuffer,
+			function: testBufferPoolCopyFromBuffer,
 		},
 	}
 
@@ -49,22 +49,25 @@ func testPageBufferPool(t *testing.T, pool parquet.PageBufferPool) {
 	}
 }
 
-func testPageBufferPoolWriteBytes(t *testing.T, pool parquet.PageBufferPool) {
+func testBufferPoolWriteBytes(t *testing.T, pool parquet.BufferPool) {
 	const content = "Hello World!"
 
-	buffer := pool.GetPageBuffer()
+	buffer := pool.GetBuffer()
+	defer pool.PutBuffer(buffer)
+
 	_, err := buffer.Write([]byte(content))
 	if err != nil {
 		t.Fatal(err)
 	}
-
 	assertBufferContent(t, buffer, content)
 }
 
-func testPageBufferPoolWriteString(t *testing.T, pool parquet.PageBufferPool) {
+func testBufferPoolWriteString(t *testing.T, pool parquet.BufferPool) {
 	const content = "Hello World!"
 
-	buffer := pool.GetPageBuffer()
+	buffer := pool.GetBuffer()
+	defer pool.PutBuffer(buffer)
+
 	_, err := io.WriteString(buffer, content)
 	if err != nil {
 		t.Fatal(err)
@@ -73,10 +76,12 @@ func testPageBufferPoolWriteString(t *testing.T, pool parquet.PageBufferPool) {
 	assertBufferContent(t, buffer, content)
 }
 
-func testPageBufferPoolCopyToBuffer(t *testing.T, pool parquet.PageBufferPool) {
+func testBufferPoolCopyToBuffer(t *testing.T, pool parquet.BufferPool) {
 	const content = "ABC"
 
-	buffer := pool.GetPageBuffer()
+	buffer := pool.GetBuffer()
+	defer pool.PutBuffer(buffer)
+
 	reader := strings.NewReader(content)
 	_, err := io.Copy(buffer, struct{ io.Reader }{reader})
 	if err != nil {
@@ -86,10 +91,12 @@ func testPageBufferPoolCopyToBuffer(t *testing.T, pool parquet.PageBufferPool) {
 	assertBufferContent(t, buffer, content)
 }
 
-func testPageBufferPoolCopyFromBuffer(t *testing.T, pool parquet.PageBufferPool) {
+func testBufferPoolCopyFromBuffer(t *testing.T, pool parquet.BufferPool) {
 	const content = "0123456789"
 
-	buffer := pool.GetPageBuffer()
+	buffer := pool.GetBuffer()
+	defer pool.PutBuffer(buffer)
+
 	if _, err := io.WriteString(buffer, content); err != nil {
 		t.Fatal(err)
 	}
diff --git a/sorting_test.go b/sorting_test.go
index 3c5236e..f34b624 100644
--- a/sorting_test.go
+++ b/sorting_test.go
@@ -74,6 +74,9 @@ func TestSortingWriterDropDuplicatedRows(t *testing.T) {
 	buffer := bytes.NewBuffer(nil)
 	writer := parquet.NewSortingWriter[Row](buffer, 99,
 		parquet.SortingWriterConfig(
+			parquet.SortingBuffers(
+				parquet.NewFileBufferPool("", "buffers.*"),
+			),
 			parquet.SortingColumns(
 				parquet.Ascending("value"),
 			),