Repository: segmentio/parquet-go
Pull request #440 (2022-12-07 23:34:25)

Reduce memory wasted by buffer pooling
With this PR, I investigated ways to reduce the memory waste caused by using buffer pools internally. At first I found an issue where in some edge cases, the buffers would always be released in a bucket different than the one they were allocated from (e.g. when the buffer size was a power of two, which was frequent considering the default target page size is 256KiB). Even with the fix, when allocating multi-megabyte buffers a lot of memory would go to waste with the allocation strategy we were using.

I would like to modify the buffer pool to use a different sizing of buckets, growing 2x the bucket size for small buffers, then 1.5x for larger buffers (greater than 256KiB).

Eventually it might be interesting to give applications more control over these settings, either through environment variable configuration or explicitly in the code, but I'd like to experiment more with this new strategy before we add more complexity.

achille-roussel: > I was considering a set of benchmarks focused on memory consumption that use the files in ./testdata. Then at least we'd have some kind of guide when making choices about PRs like this. WDYT?

*(copied from internal Slack discussions)*

I don't know if we'll successfully write such a test suite tho, memory utilization tends to be very dependent on access patterns and the shape of data; we could probably detect regressions with the benchmarks but that may not necessarily map to what we observe in production environments...

I guess the goal would be to ensure that we gate regressions before they reach production, that may reduce engineering time spent on addressing them but I assume that it wouldn't reduce the amount of time spent validating changes prior to reaching production, we'd still want to test in a staging/canary.

That's just my opinion on whether it's worth the time investment to build the benchmarks in the first place. I won't oppose adding more tests to parquet-go if they are valuable for your use case :slightly_smiling_face:

joe-elliott: agree, it's tough. i was hoping for a lower barrier of "directionally correct" so even if the tests weren't perfect they'd roughly indicate whether a change was positive, negative or neutral with regard to memory usage. the danger of course being overindexing on the results of these benchmarks.

the other reason i like the idea of the tests is it will encourage more applications to PR example files to the `./testdata` folder. We already have a Tempo file in there and it would be nice to get examples from other applications.

Regardless, thank you for the honest feedback. If I can think of something that feels worthwhile I will submit a PR and we can discuss :pray: 


Patch:
diff --git a/buffer.go b/buffer.go
index aec3b83..8f8cc76 100644
--- a/buffer.go
+++ b/buffer.go
@@ -2,7 +2,6 @@ package parquet
 
 import (
 	"log"
-	"math/bits"
 	"runtime"
 	"sort"
 	"sync"
@@ -318,26 +317,25 @@ func monitorBufferRelease(b *buffer) {
 	}
 }
 
