Implement buffer for TupleBuilder

go/libraries/doltcore/sqle/index/dolt_index.go

@@ -1149,6 +1149,7 @@ func (di *doltIndex) keysToTuple(ctx *sql.Context, keys []interface{}) (types.Tu
 }
 
 var sharePool = pool.NewBuffPool()
+var realPool = pool.NewRealBuffPool()
 
 func maybeGetKeyBuilder(idx durable.Index) *val.TupleBuilder {
 	if types.IsFormat_DOLT(idx.Format()) {

go/libraries/doltcore/sqle/index/prolly_index_iter.go

@@ -38,6 +38,7 @@ type prollyIndexIter struct {
 	// into primary index keys
 	pkMap val.OrdinalMapping
 	pkBld *val.TupleBuilder
+	pkBuf fakePool
 
 	// keyMap and valMap transform tuples from
 	// primary row storage into sql.Row's
@@ -50,6 +51,20 @@ type prollyIndexIter struct {
 
 var _ sql.RowIter = prollyIndexIter{}
 
+type fakePool []byte
+
+func NewFakePool() fakePool {
+	return make(fakePool, val.MaxTupleDataSize)
+}
+
+func (p fakePool) Get(size uint64) []byte {
+	return p[:size]
+}
+
+func (p fakePool) GetSlices(size uint64) [][]byte {
+	panic("not used")
+}
+
 // NewProllyIndexIter returns a new prollyIndexIter.
 func newProllyIndexIter(
 	ctx *sql.Context,

go/store/pool/buffer_pool.go

@@ -14,6 +14,8 @@
 
 package pool
 
+import "sync"
+
 type BuffPool interface {
 	Get(size uint64) []byte
 	GetSlices(size uint64) [][]byte
@@ -32,3 +34,30 @@ func (bp buffPool) Get(size uint64) []byte {
 func (bp buffPool) GetSlices(size uint64) [][]byte {
 	return make([][]byte, size)
 }
+
+var realBuffPool = sync.Pool{
+	New: func() any {
+		return make([]byte, 0, 65504) // This is apparently MaxTupleDataSize
+	},
+}
+
+type realBuff struct {
+	buf []byte
+	pos uint64
+}
+
+func NewRealBuffPool() *realBuff {
+	return &realBuff{
+		buf: make([]byte, 65504),
+	}
+}
+
+func (b *realBuff) Get(size uint64) []byte {
+	//clear(b.buf[:b.pos])
+	//b.pos = size
+	return b.buf[:size]
+}
+
+func (b *realBuff) GetSlices(size uint64) [][]byte {
+	panic("this is unused")
+}

go/store/val/tuple.go

@@ -68,6 +68,9 @@ var EmptyTuple = Tuple([]byte{0, 0})
 
 func NewTuple(pool pool.BuffPool, values ...[]byte) Tuple {
 	values = trimNullSuffix(values)
+	if len(values) > MaxTupleFields {
+		panic("tuple field maxIdx exceeds maximum")
+	}
 
 	var count int
 	var pos ByteSize
@@ -78,9 +81,6 @@ func NewTuple(pool pool.BuffPool, values ...[]byte) Tuple {
 		count++
 		pos += sizeOf(v)
 	}
-	if len(values) > MaxTupleFields {
-		panic("tuple field maxIdx exceeds maximum")
-	}
 	if pos > MaxTupleDataSize {
 		panic("tuple data size exceeds maximum")
 	}

go/store/val/tuple_builder.go

@@ -60,7 +60,7 @@ func (om OrdinalMapping) IsIdentityMapping() bool {
 	return true
 }
 
-var defaultTupleLengthTarget int64 = (1 << 11)
+var defaultTupleLengthTarget int64 = (1 << 11) // 2048 = 2KB
 
 type TupleBuilder struct {
 	vs                ValueStore
@@ -77,7 +77,7 @@ func NewTupleBuilder(desc *TupleDesc, vs ValueStore) *TupleBuilder {
 	return &TupleBuilder{
 		Desc:              desc,
 		fields:            make([][]byte, len(desc.Types)),
-		buf:               make([]byte, builderBufferSize),
+		buf:               make([]byte, defaultTupleLengthTarget),
 		vs:                vs,
 		tupleLengthTarget: defaultTupleLengthTarget,
 	}
@@ -522,11 +522,19 @@ func (tb *TupleBuilder) putAddr(i int, v hash.Hash) {
 }
 
 func (tb *TupleBuilder) ensureCapacity(sz ByteSize) {
-	need := int(tb.pos+int64(sz)) - len(tb.buf)
+	need := int(tb.pos+int64(sz)) - cap(tb.buf)
+	// TODO: This wastes memory when allocating a new backing array.
+	//   The initial part of the new backing array won't be referenced by anything.
+	//   tb.fields will still point to the original backing array.
+	//for i := 0; i < need; i++ {
+	//	tb.buf = append(tb.buf, byte(0))
+	//}
+
+	// We can safely replace tb.buf.
+	// tb.fields is still referencing the original backing array.
 	if need > 0 {
-		for i := 0; i < need; i++ {
-			tb.buf = append(tb.buf, byte(0))
-		}
+		tb.pos = 0
+		tb.buf = make([]byte, max(defaultTupleLengthTarget, int64(sz)))
 	}
 }
 

go/store/val/tuple_builder_test.go

@@ -17,6 +17,7 @@ package val
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"math"
 	"math/rand"
 	"testing"
@@ -354,3 +355,39 @@ func TestTupleBuilderAdaptiveEncodings(t *testing.T) {
 		})
 	}
 }
+
+type testBuilder struct {
+	fields [][]byte
+	buf    []byte
+}
+
+func TestTestBuilder(t *testing.T) {
+	tb := &testBuilder{
+		fields: make([][]byte, 20),
+		buf:    make([]byte, 10),
+	}
+	origbuf := tb.buf
+	for i := 0; i < len(tb.buf); i++ {
+		tb.buf[i] = byte(i)
+		tb.fields[i] = tb.buf[i : i+1]
+	}
+	fmt.Printf("Orig Buf: %v\n", origbuf)
+	fmt.Printf("  Fields: %v\n", tb.fields)
+
+	// Allocate new backing array
+	for i := 10; i < len(tb.fields); i++ {
+		tb.buf = append(tb.buf, byte(100+i))
+		tb.fields[i] = tb.buf[i : i+1]
+	}
+	newBuf := tb.buf
+	fmt.Printf("Orig Buf: %v\n", origbuf)
+	fmt.Printf(" New Buf: %v\n", newBuf)
+	fmt.Printf("  Fields: %v\n", tb.fields)
+
+	// Fields[0] still references the original backing array
+	origbuf[0] = 99
+	newBuf[0] = 10
+	fmt.Printf("Orig Buf: %v\n", origbuf)
+	fmt.Printf(" New Buf: %v\n", newBuf)
+	fmt.Printf("  Fields: %v\n", tb.fields)
+}

go/store/val/tuple_descriptor.go

@@ -585,13 +585,10 @@ func (td *TupleDesc) GetAddr(i int, tup Tuple) (hash.Hash, bool) {
 	return hash.New(b), true
 }
 
-func (td *TupleDesc) ExpectEncoding(i int, encodings ...Encoding) {
-	for _, enc := range encodings {
-		if enc == td.Types[i].Enc {
-			return
-		}
+func (td *TupleDesc) ExpectEncoding(i int, enc Encoding) {
+	if enc != td.Types[i].Enc {
+		panic("incorrect value encoding")
 	}
-	panic("incorrect value encoding")
 }
 
 func (td *TupleDesc) GetCell(i int, tup Tuple) (v Cell, ok bool) {