util: add an interface for Chunk to count the memory usage (#5645)

util/chunk/chunk.go

@@ -57,6 +57,17 @@ func NewChunkWithCapacity(fields []*types.FieldType, cap int) *Chunk {
 	return chk
 }
 
+// MemoryUsage returns the total memory usage of a Chunk in B.
+// We ignore the size of column.length and column.nullCount
+// since they have little effect of the total memory usage.
+func (c *Chunk) MemoryUsage() (sum int64) {
+	for _, col := range c.columns {
+		curColMemUsage := int64(unsafe.Sizeof(*col)) + int64(cap(col.nullBitmap)) + int64(cap(col.offsets)*4) + int64(cap(col.data)) + int64(cap(col.elemBuf))
+		sum += curColMemUsage
+	}
+	return
+}
+
 // addFixedLenColumn adds a fixed length column with elemLen and initial data capacity.
 func (c *Chunk) addFixedLenColumn(elemLen, initCap int) {
 	c.columns = append(c.columns, &column{

util/chunk/chunk_test.go

@@ -18,6 +18,7 @@ import (
 	"math"
 	"testing"
 	"time"
+	"unsafe"
 
 	"github.com/pingcap/check"
 	"github.com/pingcap/tidb/mysql"
@@ -394,6 +395,61 @@ func (s *testChunkSuite) TestGetDecimalDatum(c *check.C) {
 	c.Assert(decDatum.Frac(), check.Equals, decFromChk.Frac())
 }
 
+func (s *testChunkSuite) TestChunkMemoryUsage(c *check.C) {
+	fieldTypes := make([]*types.FieldType, 0, 3)
+	fieldTypes = append(fieldTypes, &types.FieldType{Tp: mysql.TypeFloat})
+	fieldTypes = append(fieldTypes, &types.FieldType{Tp: mysql.TypeVarchar})
+	fieldTypes = append(fieldTypes, &types.FieldType{Tp: mysql.TypeJSON})
+	fieldTypes = append(fieldTypes, &types.FieldType{Tp: mysql.TypeDatetime})
+	fieldTypes = append(fieldTypes, &types.FieldType{Tp: mysql.TypeDuration})
+
+	initCap := 10
+	chk := NewChunkWithCapacity(fieldTypes, initCap)
+
+	//cap(c.nullBitmap) + cap(c.offsets)*4 + cap(c.data) + cap(c.elemBuf)
+	colUsage := make([]int, len(fieldTypes))
+	colUsage[0] = initCap>>3 + 0 + initCap*4 + 4
+	colUsage[1] = initCap>>3 + (initCap+1)*4 + initCap*4 + 0
+	colUsage[2] = initCap>>3 + (initCap+1)*4 + initCap*4 + 0
+	colUsage[3] = initCap>>3 + 0 + initCap*16 + 16
+	colUsage[4] = initCap>>3 + 0 + initCap*16 + 16
+
+	expectedUsage := 0
+	for i := range colUsage {
+		expectedUsage += colUsage[i] + int(unsafe.Sizeof(*chk.columns[i]))
+	}
+	memUsage := chk.MemoryUsage()
+	c.Assert(memUsage, check.Equals, int64(expectedUsage))
+
+	jsonObj, err := json.ParseBinaryFromString("1")
+	c.Assert(err, check.IsNil)
+	timeObj := types.Time{Time: types.FromGoTime(time.Now()), Fsp: 0, Type: mysql.TypeDatetime}
+	durationObj := types.Duration{Duration: math.MaxInt64, Fsp: 0}
+
+	chk.AppendFloat32(0, 12.4)
+	chk.AppendString(1, "123")
+	chk.AppendJSON(2, jsonObj)
+	chk.AppendTime(3, timeObj)
+	chk.AppendDuration(4, durationObj)
+
+	memUsage = chk.MemoryUsage()
+	c.Assert(memUsage, check.Equals, int64(expectedUsage))
+
+	chk.AppendFloat32(0, 12.4)
+	chk.AppendString(1, "123111111111111111111111111111111111111111111111")
+	chk.AppendJSON(2, jsonObj)
+	chk.AppendTime(3, timeObj)
+	chk.AppendDuration(4, durationObj)
+
+	memUsage = chk.MemoryUsage()
+	colUsage[1] = initCap>>3 + (initCap+1)*4 + cap(chk.columns[1].data) + 0
+	expectedUsage = 0
+	for i := range colUsage {
+		expectedUsage += colUsage[i] + int(unsafe.Sizeof(*chk.columns[i]))
+	}
+	c.Assert(memUsage, check.Equals, int64(expectedUsage))
+}
+
 func BenchmarkAppendInt(b *testing.B) {
 	b.ReportAllocs()
 	chk := newChunk(8)