sql row buffer for merge iters

Author: James Cor

sql/expression/function/aggregation/window_partition.go

@@ -170,7 +170,7 @@ func (i *WindowPartitionIter) materializeInput(ctx *sql.Context) (sql.WindowBuff
 			}
 			return nil, nil, err
 		}
-		input = append(input, append(row, j))
+		input = append(input, append(append(sql.Row(nil), row...), j))
 		j++
 	}
 

sql/rowexec/merge_join.go

@@ -46,8 +46,11 @@ func newMergeJoinIter(ctx *sql.Context, b sql.NodeExecBuilder, j *plan.JoinNode,
 		return nil, err
 	}
 
-	fullRow := make(sql.Row, len(row)+len(j.Left().Schema())+len(j.Right().Schema()))
-	fullRow[0] = row
+	//fullRow := make(sql.Row, len(row)+len(j.Left().Schema())+len(j.Right().Schema()))
+
+	rowBuf := sql.RowBufPool.Get().(*sql.RowBuffer)
+	fullRow := rowBuf.Get(len(row) + len(j.Left().Schema()) + len(j.Right().Schema()))
+	//fullRow[0] = row
 	if len(row) > 0 {
 		copy(fullRow[0:], row[:])
 	}
@@ -83,6 +86,7 @@ func newMergeJoinIter(ctx *sql.Context, b sql.NodeExecBuilder, j *plan.JoinNode,
 		leftRowLen:  len(j.Left().Schema()),
 		rightRowLen: len(j.Right().Schema()),
 		isReversed:  j.IsReversed,
+		rowBuffer:   rowBuf,
 	}
 	return iter, nil
 }
@@ -128,6 +132,8 @@ type mergeJoinIter struct {
 	leftRowLen  int
 	rightRowLen int
 	parentLen   int
+
+	rowBuffer *sql.RowBuffer
 }
 
 var _ sql.RowIter = (*mergeJoinIter)(nil)
@@ -314,7 +320,7 @@ func (i *mergeJoinIter) Next(ctx *sql.Context) (sql.Row, error) {
 }
 
 func (i *mergeJoinIter) copyReturnRow() sql.Row {
-	ret := make(sql.Row, len(i.fullRow))
+	ret := i.rowBuffer.Get(len(i.fullRow))
 	copy(ret, i.fullRow)
 	return ret
 }

sql/rows.go

@@ -18,6 +18,7 @@ import (
 	"fmt"
 	"io"
 	"strings"
+	"sync"
 
 	"github.com/dolthub/vitess/go/vt/proto/query"
 
@@ -83,6 +84,42 @@ func FormatRow(row Row) string {
 	return sb.String()
 }
 
+const defaultRowBuffCap = 4096
+
+type RowBuffer struct {
+	i   int
+	buf Row
+}
+
+func NewRowBuffer() *RowBuffer {
+	return &RowBuffer{
+		buf: make(Row, defaultRowBuffCap),
+	}
+}
+
+func (b *RowBuffer) Get(n int) (res Row) {
+	newI := b.i + n
+	if newI >= len(b.buf) {
+		//b.buf = append(b.buf, b.buf) // TODO: not sure if this is correct, but it seems faster
+		buf := make(Row, len(b.buf)*2)
+		copy(b.buf, buf)
+		b.buf = buf
+	}
+	res = b.buf[b.i:newI]
+	b.i = newI
+	return
+}
+
+func (b *RowBuffer) Reset() {
+	b.i = 0
+}
+
+var RowBufPool = sync.Pool{
+	New: func() any {
+		return NewRowBuffer()
+	},
+}
+
 // RowIter is an iterator that produces rows.
 // TODO: most row iters need to be Disposable for CachedResult safety
 type RowIter interface {