diff --git a/util/chunk/chunk.go b/util/chunk/chunk.go index 6539f0ccc7..eecf4e5caf 100644 --- a/util/chunk/chunk.go +++ b/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{ diff --git a/util/chunk/chunk_test.go b/util/chunk/chunk_test.go index 55de67a96a..95bae758d4 100644 --- a/util/chunk/chunk_test.go +++ b/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)