Repository: segmentio/parquet-go
Pull request #354 (2022-09-29 14:18:39)

Feature: Leveled buffers
This PR introduces leveled buffering to the pools in buffer.go. Each pool internally is divided into 16 different pools. The smallest pool is guaranteed to have buffers in size from 1024->2047. The next is 2048->4095, etc. When a buffer is requested the appropriate subpool is used to guarantee that the buffer will be large enough. This prevents all buffers in the pools to trend towards max requested size. The largest pool contains buffers of 32MB and larger. This is the one sub pool that does not guarantee that it will return a buffer of the appropriate size and a resize may be necessary.

We have seen a total heap reduction of ~25% in the components we tested.

Also, fixed a bug where `resize()` would drop a buffer.

abraithwaite: For my own understanding, this is essentially creating room by allowing smaller buffers to still be created without claiming ownership of a big buffer even when there's very large buffers being allocated for I/O.  Is that understanding correct?

Also, without anything to prevent resize of buffers by the owner when it's claimed from a pool (not sure we'd want to do this), it's possible that a buffer would be taken from a smaller pool and returned to a larger pool, yeah?  I presume this wouldn't be a big issue in practice with the size hints, but I'm curious how we're thinking about that from a `bufferPool` consumer perspective.

Very cool though, thanks!  I think we should let Achille take a look before merging, but overall I think it looks great.


joe-elliott: > For my own understanding, this is essentially creating room by allowing smaller buffers to still be created without claiming ownership of a big buffer even when there's very large buffers being allocated for I/O. Is that understanding correct?

Yup, this will choose more appropriately sized buffers when they are requested. This prevents all buffers trending toward the largest size. It also significantly reduces allocs b/c previously it would be very common to receive a smaller buffer then requested.

> Also, without anything to prevent resize of buffers by the owner when it's claimed from a pool (not sure we'd want to do this), it's possible that a buffer would be taken from a smaller pool and returned to a larger pool, yeah?

Yup. I wanted to build something with as small of an impact to the codebase as possible. So all the existing code that pulls from the pool and then attempts to resize still continues working. This implementation still does have one case in which this is necessary: when you are dealing in buffers > 32MB.

Going to mark this draft real fast. There is one more change I want to make and test. As written if you request a 2047 byte buffer it will return one of exactly this size and then on `put()` it will place that buffer in the 1024 byte bucket. This is a bit wasteful. I'm going to change these lines:

> if sz < basePoolIncrement {
> 	sz = basePoolIncrement
> }

To align to the next pool up.



Patch:
diff --git a/buffer.go b/buffer.go
index 0012655..6527114 100644
--- a/buffer.go
+++ b/buffer.go
@@ -1,6 +1,7 @@
 package parquet
 
 import (
+	"math/bits"
 	"sort"
 	"sync"
 	"sync/atomic"
@@ -306,45 +307,47 @@ func (b *buffer) unref() {
 	}
 }
 
-func (b *buffer) reset() {
-	b.data = b.data[:0]
-}
-
-func (b *buffer) resize(size int) {
-	if cap(b.data) < size {
-		const pageSize = 4096
-		minSize := 2 * cap(b.data)
-		bufferSize := ((size + (pageSize - 1)) / pageSize) * pageSize
-		if bufferSize < minSize {
-			bufferSize = minSize
-		}
-		b.data = make([]byte, size, bufferSize)
-	} else {
-		b.data = b.data[:size]
-	}
-}
-
-func (b *buffer) clone() (clone *buffer) {
-	if b.pool != nil {
-		clone = b.pool.get()
-	} else {
-		clone = &buffer{refc: 1}
-	}
-	clone.data = append(clone.data, b.data...)
-	return clone
-}
+// bufferPool holds a slice of sync.pools used for levelled buffering.
+// the table below shows the pools used for different buffer sizes when both getting
+// and putting a buffer. when allocating a new buffer from a given pool we always choose the
+// min of the put range to guarantee that all gets will have an adequately sized buffer.
+//
+// [pool] : <get range>  : <put range>  : <alloc size>
+// [0]    : 0    -> 1023 : 1024 -> 2047 : 1024
+// [1]    : 1024 -> 2047 : 2048 -> 4095 : 2048
+// [2]    : 2048 -> 4095 : 4096 -> 8191 : 4096
+// ...
+const numPoolBuckets = 16
+const basePoolIncrement = 1024
 
 type bufferPool struct {
-	pool sync.Pool
+	pool [numPoolBuckets]sync.Pool
 }
 
