util: add PreAlloc4Row and Insert for Chunk and List

util/chunk/chunk.go

@@ -15,6 +15,7 @@ package chunk
 
 import (
 	"encoding/binary"
+	"reflect"
 	"unsafe"
 
 	"github.com/cznic/mathutil"
@@ -277,6 +278,60 @@ func (c *Chunk) AppendPartialRow(colIdx int, row Row) {
 	}
 }
 
+// PreAlloc4Row pre-allocates the memory space for a Row.
+// The null elem info will be pre-written.
+func (c *Chunk) PreAlloc4Row(row Row) {
+	for i, rowCol := range row.c.columns {
+		chkCol := c.columns[i]
+		chkCol.appendNullBitmap(!rowCol.isNull(row.idx))
+		if rowCol.isFixed() {
+			elemLen := len(rowCol.elemBuf)
+			if len(chkCol.data)+elemLen >= cap(chkCol.data) {
+				chkCol.data = make([]byte, len(chkCol.data)+elemLen, 2*cap(chkCol.data))
+			} else {
+				(*reflect.SliceHeader)(unsafe.Pointer(&chkCol.data)).Len = len(chkCol.data) + elemLen
+			}
+		} else {
+			elemLen := int(rowCol.offsets[row.idx+1] - rowCol.offsets[row.idx])
+			if len(chkCol.data)+elemLen >= cap(chkCol.data) {
+				chkCol.data = make([]byte, len(chkCol.data)+elemLen, 2*cap(chkCol.data))
+			} else {
+				(*reflect.SliceHeader)(unsafe.Pointer(&chkCol.data)).Len = len(chkCol.data) + elemLen
+			}
+			chkCol.offsets = append(chkCol.offsets, int32(len(chkCol.data)))
+		}
+		chkCol.length++
+	}
+	c.numVirtualRows++
+}
+
+// Insert inserts `row` on the position specified by `rowIdx`.
+// Note: Insert will cover the origin data, it should be called after
+// PreAlloc4Row.
+func (c *Chunk) Insert(rowIdx int, row Row) {
+	// Check data length between row and the origin data for every column.
+	// Cover the origin data if the upper check is valid.
+	for i, rowCol := range row.c.columns {
+		chkCol := c.columns[i]
+		if chkCol.isFixed() != rowCol.isFixed() {
+			panic("unexcepted error happens during Chunk.Insert")
+		}
+		var srcStart, srcEnd, destStart, destEnd int
+		if rowCol.isFixed() {
+			srcElemLen, destElemLen := len(rowCol.elemBuf), len(chkCol.elemBuf)
+			srcStart, destStart = row.idx*srcElemLen, rowIdx*destElemLen
+			srcEnd, destEnd = srcStart+srcElemLen, destStart+destElemLen
+		} else {
+			srcStart, srcEnd = int(rowCol.offsets[row.idx]), int(rowCol.offsets[row.idx+1])
+			destStart, destEnd = int(chkCol.offsets[rowIdx]), int(chkCol.offsets[rowIdx+1])
+		}
+		if destEnd-destStart != srcEnd-srcStart {
+			panic("unexcepted error happens during Chunk.Insert")
+		}
+		copy(chkCol.data[destStart:destEnd], rowCol.data[srcStart:srcEnd])
+	}
+}
+
 // Append appends rows in [begin, end) in another Chunk to a Chunk.
 func (c *Chunk) Append(other *Chunk, begin, end int) {
 	for colID, src := range other.columns {

util/chunk/chunk_test.go

@@ -18,6 +18,8 @@ import (
 	"fmt"
 	"math"
 	"strconv"
+	"strings"
+	"sync"
 	"testing"
 	"time"
 	"unsafe"
@@ -517,6 +519,107 @@ func (s *testChunkSuite) TestSwapColumn(c *check.C) {
 	checkRef()
 }
 
+func (s *testChunkSuite) TestPreAlloc4RowAndInsert(c *check.C) {
+	fieldTypes := make([]*types.FieldType, 0, 4)
+	fieldTypes = append(fieldTypes, &types.FieldType{Tp: mysql.TypeFloat})
+	fieldTypes = append(fieldTypes, &types.FieldType{Tp: mysql.TypeLonglong})
+	fieldTypes = append(fieldTypes, &types.FieldType{Tp: mysql.TypeNewDecimal})
+	fieldTypes = append(fieldTypes, &types.FieldType{Tp: mysql.TypeVarchar})
+
+	srcChk := NewChunkWithCapacity(fieldTypes, 10)
+	for i := int64(0); i < 10; i++ {
+		srcChk.AppendFloat32(0, float32(i))
+		srcChk.AppendInt64(1, i)
+		srcChk.AppendMyDecimal(2, types.NewDecFromInt(i))
+		srcChk.AppendString(3, strings.Repeat(strconv.FormatInt(i, 10), int(i)))
+	}
+
+	destChk := NewChunkWithCapacity(fieldTypes, 3)
+
+	// Test Chunk.PreAlloc4Row.
+	for i := 0; i < srcChk.NumRows(); i++ {
+		c.Assert(destChk.NumRows(), check.Equals, i)
+		destChk.PreAlloc4Row(srcChk.GetRow(i))
+	}
+	for i, srcCol := range srcChk.columns {
+		destCol := destChk.columns[i]
+		c.Assert(len(srcCol.elemBuf), check.Equals, len(destCol.elemBuf))
+		c.Assert(len(srcCol.data), check.Equals, len(destCol.data))
+		c.Assert(len(srcCol.offsets), check.Equals, len(destCol.offsets))
+		c.Assert(len(srcCol.nullBitmap), check.Equals, len(destCol.nullBitmap))
+		c.Assert(srcCol.length, check.Equals, destCol.length)
+		c.Assert(srcCol.nullCount, check.Equals, destCol.nullCount)
+
+		for _, val := range destCol.data {
+			c.Assert(val == 0, check.IsTrue)
+		}
+		for j, val := range srcCol.offsets {
+			c.Assert(val, check.Equals, destCol.offsets[j])
+		}
+		for j, val := range srcCol.nullBitmap {
+			c.Assert(val, check.Equals, destCol.nullBitmap[j])
+		}
+		for _, val := range destCol.elemBuf {
+			c.Assert(val == 0, check.IsTrue)
+		}
+	}
+
+	// Test Chunk.Insert.
+	for i := srcChk.NumRows() - 1; i >= 0; i-- {
+		destChk.Insert(i, srcChk.GetRow(i))
+	}
+	for i, srcCol := range srcChk.columns {
+		destCol := destChk.columns[i]
+
+		for j, val := range srcCol.data {
+			c.Assert(val, check.Equals, destCol.data[j])
+		}
+		for j, val := range srcCol.offsets {
+			c.Assert(val, check.Equals, destCol.offsets[j])
+		}
+		for j, val := range srcCol.nullBitmap {
+			c.Assert(val, check.Equals, destCol.nullBitmap[j])
+		}
+		for _, val := range destCol.elemBuf {
+			c.Assert(val == 0, check.IsTrue)
+		}
+	}
+
+	// Test parallel Chunk.Insert.
+	destChk.Reset()
+	startWg, endWg := &sync.WaitGroup{}, &sync.WaitGroup{}
+	startWg.Add(1)
+	for i := 0; i < srcChk.NumRows(); i++ {
+		destChk.PreAlloc4Row(srcChk.GetRow(i))
+		endWg.Add(1)
+		go func(rowIdx int) {
+			defer func() {
+				endWg.Done()
+			}()
+			startWg.Wait()
+			destChk.Insert(rowIdx, srcChk.GetRow(rowIdx))
+		}(i)
+	}
+	startWg.Done()
+	endWg.Wait()
+	for i, srcCol := range srcChk.columns {
+		destCol := destChk.columns[i]
+
+		for j, val := range srcCol.data {
+			c.Assert(val, check.Equals, destCol.data[j])
+		}
+		for j, val := range srcCol.offsets {
+			c.Assert(val, check.Equals, destCol.offsets[j])
+		}
+		for j, val := range srcCol.nullBitmap {
+			c.Assert(val, check.Equals, destCol.nullBitmap[j])
+		}
+		for _, val := range destCol.elemBuf {
+			c.Assert(val == 0, check.IsTrue)
+		}
+	}
+}
+
 func BenchmarkAppendInt(b *testing.B) {
 	b.ReportAllocs()
 	chk := newChunk(8)

util/chunk/list.go

@@ -140,6 +140,32 @@ func (l *List) Reset() {
 	l.consumedIdx = -1
 }
 
+// PreAlloc4Row pre-allocate the storage memory for a Row.
+func (l *List) PreAlloc4Row(row Row) (ptr RowPtr) {
+	chkIdx := len(l.chunks) - 1
+	if chkIdx == -1 || l.chunks[chkIdx].NumRows() >= l.chunks[chkIdx].Capacity() || chkIdx == l.consumedIdx {
+		newChk := l.allocChunk()
+		l.chunks = append(l.chunks, newChk)
+		if chkIdx != l.consumedIdx {
+			l.memTracker.Consume(l.chunks[chkIdx].MemoryUsage())
+			l.consumedIdx = chkIdx
+		}
+		chkIdx++
+	}
+	chk := l.chunks[chkIdx]
+	rowIdx := chk.NumRows()
+	chk.PreAlloc4Row(row)
+	l.length++
+	return RowPtr{ChkIdx: uint32(chkIdx), RowIdx: uint32(rowIdx)}
+}
+
+// Insert inserts `row` on the position specified by `ptr`.
+// Note: Insert will cover the origin data, it should be called after
+// PreAlloc4Row.
+func (l *List) Insert(ptr RowPtr, row Row) {
+	l.chunks[ptr.ChkIdx].Insert(int(ptr.RowIdx), row)
+}
+
 // ListWalkFunc is used to walk the list.
 // If error is returned, it will stop walking.
 type ListWalkFunc = func(row Row) error

util/chunk/list_test.go

@@ -22,6 +22,8 @@ import (
 	"github.com/pingcap/tidb/mysql"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/types/json"
+	"strconv"
+	"strings"
 )
 
 func (s *testChunkSuite) TestList(c *check.C) {
@@ -114,6 +116,47 @@ func (s *testChunkSuite) TestListMemoryUsage(c *check.C) {
 	c.Assert(list.GetMemTracker().BytesConsumed(), check.Equals, memUsage+srcChk.MemoryUsage())
 }
 
+func (s *testChunkSuite) TestListPrePreAlloc4RowAndInsert(c *check.C) {
+	fieldTypes := make([]*types.FieldType, 0, 4)
+	fieldTypes = append(fieldTypes, &types.FieldType{Tp: mysql.TypeFloat})
+	fieldTypes = append(fieldTypes, &types.FieldType{Tp: mysql.TypeLonglong})
+	fieldTypes = append(fieldTypes, &types.FieldType{Tp: mysql.TypeNewDecimal})
+	fieldTypes = append(fieldTypes, &types.FieldType{Tp: mysql.TypeVarchar})
+
+	srcChk := NewChunkWithCapacity(fieldTypes, 10)
+	for i := int64(0); i < 10; i++ {
+		srcChk.AppendFloat32(0, float32(i))
+		srcChk.AppendInt64(1, i)
+		srcChk.AppendMyDecimal(2, types.NewDecFromInt(i))
+		srcChk.AppendString(3, strings.Repeat(strconv.FormatInt(i, 10), int(i)))
+	}
+
+	srcList := NewList(fieldTypes, 3, 3)
+	destList := NewList(fieldTypes, 5, 5)
+	destRowPtr := make([]RowPtr, srcChk.NumRows())
+	for i := 0; i < srcChk.NumRows(); i++ {
+		srcList.AppendRow(srcChk.GetRow(i))
+		destRowPtr[i] = destList.PreAlloc4Row(srcChk.GetRow(i))
+	}
+
+	c.Assert(srcList.NumChunks(), check.Equals, 4)
+	c.Assert(destList.NumChunks(), check.Equals, 2)
+
+	iter4Src := NewIterator4List(srcList)
+	for row, i := iter4Src.Begin(), 0; row != iter4Src.End(); row, i = iter4Src.Next(), i+1 {
+		destList.Insert(destRowPtr[i], row)
+	}
+
+	iter4Dest := NewIterator4List(destList)
+	srcRow, destRow := iter4Src.Begin(), iter4Dest.Begin()
+	for ; srcRow != iter4Src.End(); srcRow, destRow = iter4Src.Next(), iter4Dest.Next() {
+		c.Assert(srcRow.GetFloat32(0), check.Equals, destRow.GetFloat32(0))
+		c.Assert(srcRow.GetInt64(1), check.Equals, destRow.GetInt64(1))
+		c.Assert(srcRow.GetMyDecimal(2).Compare(destRow.GetMyDecimal(2)) == 0, check.IsTrue)
+		c.Assert(srcRow.GetString(3), check.Equals, destRow.GetString(3))
+	}
+}
+
 func BenchmarkListMemoryUsage(b *testing.B) {
 	fieldTypes := make([]*types.FieldType, 0, 4)
 	fieldTypes = append(fieldTypes, &types.FieldType{Tp: mysql.TypeFloat})