Repository: segmentio/parquet-go
Pull request #130 (2022-04-08 01:47:55)

fix file buffer pool
While working on using the file buffer pool from this package, I ran into an issue caused by the fact that `os.File` has a single cursor for read AND write operations, which resulted in programs being unable to read back the data that they had written to the buffers.

This PR addresses the issue by adding a second cursor to track the position of read operations, mimicking the behavior of `bytes.Buffer` on top of disk storage.

I also added tests, documentation, and fixed a few typos.


Patch:
diff --git a/.words b/.words
index abc9140..d7ff365 100644
--- a/.words
+++ b/.words
@@ -4,7 +4,24 @@ Twilio
 bottlenecked
 decompressors
 int96
+millis
+nanos
 reindexing
 repositions
 schemas
 ColumnPages
+PageIndex
+Zstandard
+xxHash
+cardinality
+enums
+32bit
+dic
+Blart
+Versenwald
+purego
+stdlib
+unscaled
+cespare
+bitset
+checksumming
diff --git a/README.md b/README.md
index 82a1d19..6f3bc68 100644
--- a/README.md
+++ b/README.md
@@ -496,6 +496,41 @@ achieved by using an intermediary buffering layer with `parquet.Buffer`.
 See [parquet.RowGroup](https://pkg.go.dev/github.com/segmentio/parquet-go#RowGroup)
 for the full interface documentation.
 
+### C. Using on-disk page buffers
+
+When generating parquet files, the writer needs to buffer all pages before it
+can create the row group. This may require significant amounts of memory as the
+entire file content must be buffered prior to generating it. In some cases, the
+files might even be larger than the amount of memory available to the program.
+
+The `parquet.Writer` can be configured to use disk storage instead as a scratch
+buffer when generating files, by configuring a different page buffer pool using
+the `parquet.ColumnPageBuffers` option and `parquet.PageBufferPool` interface.
+
+The parquet-go package provides an implementation of the interface which uses
+temporary files to store pages while a file is generated, allowing programs to
+use local storage as swap space to hold pages and keep memory utilization to a
+minimum. The following example demonstrates how to configure a parquet writer
+to use on-disk page buffers:
+
+```go
+writer := parquet.NewWriter(output,
+    parquet.ColumnPageBuffers(
+        parquet.NewFileBufferPool("", "buffers.*"),
+    ),
+)
+```
+
+When a row group is complete, pages buffered to disk need to be copied back to
+the output file. This results in doubling I/O operations and storage space
+requirements (the system needs to have enough free disk space to hold two copies
+of the file). The resulting write amplification can often be optimized away by
+the kernel if the file system supports copy-on-write of disk pages since copies
+between `os.File` instances are optimized using `copy_file_range(2)` (on linux).
+
+See [parquet.PageBufferPool](https://pkg.go.dev/github.com/segmentio/parquet-go#PageBufferPool)
+for the full interface documentation.
+
 ## Maintenance
 
 The project is hosted and maintained by Twilio; we welcome external contributors
diff --git a/format/parquet.go b/format/parquet.go
index e81f532..0461079 100644
--- a/format/parquet.go
+++ b/format/parquet.go
@@ -6,7 +6,7 @@ import (
 	"github.com/segmentio/parquet-go/deprecated"
 )
 
-// Types supported by Parquet. These types are intended to be used incombination
+// Types supported by Parquet. These types are intended to be used in combination
 // with the encodings to control the on disk storage format. For example INT16
 // is not included as a type since a good encoding of INT32 would handle this.
 type Type int32
@@ -917,7 +917,7 @@ type ColumnIndex struct {
 	MinValues [][]byte `thrift:"2,required"`
 	MaxValues [][]byte `thrift:"3,required"`
 
-	// Stores whether both min_values and max_values are orderd and if so, in
+	// Stores whether both min_values and max_values are ordered and if so, in
 	// which direction. This allows readers to perform binary searches in both
 	// lists. Readers cannot assume that max_values[i] <= min_values[i+1], even
 	// if the lists are ordered.
@@ -992,7 +992,7 @@ type FileMetaData struct {
 	//
 	// Without column_orders, the meaning of the min_value and max_value fields
 	// in the Statistics object and the ColumnIndex object is undefined. To ensure
-	// well-defined behaviour, if these fields are written to a Parquet file,
+	// well-defined behavior, if these fields are written to a Parquet file,
 	// column_orders must be written as well.
 	//
 	// The obsolete min and max fields in the Statistics object are always sorted
diff --git a/page.go b/page.go
index 4c7e77d..2059ebb 100644
--- a/page.go
+++ b/page.go
@@ -1,12 +1,8 @@
 package parquet
 
 import (
-	"bytes"
 	"fmt"
 	"io"
-	"os"
-	"path/filepath"
-	"sync"
 
 	"github.com/segmentio/parquet-go/deprecated"
 	"github.com/segmentio/parquet-go/encoding"
@@ -811,75 +807,3 @@ func (r *fixedLenByteArrayPageReader) ReadValues(values []Value) (n int, err err
 	}
 	return n, err
 }
-
-type PageBufferPool interface {
-	GetPageBuffer() io.ReadWriter
-	PutPageBuffer(io.ReadWriter)
-}
-
-func NewPageBufferPool() PageBufferPool { return new(pageBufferPool) }
-
-type pageBufferPool struct{ sync.Pool }
-
-func (pool *pageBufferPool) GetPageBuffer() io.ReadWriter {
-	b, _ := pool.Get().(*bytes.Buffer)
-	if b == nil {
-		b = new(bytes.Buffer)
-	} else {
-		b.Reset()
-	}
-	return b
-}
-
-func (pool *pageBufferPool) PutPageBuffer(buf io.ReadWriter) {
-	if b, _ := buf.(*bytes.Buffer); b != nil {
-		pool.Put(b)
-	}
-}
-
-type fileBufferPool struct {
-	err     error
-	tempdir string
-	pattern string
-}
-
-func NewFileBufferPool(tempdir, pattern string) PageBufferPool {
-	pool := &fileBufferPool{
-		tempdir: tempdir,
-		pattern: pattern,
-	}
-	pool.tempdir, pool.err = filepath.Abs(pool.tempdir)
-	return pool
-}
-
-func (pool *fileBufferPool) GetPageBuffer() io.ReadWriter {
-	if pool.err != nil {
-		return &errorBuffer{err: pool.err}
-	}
-	f, err := os.CreateTemp(pool.tempdir, pool.pattern)
-	if err != nil {
-		return &errorBuffer{err: err}
-	}
-	return f
-}
-
-func (pool *fileBufferPool) PutPageBuffer(buf io.ReadWriter) {
-	if f, _ := buf.(*os.File); f != nil {
-		defer f.Close()
-		os.Remove(f.Name())
-	}
-}
-
-type errorBuffer struct{ err error }
-
-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 }
-
-var (
-	defaultPageBufferPool pageBufferPool
-
-	_ io.ReaderFrom = (*errorBuffer)(nil)
-	_ io.WriterTo   = (*errorBuffer)(nil)
-)
diff --git a/page_buffer.go b/page_buffer.go
new file mode 100644
index 0000000..8dcd158
--- /dev/null
+++ b/page_buffer.go
@@ -0,0 +1,140 @@
+package parquet
+
+import (
+	"bytes"
+	"io"
+	"os"
+	"path/filepath"
+	"sync"
+)
+
+// PageBufferPool is an interface abstracting the underlying implementation of
+// page buffer pools.
+//
+// The parquet-go package provides two implementations of this interface, one
+// backed by in-memory buffers (on the Go heap), and the other using temporary
+// files on disk.
+//
+// Applications which need finer grain control over the allocation and retention
+// 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
+// goroutines.
+type PageBufferPool interface {
+	// GetPageBuffer is called when a parquet writer needs to acquires a new
+	// page buffer from the pool.
+	GetPageBuffer() io.ReadWriter
+
+	// PutPageBuffer 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
+	// pool, and that it will not use them anymore after the call.
+	PutPageBuffer(io.ReadWriter)
+}
+
+// NewPageBufferPool 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) }
+
+type pageBufferPool struct{ sync.Pool }
+
+func (pool *pageBufferPool) GetPageBuffer() io.ReadWriter {
+	b, _ := pool.Get().(*bytes.Buffer)
+	if b == nil {
+		b = new(bytes.Buffer)
+	} else {
+		b.Reset()
+	}
+	return b
+}
+
+func (pool *pageBufferPool) PutPageBuffer(buf io.ReadWriter) {
+	if b, _ := buf.(*bytes.Buffer); b != nil {
+		pool.Put(b)
+	}
+}
+
+type fileBufferPool struct {
+	err     error
+	tempdir string
+	pattern string
+}
+
+// NewFileBufferPool creates a new on-disk page buffer pool.
+func NewFileBufferPool(tempdir, pattern string) PageBufferPool {
+	pool := &fileBufferPool{
+		tempdir: tempdir,
+		pattern: pattern,
+	}
+	pool.tempdir, pool.err = filepath.Abs(pool.tempdir)
+	return pool
+}
+
+func (pool *fileBufferPool) GetPageBuffer() io.ReadWriter {
+	if pool.err != nil {
+		return &errorBuffer{err: pool.err}
+	}
+	f, err := os.CreateTemp(pool.tempdir, pool.pattern)
+	if err != nil {
+		return &errorBuffer{err: err}
+	}
+	return &fileBuffer{file: f}
+}
+
+func (pool *fileBufferPool) PutPageBuffer(buf io.ReadWriter) {
+	if f, _ := buf.(*fileBuffer); f != nil {
+		defer f.file.Close()
+		os.Remove(f.file.Name())
+	}
+}
+
+type fileBuffer struct {
+	file *os.File
+	seek int64
+}
+
+func (buf *fileBuffer) Read(b []byte) (int, error) {
+	// The *os.File tracks a single cursor which we use for write operations to
+	// support appending to the buffer. We need a second cursor for reads which
+	// is tracked by the buf.seek field, using ReadAt to read from the file at
+	// the current read position.
+	n, err := buf.file.ReadAt(b, buf.seek)
+	buf.seek += int64(n)
+	return n, err
+}
+
+func (buf *fileBuffer) ReadFrom(r io.Reader) (int64, error) {
+	return buf.file.ReadFrom(r)
+}
+
+func (buf *fileBuffer) Write(b []byte) (int, error) {
+	return buf.file.Write(b)
+}
+
+func (buf *fileBuffer) WriteString(s string) (int, error) {
+	return buf.file.WriteString(s)
+}
+
+type errorBuffer struct{ err error }
+
+func (buf *errorBuffer) Read([]byte) (int, error)          { return 0, buf.err }
+func (buf *errorBuffer) Write([]byte) (int, error)         { return 0, buf.err }
+func (buf *errorBuffer) WriteString(string) (int, error)   { return 0, buf.err }
+func (buf *errorBuffer) ReadFrom(io.Reader) (int64, error) { return 0, buf.err }
+func (buf *errorBuffer) WriteTo(io.Writer) (int64, error)  { return 0, buf.err }
+
+var (
+	defaultPageBufferPool pageBufferPool
+
+	_ io.ReaderFrom   = (*fileBuffer)(nil)
+	_ io.StringWriter = (*fileBuffer)(nil)
+
+	_ io.ReaderFrom   = (*errorBuffer)(nil)
+	_ io.WriterTo     = (*errorBuffer)(nil)
+	_ io.StringWriter = (*errorBuffer)(nil)
+)


Test patch:
diff --git a/page_buffer_test.go b/page_buffer_test.go
new file mode 100644
index 0000000..6ae810a
--- /dev/null
+++ b/page_buffer_test.go
@@ -0,0 +1,112 @@
+package parquet_test
+
+import (
+	"bytes"
+	"io"
+	"strings"
+	"testing"
+	"testing/iotest"
+
+	"github.com/segmentio/parquet-go"
+)
+
+func TestPageBufferPool(t *testing.T) {
+	testPageBufferPool(t, parquet.NewPageBufferPool())
+}
+
+func TestFileBufferPool(t *testing.T) {
+	testPageBufferPool(t, parquet.NewFileBufferPool("", "buffers.*"))
+}
+
+func testPageBufferPool(t *testing.T, pool parquet.PageBufferPool) {
+	tests := []struct {
+		scenario string
+		function func(*testing.T, parquet.PageBufferPool)
+	}{
+		{
+			scenario: "write bytes",
+			function: testPageBufferPoolWriteBytes,
+		},
+
+		{
+			scenario: "write string",
+			function: testPageBufferPoolWriteString,
+		},
+
+		{
+			scenario: "copy to buffer",
+			function: testPageBufferPoolCopyToBuffer,
+		},
+
+		{
+			scenario: "copy from buffer",
+			function: testPageBufferPoolCopyFromBuffer,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.scenario, func(t *testing.T) { test.function(t, pool) })
+	}
+}
+
+func testPageBufferPoolWriteBytes(t *testing.T, pool parquet.PageBufferPool) {
+	const content = "Hello World!"
+
+	buffer := pool.GetPageBuffer()
+	_, err := buffer.Write([]byte(content))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assertBufferContent(t, buffer, content)
+}
+
+func testPageBufferPoolWriteString(t *testing.T, pool parquet.PageBufferPool) {
+	const content = "Hello World!"
+
+	buffer := pool.GetPageBuffer()
+	_, err := io.WriteString(buffer, content)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assertBufferContent(t, buffer, content)
+}
+
+func testPageBufferPoolCopyToBuffer(t *testing.T, pool parquet.PageBufferPool) {
+	const content = "ABC"
+
+	buffer := pool.GetPageBuffer()
+	reader := strings.NewReader(content)
+	_, err := io.Copy(buffer, struct{ io.Reader }{reader})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assertBufferContent(t, buffer, content)
+}
+
+func testPageBufferPoolCopyFromBuffer(t *testing.T, pool parquet.PageBufferPool) {
+	const content = "0123456789"
+
+	buffer := pool.GetPageBuffer()
+	if _, err := io.WriteString(buffer, content); err != nil {
+		t.Fatal(err)
+	}
+
+	writer := new(bytes.Buffer)
+	_, err := io.Copy(struct{ io.Writer }{writer}, buffer)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assertBufferContent(t, writer, content)
+}
+
+func assertBufferContent(t *testing.T, b io.Reader, s string) {
+	t.Helper()
+
+	if err := iotest.TestReader(b, []byte(s)); err != nil {
+		t.Error(err)
+	}
+}