-// 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 {
-	bucket [numPoolBuckets]sync.Pool
-}
-
-func (p *bufferPool) newBuffer(size int) *buffer {
+	// Buckets are split in two groups for short and large buffers. In the short
+	// buffer group (below 256KB), the growth rate between each bucket is 2. The
+	// growth rate changes to 1.5 in the larger buffer group.
+	//
+	// Short buffer buckets:
+	// ---------------------
+	//   4K, 8K, 16K, 32K, 64K, 128K, 256K
+	//
+	// Large buffer buckets:
+	// ---------------------
+	//   364K, 546K, 819K ...
+	//
+	buckets [bufferPoolBucketCount]sync.Pool
+}
+
+func (p *bufferPool) newBuffer(bufferSize, bucketSize int) *buffer {
 	b := &buffer{
-		data: make([]byte, size),
+		data: make([]byte, bufferSize, bucketSize),
 		refc: 1,
 		pool: p,
 	}
@@ -350,25 +348,21 @@ func (p *bufferPool) newBuffer(size int) *buffer {
 
 // get returns a buffer from the levelled buffer pool. size is used to choose
 // the appropriate pool.
-func (p *bufferPool) get(size int) *buffer {
-	i := levelledPoolIndex(size)
-	b, _ := p.bucket[i].Get().(*buffer)
-	switch {
-	case b == nil:
-		// align size to the pool
-		bucketSize := basePoolIncrement << i
-		if size > bucketSize { // this can occur when the buffer requested is larger than the largest pool
-			bucketSize = size
-		}
-		b = p.newBuffer(bucketSize)
-	case size <= cap(b.data):
+func (p *bufferPool) get(bufferSize int) *buffer {
+	bucketIndex, bucketSize := bufferPoolBucketIndexAndSizeOfGet(bufferSize)
+
+	b := (*buffer)(nil)
+	if bucketIndex >= 0 {
+		b, _ = p.buckets[bucketIndex].Get().(*buffer)
+	}
+
+	if b == nil {
+		b = p.newBuffer(bufferSize, bucketSize)
+	} else {
+		b.data = b.data[:bufferSize]
 		b.ref()
-	default:
-		// if the buffer comes from the largest pool it may not be big enough
-		p.bucket[i].Put(b)
-		b = p.newBuffer(size)
 	}
-	b.data = b.data[:size]
+
 	if debug.TRACEBUF > 0 {
 		b.stack = b.stack[:runtime.Stack(b.stack[:cap(b.stack)], false)]
 	}
@@ -382,27 +376,53 @@ func (p *bufferPool) put(b *buffer) {
 	if b.refCount() != 0 {
 		panic("BUG: buffer returned to pool with a non-zero reference count")
 	}
-	// if this slice is somehow less then our min pool size, just drop it
-	size := cap(b.data)
-	if size < basePoolIncrement {
-		return
+	if bucketIndex, _ := bufferPoolBucketIndexAndSizeOfPut(cap(b.data)); bucketIndex >= 0 {
+		p.buckets[bucketIndex].Put(b)
 	}
-	i := levelledPoolIndex(size / 2) // divide by 2 to put the buffer in the level below so it will always be large enough
-	p.bucket[i].Put(b)
 }
 
-// levelledPoolIndex returns the index of the pool to use for a buffer of the
-// given size. It never returns an index that will panic.
-func levelledPoolIndex(size int) int {
-	i := size / basePoolIncrement
-	i = 32 - bits.LeadingZeros32(uint32(i)) // log2
-	if i >= numPoolBuckets {
-		i = numPoolBuckets - 1
+const (
+	bufferPoolBucketCount         = 32
+	bufferPoolMinSize             = 4096
+	bufferPoolLastShortBucketSize = 262144
+)
+
+func bufferPoolNextSize(size int) int {
+	if size < bufferPoolLastShortBucketSize {
+		return size * 2
+	} else {
+		return size + (size / 2)
+	}
+}
+
+func bufferPoolBucketIndexAndSizeOfGet(size int) (int, int) {
+	limit := bufferPoolMinSize
+
+	for i := 0; i < bufferPoolBucketCount; i++ {
+		if size <= limit {
+			return i, limit
+		}
+		limit = bufferPoolNextSize(limit)
 	}
-	if i < 0 {
-		i = 0
+
+	return -1, size
+}
+
+func bufferPoolBucketIndexAndSizeOfPut(size int) (int, int) {
+	// When releasing buffers, some may have a capacity that is not one of the
+	// bucket sizes (due to the use of append for example). In this case, we
+	// have to put the buffer is the highest bucket with a size less or equal
+	// to the buffer capacity.
+	if limit := bufferPoolMinSize; size >= limit {
+		for i := 0; i < bufferPoolBucketCount; i++ {
+			n := bufferPoolNextSize(limit)
+			if size < n {
+				return i, limit
+			}
+			limit = n
+		}
 	}
-	return i
+	return -1, size
 }
 
 var (


Test patch:
diff --git a/buffer_internal_test.go b/buffer_internal_test.go
index 5502d4b..8d6a98e 100644
--- a/buffer_internal_test.go
+++ b/buffer_internal_test.go
@@ -1,7 +1,7 @@
 package parquet
 
 import (
-	"math"
+	"fmt"
 	"math/rand"
 	"testing"
 )
@@ -12,46 +12,45 @@ func TestBufferAlwaysCorrectSize(t *testing.T) {
 		n := rand.Intn(1024 * 1024)
 		b := p.get(n)
 		if len(b.data) != n {
-			t.Errorf("Expected buffer of size %d, got %d", n, len(b.data))
+			t.Fatalf("Expected buffer of size %d, got %d", n, len(b.data))
 		}
 		b.unref()
 	}
 }
 
-func TestLevelledPoolIndex(t *testing.T) {
-	tcs := []struct {
-		size     int
-		expected int
+func TestBufferPoolBucketIndexAndSizeOf(t *testing.T) {
+	tests := []struct {
+		size        int
+		bucketIndex int
+		bucketSize  int
 	}{
-		{
-			size:     1023,
-			expected: 0,
-		},
-		{
-			size:     1024,
-			expected: 1,
-		},
-		{
-			size:     -1,
-			expected: 0,
-		},
-		{
-			size:     16*1024*1024 - 1,
-			expected: 14,
-		},
-		{
-			size:     16 * 1024 * 1024,
-			expected: 15,
-		},
-		{
-			size:     math.MaxInt,
-			expected: 15,
-		},
+		{size: 0, bucketIndex: 0, bucketSize: 4096},
+		{size: 1, bucketIndex: 0, bucketSize: 4096},
+		{size: 2049, bucketIndex: 0, bucketSize: 4096},
+		{size: 4096, bucketIndex: 0, bucketSize: 4096},
+		{size: 4097, bucketIndex: 1, bucketSize: 8192},
+		{size: 8192, bucketIndex: 1, bucketSize: 8192},
+		{size: 8193, bucketIndex: 2, bucketSize: 16384},
+		{size: 16384, bucketIndex: 2, bucketSize: 16384},
+		{size: 16385, bucketIndex: 3, bucketSize: 32768},
+		{size: 32768, bucketIndex: 3, bucketSize: 32768},
+		{size: 32769, bucketIndex: 4, bucketSize: 65536},
+		{size: 262143, bucketIndex: 6, bucketSize: 262144},
+		{size: 262144, bucketIndex: 6, bucketSize: 262144},
+		{size: 262145, bucketIndex: 7, bucketSize: 393216},
 	}
 
-	for _, tc := range tcs {
-		if actual := levelledPoolIndex(tc.size); actual != tc.expected {
-			t.Errorf("Expected index %d for size %d, got %d", tc.expected, tc.size, actual)
-		}
+	for _, test := range tests {
+		t.Run(fmt.Sprintf("size=%d", test.size), func(t *testing.T) {
+			bucketIndex, bucketSize := bufferPoolBucketIndexAndSizeOfGet(test.size)
+
+			if bucketIndex != test.bucketIndex {
+				t.Errorf("wrong bucket index, want %d but got %d", test.bucketIndex, bucketIndex)
+			}
+
+			if bucketSize != test.bucketSize {
+				t.Errorf("wrong bucket size, want %d but got %d", test.bucketSize, bucketSize)
+			}
+		})
 	}
 }