-func (p *bufferPool) get() *buffer {
-	b, _ := p.pool.Get().(*buffer)
+// get returns a buffer from the levelled buffer pool. sz is used to choose the appropriate pool
+func (p *bufferPool) get(sz int) *buffer {
+	i := levelledPoolIndex(sz)
+	b, _ := p.pool[i].Get().(*buffer)
 	if b == nil {
-		b = &buffer{pool: p}
-	} else {
-		b.reset()
+		// align size to the pool
+		poolSize := basePoolIncrement << i
+		if sz > poolSize { // this can occur when the buffer requested is larger than the largest pool
+			poolSize = sz
+		}
+		b = &buffer{
+			data: make([]byte, 0, poolSize),
+			pool: p,
+		}
 	}
+	// if the buffer comes from the largest pool it may not be big enough
+	if cap(b.data) < sz {
+		p.pool[i].Put(b)
+		b = &buffer{
+			data: make([]byte, 0, sz),
+			pool: p,
+		}
+	}
+	b.data = b.data[:sz]
 	b.ref()
 	return b
 }
@@ -353,7 +356,27 @@ func (p *bufferPool) put(b *buffer) {
 	if b.pool != p {
 		panic("BUG: buffer returned to a different pool than the one it was allocated from")
 	}
-	p.pool.Put(b)
+	// if this slice is somehow less then our min pool size, just drop it
+	sz := cap(b.data)
+	if sz < basePoolIncrement {
+		return
+	}
+	i := levelledPoolIndex(sz / 2) // divide by 2 to put the buffer in the level below so it will always be large enough
+	p.pool[i].Put(b)
+}
+
+// levelledPoolIndex returns the index of the pool to use for a buffer of size sz. it never returns
+// an index that will panic
+func levelledPoolIndex(sz int) int {
+	i := sz / basePoolIncrement
+	i = 32 - bits.LeadingZeros32(uint32(i)) // log2
+	if i >= numPoolBuckets {
+		i = numPoolBuckets - 1
+	}
+	if i < 0 {
+		i = 0
+	}
+	return i
 }
 
 var (
diff --git a/buffer_go18.go b/buffer_go18.go
index c073abf..ff90540 100644
--- a/buffer_go18.go
+++ b/buffer_go18.go
@@ -22,7 +22,8 @@ type GenericBuffer[T any] struct {
 // The type parameter T should be a map, struct, or any. Any other types will
 // cause a panic at runtime. Type checking is a lot more effective when the
 // generic parameter is a struct type, using map and interface types is somewhat
-// similar to using a Writer.
+// similar to using a Writer.  If using an interface type for the type parameter,
+// then providing a schema at instantiation is required.
 //
 // If the option list may explicitly declare a schema, it must be compatible
 // with the schema generated from T.
@@ -37,6 +38,10 @@ func NewGenericBuffer[T any](options ...RowGroupOption) *GenericBuffer[T] {
 		config.Schema = schemaOf(dereference(t))
 	}
 
+	if config.Schema == nil {
+		panic("generic buffer must be instantiated with schema or concrete type.")
+	}
+
 	buf := &GenericBuffer[T]{
 		base: Buffer{config: config},
 	}
diff --git a/column.go b/column.go
index 08ccfba..74c09ea 100644
--- a/column.go
+++ b/column.go
@@ -332,8 +332,13 @@ func (cl *columnLoader) open(file *File, path []string) (*Column, error) {
 			// the page headers to determine which compression and encodings are
 			// applied.
 			for _, encoding := range c.chunks[0].MetaData.Encoding {
-				c.encoding = LookupEncoding(encoding)
-				break
+				if c.encoding == nil {
+					c.encoding = LookupEncoding(encoding)
+				}
+				if encoding != format.Plain && encoding != format.RLE {
+					c.encoding = LookupEncoding(encoding)
+					break
+				}
 			}
 			c.compression = LookupCompressionCodec(c.chunks[0].MetaData.Codec)
 		}
@@ -488,10 +493,7 @@ func schemaRepetitionTypeOf(s *format.SchemaElement) format.FieldRepetitionType
 }
 
 func (c *Column) decompress(compressedPageData []byte, uncompressedPageSize int32) (page *buffer, err error) {
-	page = uncompressedPageBufferPool.get()
-	if uncompressedPageSize > 0 {
-		page.resize(int(uncompressedPageSize))
-	}
+	page = uncompressedPageBufferPool.get(int(uncompressedPageSize))
 	page.data, err = c.compression.Decode(page.data, compressedPageData)
 	if err != nil {
 		page.unref()
@@ -628,15 +630,13 @@ func (c *Column) decodeDataPage(header DataPageHeader, numValues int, repetition
 
 	pageKind := pageType.Kind()
 	if pageKind >= 0 && int(pageKind) < len(pageValuesBufferPool) {
-		vbuf = pageValuesBufferPool[pageKind].get()
+		vbuf = pageValuesBufferPool[pageKind].get(int(pageType.EstimateSize(numValues)))
 		defer vbuf.unref()
-		vbuf.resize(int(pageType.EstimateSize(numValues)))
 		pageValues = vbuf.data
 	}
 	if pageKind == ByteArray {
-		obuf = pageOffsetsBufferPool.get()
+		obuf = pageOffsetsBufferPool.get(4 * (numValues + 1))
 		defer obuf.unref()
-		obuf.resize(4 * (numValues + 1))
 		pageOffsets = unsafecast.BytesToUint32(obuf.data)
 	}
 
@@ -707,7 +707,7 @@ func decodeLevelsV2(enc encoding.Encoding, numValues int, data []byte, length in
 }
 
 func decodeLevels(enc encoding.Encoding, numValues int, data []byte) (levels *buffer, err error) {
-	levels = levelsBufferPool.get()
+	levels = levelsBufferPool.get(numValues)
 	levels.data, err = enc.DecodeLevels(levels.data, data)
 	if err != nil {
 		levels.unref()
diff --git a/file.go b/file.go
index f19a814..7a47570 100644
--- a/file.go
+++ b/file.go
@@ -567,11 +567,9 @@ func (f *filePages) readDictionary() error {
 		return err
 	}
 
-	page := compressedPageBufferPool.get()
+	page := compressedPageBufferPool.get(int(header.CompressedPageSize))
 	defer page.unref()
 
-	page.resize(int(header.CompressedPageSize))
-
 	if _, err := io.ReadFull(rbuf, page.data); err != nil {
 		return err
 	}
@@ -619,11 +617,9 @@ func (f *filePages) readDataPageV2(header *format.PageHeader, page *buffer) (Pag
 }
 
 func (f *filePages) readPage(header *format.PageHeader, reader *bufio.Reader) (*buffer, error) {
-	page := compressedPageBufferPool.get()
+	page := compressedPageBufferPool.get(int(header.CompressedPageSize))
 	defer page.unref()
 
-	page.resize(int(header.CompressedPageSize))
-
 	if _, err := io.ReadFull(reader, page.data); err != nil {
 		return nil, err
 	}
diff --git a/internal/debug/finalizer_off.go b/internal/debug/finalizer_off.go
new file mode 100644
index 0000000..acabe9c
--- /dev/null
+++ b/internal/debug/finalizer_off.go
@@ -0,0 +1,6 @@
+//go:build debug
+
+package debug
+
+// SetFinalizer is a no-op when the debug tag is specified.
+func SetFinalizer(interface{}, interface{}) {}
diff --git a/internal/debug/finalizer_on.go b/internal/debug/finalizer_on.go
new file mode 100644
index 0000000..fc188cc
--- /dev/null
+++ b/internal/debug/finalizer_on.go
@@ -0,0 +1,7 @@
+//go:build !debug
+
+package debug
+
+import "runtime"
+
+func SetFinalizer(obj, finalizer interface{}) { runtime.SetFinalizer(obj, finalizer) }
diff --git a/page.go b/page.go
index 9c991ac..0efeeac 100644
--- a/page.go
+++ b/page.go
@@ -4,11 +4,11 @@ import (
 	"bytes"
 	"fmt"
 	"io"
-	"runtime"
 
 	"github.com/segmentio/parquet-go/deprecated"
 	"github.com/segmentio/parquet-go/encoding"
 	"github.com/segmentio/parquet-go/internal/bitpack"
+	"github.com/segmentio/parquet-go/internal/debug"
 )
 
 // Page values represent sequences of parquet values. From the Parquet
@@ -129,7 +129,7 @@ func AsyncPages(pages Pages) Pages {
 	// If the pages object gets garbage collected without Close being called,
 	// this finalizer would ensure that the goroutine is stopped and doesn't
 	// leak.
-	runtime.SetFinalizer(p, func(p *asyncPages) { p.Close() })
+	debug.SetFinalizer(p, func(p *asyncPages) { p.Close() })
 	return p
 }
 
diff --git a/row_group.go b/row_group.go
index a026ad8..698cbf6 100644
--- a/row_group.go
+++ b/row_group.go
@@ -3,7 +3,8 @@ package parquet
 import (
 	"fmt"
 	"io"
-	"runtime"
+
+	"github.com/segmentio/parquet-go/internal/debug"
 )
 
 // RowGroup is an interface representing a parquet row group. From the Parquet
@@ -294,7 +295,7 @@ func (r *rowGroupRows) init() {
 	// This finalizer is used to ensure that the goroutines started by calling
 	// init on the underlying page readers will be shutdown in the event that
 	// Close isn't called and the rowGroupRows object is garbage collected.
-	runtime.SetFinalizer(r, func(r *rowGroupRows) { r.Close() })
+	debug.SetFinalizer(r, func(r *rowGroupRows) { r.Close() })
 }
 
 func (r *rowGroupRows) clear() {
diff --git a/writer_go18.go b/writer_go18.go
index b1f9598..92e0d0e 100644
--- a/writer_go18.go
+++ b/writer_go18.go
@@ -78,6 +78,10 @@ func NewGenericWriter[T any](output io.Writer, options ...WriterOption) *Generic
 		config.Schema = schema
 	}
 
+	if config.Schema == nil {
+		panic("generic writer must be instantiated with schema or concrete type.")
+	}
+
 	return &GenericWriter[T]{
 		base: Writer{
 			output: output,


Test patch:
diff --git a/buffer_go18_test.go b/buffer_go18_test.go
index 038f6df..aea08f9 100644
--- a/buffer_go18_test.go
+++ b/buffer_go18_test.go
@@ -182,11 +182,6 @@ func TestIssue327(t *testing.T) {
 	})
 }
 
-func TestIssue343(t *testing.T) {
-	// must not panic
-	_ = parquet.NewGenericBuffer[any]()
-}
-
 func TestIssue346(t *testing.T) {
 	type TestType struct {
 		Key int
@@ -199,3 +194,24 @@ func TestIssue346(t *testing.T) {
 	data[0] = TestType{Key: 0}
 	_, _ = buffer.Write(data)
 }
+
+func TestIssue347(t *testing.T) {
+	type TestType struct {
+		Key int
+	}
+
+	// instantiating with concrete type shouldn't panic
+	_ = parquet.NewGenericBuffer[TestType]()
+
+	// instantiating with schema and interface type parameter shouldn't panic
+	schema := parquet.SchemaOf(TestType{})
+	_ = parquet.NewGenericBuffer[any](schema)
+
+	defer func() {
+		if r := recover(); r == nil {
+			t.Errorf("instantiating generic buffer without schema and with interface " +
+				"type parameter should panic")
+		}
+	}()
+	_ = parquet.NewGenericBuffer[any]()
+}
diff --git a/buffer_internal_test.go b/buffer_internal_test.go
new file mode 100644
index 0000000..e2a0f42
--- /dev/null
+++ b/buffer_internal_test.go
@@ -0,0 +1,57 @@
+package parquet
+
+import (
+	"math"
+	"math/rand"
+	"testing"
+)
+
+func TestBufferAlwaysCorrectSize(t *testing.T) {
+	var p bufferPool
+	for i := 0; i < 1000; i++ {
+		sz := rand.Intn(1024 * 1024)
+		buff := p.get(sz)
+		if len(buff.data) != sz {
+			t.Errorf("Expected buffer of size %d, got %d", sz, len(buff.data))
+		}
+		p.put(buff)
+	}
+}
+
+func TestLevelledPoolIndex(t *testing.T) {
+	tcs := []struct {
+		sz       int
+		expected int
+	}{
+		{
+			sz:       1023,
+			expected: 0,
+		},
+		{
+			sz:       1024,
+			expected: 1,
+		},
+		{
+			sz:       -1,
+			expected: 0,
+		},
+		{
+			sz:       16*1024*1024 - 1,
+			expected: 14,
+		},
+		{
+			sz:       16 * 1024 * 1024,
+			expected: 15,
+		},
+		{
+			sz:       math.MaxInt,
+			expected: 15,
+		},
+	}
+
+	for _, tc := range tcs {
+		if actual := levelledPoolIndex(tc.sz); actual != tc.expected {
+			t.Errorf("Expected index %d for size %d, got %d", tc.expected, tc.sz, actual)
+		}
+	}
+}
diff --git a/writer_go18_test.go b/writer_go18_test.go
index 9036b91..c26cb97 100644
--- a/writer_go18_test.go
+++ b/writer_go18_test.go
@@ -293,3 +293,25 @@ func TestIssue302(t *testing.T) {
 		t.Run(test.name, test.fn)
 	}
 }
+
+func TestIssue347Writer(t *testing.T) {
+	type TestType struct {
+		Key int
+	}
+
+	b := new(bytes.Buffer)
+	// instantiating with concrete type shouldn't panic
+	_ = parquet.NewGenericWriter[TestType](b)
+
+	// instantiating with schema and interface type parameter shouldn't panic
+	schema := parquet.SchemaOf(TestType{})
+	_ = parquet.NewGenericWriter[any](b, schema)
+
+	defer func() {
+		if r := recover(); r == nil {
+			t.Errorf("instantiating generic buffer without schema and with interface " +
+				"type parameter should panic")
+		}
+	}()
+	_ = parquet.NewGenericWriter[any](b)
+}