executor: track the memory usage of aggregate more accurate (#34732)

ref #33877

executor/aggfuncs/aggfunc_test.go

@@ -165,13 +165,10 @@ func distinctUpdateMemDeltaGens(srcChk *chunk.Chunk, dataType *types.FieldType)
 		switch dataType.GetType() {
 		case mysql.TypeLonglong:
 			val = strconv.FormatInt(row.GetInt64(0), 10)
-			memDelta = aggfuncs.DefInt64Size
 		case mysql.TypeFloat:
 			val = strconv.FormatFloat(float64(row.GetFloat32(0)), 'f', 6, 64)
-			memDelta = aggfuncs.DefFloat64Size
 		case mysql.TypeDouble:
 			val = strconv.FormatFloat(row.GetFloat64(0), 'f', 6, 64)
-			memDelta = aggfuncs.DefFloat64Size
 		case mysql.TypeNewDecimal:
 			decimal := row.GetMyDecimal(0)
 			hash, err := decimal.ToHashKey()
@@ -190,7 +187,6 @@ func distinctUpdateMemDeltaGens(srcChk *chunk.Chunk, dataType *types.FieldType)
 			memDelta = 16
 		case mysql.TypeDuration:
 			val = strconv.FormatInt(row.GetInt64(0), 10)
-			memDelta = aggfuncs.DefInt64Size
 		case mysql.TypeJSON:
 			jsonVal := row.GetJSON(0)
 			bytes := make([]byte, 0)

executor/aggfuncs/func_avg.go

@@ -246,6 +246,7 @@ func (e *avgOriginal4DistinctDecimal) UpdatePartialResult(sctx sessionctx.Contex
 			continue
 		}
 		memDelta += p.valSet.Insert(decStr)
+		memDelta += int64(len(decStr))
 		newSum := new(types.MyDecimal)
 		err = types.DecimalAdd(&p.sum, input, newSum)
 		if err != nil {

executor/aggfuncs/func_avg_test.go

@@ -20,8 +20,8 @@ import (
 	"github.com/pingcap/tidb/executor/aggfuncs"
 	"github.com/pingcap/tidb/parser/ast"
 	"github.com/pingcap/tidb/parser/mysql"
+	"github.com/pingcap/tidb/util/hack"
 	"github.com/pingcap/tidb/util/mock"
-	"github.com/pingcap/tidb/util/set"
 )
 
 func TestMergePartialResult4Avg(t *testing.T) {
@@ -50,11 +50,11 @@ func TestMemAvg(t *testing.T) {
 		buildAggMemTester(ast.AggFuncAvg, mysql.TypeNewDecimal, 5,
 			aggfuncs.DefPartialResult4AvgDecimalSize, defaultUpdateMemDeltaGens, false),
 		buildAggMemTester(ast.AggFuncAvg, mysql.TypeNewDecimal, 5,
-			aggfuncs.DefPartialResult4AvgDistinctDecimalSize+set.DefStringSetBucketMemoryUsage, distinctUpdateMemDeltaGens, true),
+			aggfuncs.DefPartialResult4AvgDistinctDecimalSize+hack.DefBucketMemoryUsageForSetString, distinctUpdateMemDeltaGens, true),
 		buildAggMemTester(ast.AggFuncAvg, mysql.TypeDouble, 5,
 			aggfuncs.DefPartialResult4AvgFloat64Size, defaultUpdateMemDeltaGens, false),
 		buildAggMemTester(ast.AggFuncAvg, mysql.TypeDouble, 5,
-			aggfuncs.DefPartialResult4AvgDistinctFloat64Size+set.DefFloat64SetBucketMemoryUsage, distinctUpdateMemDeltaGens, true),
+			aggfuncs.DefPartialResult4AvgDistinctFloat64Size+hack.DefBucketMemoryUsageForSetFloat64, distinctUpdateMemDeltaGens, true),
 	}
 	for _, test := range tests {
 		testAggMemFunc(t, test)

executor/aggfuncs/func_count_distinct.go

@@ -189,6 +189,7 @@ func (e *countOriginalWithDistinct4Decimal) UpdatePartialResult(sctx sessionctx.
 			continue
 		}
 		memDelta += p.valSet.Insert(decStr)
+		memDelta += int64(len(decStr))
 	}
 
 	return memDelta, nil
@@ -286,6 +287,7 @@ func (e *countOriginalWithDistinct4String) UpdatePartialResult(sctx sessionctx.C
 		}
 		input = stringutil.Copy(input)
 		memDelta += p.valSet.Insert(input)
+		memDelta += int64(len(input))
 	}
 
 	return memDelta, nil
@@ -348,6 +350,7 @@ func (e *countOriginalWithDistinct) UpdatePartialResult(sctx sessionctx.Context,
 			continue
 		}
 		memDelta += p.valSet.Insert(encodedString)
+		memDelta += int64(len(encodedString))
 	}
 
 	return memDelta, nil

executor/aggfuncs/func_count_test.go

@@ -25,8 +25,8 @@ import (
 	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/sessionctx/stmtctx"
 	"github.com/pingcap/tidb/types"
+	"github.com/pingcap/tidb/util/hack"
 	"github.com/pingcap/tidb/util/mock"
-	"github.com/pingcap/tidb/util/set"
 	"github.com/stretchr/testify/require"
 )
 
@@ -130,21 +130,21 @@ func TestMemCount(t *testing.T) {
 		buildAggMemTester(ast.AggFuncCount, mysql.TypeDuration, 5,
 			aggfuncs.DefPartialResult4CountSize, defaultUpdateMemDeltaGens, false),
 		buildAggMemTester(ast.AggFuncCount, mysql.TypeLonglong, 5,
-			aggfuncs.DefPartialResult4CountDistinctIntSize+set.DefInt64SetBucketMemoryUsage, distinctUpdateMemDeltaGens, true),
+			aggfuncs.DefPartialResult4CountDistinctIntSize+hack.DefBucketMemoryUsageForSetInt64, distinctUpdateMemDeltaGens, true),
 		buildAggMemTester(ast.AggFuncCount, mysql.TypeFloat, 5,
-			aggfuncs.DefPartialResult4CountDistinctRealSize+set.DefFloat64SetBucketMemoryUsage, distinctUpdateMemDeltaGens, true),
+			aggfuncs.DefPartialResult4CountDistinctRealSize+hack.DefBucketMemoryUsageForSetFloat64, distinctUpdateMemDeltaGens, true),
 		buildAggMemTester(ast.AggFuncCount, mysql.TypeDouble, 5,
-			aggfuncs.DefPartialResult4CountDistinctRealSize+set.DefFloat64SetBucketMemoryUsage, distinctUpdateMemDeltaGens, true),
+			aggfuncs.DefPartialResult4CountDistinctRealSize+hack.DefBucketMemoryUsageForSetFloat64, distinctUpdateMemDeltaGens, true),
 		buildAggMemTester(ast.AggFuncCount, mysql.TypeNewDecimal, 5,
-			aggfuncs.DefPartialResult4CountDistinctDecimalSize+set.DefStringSetBucketMemoryUsage, distinctUpdateMemDeltaGens, true),
+			aggfuncs.DefPartialResult4CountDistinctDecimalSize+hack.DefBucketMemoryUsageForSetString, distinctUpdateMemDeltaGens, true),
 		buildAggMemTester(ast.AggFuncCount, mysql.TypeString, 5,
-			aggfuncs.DefPartialResult4CountDistinctStringSize+set.DefStringSetBucketMemoryUsage, distinctUpdateMemDeltaGens, true),
+			aggfuncs.DefPartialResult4CountDistinctStringSize+hack.DefBucketMemoryUsageForSetString, distinctUpdateMemDeltaGens, true),
 		buildAggMemTester(ast.AggFuncCount, mysql.TypeDate, 5,
-			aggfuncs.DefPartialResult4CountWithDistinctSize+set.DefStringSetBucketMemoryUsage, distinctUpdateMemDeltaGens, true),
+			aggfuncs.DefPartialResult4CountWithDistinctSize+hack.DefBucketMemoryUsageForSetString, distinctUpdateMemDeltaGens, true),
 		buildAggMemTester(ast.AggFuncCount, mysql.TypeDuration, 5,
-			aggfuncs.DefPartialResult4CountDistinctDurationSize+set.DefInt64SetBucketMemoryUsage, distinctUpdateMemDeltaGens, true),
+			aggfuncs.DefPartialResult4CountDistinctDurationSize+hack.DefBucketMemoryUsageForSetInt64, distinctUpdateMemDeltaGens, true),
 		buildAggMemTester(ast.AggFuncCount, mysql.TypeJSON, 5,
-			aggfuncs.DefPartialResult4CountWithDistinctSize+set.DefStringSetBucketMemoryUsage, distinctUpdateMemDeltaGens, true),
+			aggfuncs.DefPartialResult4CountWithDistinctSize+hack.DefBucketMemoryUsageForSetString, distinctUpdateMemDeltaGens, true),
 		buildAggMemTester(ast.AggFuncApproxCountDistinct, mysql.TypeLonglong, 5,
 			aggfuncs.DefPartialResult4ApproxCountDistinctSize, approxCountDistinctUpdateMemDeltaGens, true),
 		buildAggMemTester(ast.AggFuncApproxCountDistinct, mysql.TypeString, 5,

executor/aggfuncs/func_group_concat.go

@@ -248,6 +248,7 @@ func (e *groupConcatDistinct) UpdatePartialResult(sctx sessionctx.Context, rowsI
 			continue
 		}
 		memDelta += p.valSet.Insert(joinedVal)
+		memDelta += int64(len(joinedVal))
 		// write separator
 		if p.buffer == nil {
 			p.buffer = &bytes.Buffer{}
@@ -577,6 +578,7 @@ func (e *groupConcatDistinctOrder) UpdatePartialResult(sctx sessionctx.Context,
 			continue
 		}
 		memDelta += p.valSet.Insert(joinedVal)
+		memDelta += int64(len(joinedVal))
 		sortRow := sortRow{
 			buffer:  buffer,
 			byItems: make([]*types.Datum, 0, len(e.byItems)),

executor/aggfuncs/func_group_concat_test.go

@@ -66,13 +66,13 @@ func TestMemGroupConcat(t *testing.T) {
 	multiArgsTest1 := buildMultiArgsAggMemTester(ast.AggFuncGroupConcat, []byte{mysql.TypeString, mysql.TypeString}, mysql.TypeString, 5,
 		aggfuncs.DefPartialResult4GroupConcatSize+aggfuncs.DefBytesBufferSize, groupConcatMultiArgsUpdateMemDeltaGens, false)
 	multiArgsTest2 := buildMultiArgsAggMemTester(ast.AggFuncGroupConcat, []byte{mysql.TypeString, mysql.TypeString}, mysql.TypeString, 5,
-		aggfuncs.DefPartialResult4GroupConcatDistinctSize+aggfuncs.DefBytesBufferSize+set.DefStringSetBucketMemoryUsage, groupConcatDistinctMultiArgsUpdateMemDeltaGens, true)
+		aggfuncs.DefPartialResult4GroupConcatDistinctSize+aggfuncs.DefBytesBufferSize+hack.DefBucketMemoryUsageForSetString, groupConcatDistinctMultiArgsUpdateMemDeltaGens, true)
 
 	multiArgsTest3 := buildMultiArgsAggMemTester(ast.AggFuncGroupConcat, []byte{mysql.TypeString, mysql.TypeString}, mysql.TypeString, 5,
 		aggfuncs.DefPartialResult4GroupConcatOrderSize+aggfuncs.DefTopNRowsSize, groupConcatOrderMultiArgsUpdateMemDeltaGens, false)
 	multiArgsTest3.multiArgsAggTest.orderBy = true
 	multiArgsTest4 := buildMultiArgsAggMemTester(ast.AggFuncGroupConcat, []byte{mysql.TypeString, mysql.TypeString}, mysql.TypeString, 5,
-		aggfuncs.DefPartialResult4GroupConcatOrderDistinctSize+aggfuncs.DefTopNRowsSize+set.DefStringSetBucketMemoryUsage, groupConcatDistinctOrderMultiArgsUpdateMemDeltaGens, true)
+		aggfuncs.DefPartialResult4GroupConcatOrderDistinctSize+aggfuncs.DefTopNRowsSize+hack.DefBucketMemoryUsageForSetString, groupConcatDistinctOrderMultiArgsUpdateMemDeltaGens, true)
 	multiArgsTest4.multiArgsAggTest.orderBy = true
 
 	multiArgsTests := []multiArgsAggMemTest{multiArgsTest1, multiArgsTest2, multiArgsTest3, multiArgsTest4}

executor/aggfuncs/func_json_objectagg.go

@@ -29,8 +29,6 @@ import (
 const (
 	// DefPartialResult4JsonObjectAgg is the size of partialResult4JsonObject
 	DefPartialResult4JsonObjectAgg = int64(unsafe.Sizeof(partialResult4JsonObjectAgg{}))
-	// DefMapStringInterfaceBucketSize = bucketSize*(1+unsafe.Sizeof(string) + unsafe.Sizeof(interface{}))+2*ptrSize
-	DefMapStringInterfaceBucketSize = 8*(1+16+16) + 16
 )
 
 type jsonObjectAgg struct {
@@ -46,7 +44,7 @@ func (e *jsonObjectAgg) AllocPartialResult() (pr PartialResult, memDelta int64)
 	p := partialResult4JsonObjectAgg{}
 	p.entries = make(map[string]interface{})
 	p.bInMap = 0
-	return PartialResult(&p), DefPartialResult4JsonObjectAgg + (1<<p.bInMap)*DefMapStringInterfaceBucketSize
+	return PartialResult(&p), DefPartialResult4JsonObjectAgg + (1<<p.bInMap)*hack.DefBucketMemoryUsageForMapStringToAny
 }
 
 func (e *jsonObjectAgg) ResetPartialResult(pr PartialResult) {
@@ -114,7 +112,7 @@ func (e *jsonObjectAgg) UpdatePartialResult(sctx sessionctx.Context, rowsInGroup
 			if _, ok := p.entries[keyString]; !ok {
 				memDelta += int64(len(keyString)) + getValMemDelta(realVal)
 				if len(p.entries)+1 > (1<<p.bInMap)*hack.LoadFactorNum/hack.LoadFactorDen {
-					memDelta += (1 << p.bInMap) * DefMapStringInterfaceBucketSize
+					memDelta += (1 << p.bInMap) * hack.DefBucketMemoryUsageForMapStringToAny
 					p.bInMap++
 				}
 			}
@@ -162,7 +160,7 @@ func (e *jsonObjectAgg) MergePartialResult(sctx sessionctx.Context, src, dst Par
 		p2.entries[k] = v
 		memDelta += int64(len(k)) + getValMemDelta(v)
 		if len(p2.entries)+1 > (1<<p2.bInMap)*hack.LoadFactorNum/hack.LoadFactorDen {
-			memDelta += (1 << p2.bInMap) * DefMapStringInterfaceBucketSize
+			memDelta += (1 << p2.bInMap) * hack.DefBucketMemoryUsageForMapStringToAny
 			p2.bInMap++
 		}
 	}

executor/aggfuncs/func_json_objectagg_test.go

@@ -22,6 +22,7 @@ import (
 	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/types/json"
+	"github.com/pingcap/tidb/util/hack"
 	"github.com/pingcap/tidb/util/mock"
 )
 
@@ -146,8 +147,8 @@ func TestMemJsonObjectagg(t *testing.T) {
 		}
 
 		tests := []multiArgsAggMemTest{
-			buildMultiArgsAggMemTester(ast.AggFuncJsonObjectAgg, argTypes, mysql.TypeJSON, numRows, aggfuncs.DefPartialResult4JsonObjectAgg+aggfuncs.DefMapStringInterfaceBucketSize, defaultMultiArgsMemDeltaGens, true),
-			buildMultiArgsAggMemTester(ast.AggFuncJsonObjectAgg, argTypes, mysql.TypeJSON, numRows, aggfuncs.DefPartialResult4JsonObjectAgg+aggfuncs.DefMapStringInterfaceBucketSize, defaultMultiArgsMemDeltaGens, false),
+			buildMultiArgsAggMemTester(ast.AggFuncJsonObjectAgg, argTypes, mysql.TypeJSON, numRows, aggfuncs.DefPartialResult4JsonObjectAgg+hack.DefBucketMemoryUsageForMapStringToAny, defaultMultiArgsMemDeltaGens, true),
+			buildMultiArgsAggMemTester(ast.AggFuncJsonObjectAgg, argTypes, mysql.TypeJSON, numRows, aggfuncs.DefPartialResult4JsonObjectAgg+hack.DefBucketMemoryUsageForMapStringToAny, defaultMultiArgsMemDeltaGens, false),
 		}
 		for _, test := range tests {
 			testMultiArgsAggMemFunc(t, test)

executor/aggfuncs/func_sum.go

@@ -358,6 +358,7 @@ func (e *sum4DistinctDecimal) UpdatePartialResult(sctx sessionctx.Context, rowsI
 			continue
 		}
 		memDelta += p.valSet.Insert(decStr)
+		memDelta += int64(len(decStr))
 		if p.isNull {
 			p.val = *input
 			p.isNull = false

executor/aggfuncs/func_sum_test.go

@@ -22,7 +22,7 @@ import (
 	"github.com/pingcap/tidb/parser/ast"
 	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/types"
-	"github.com/pingcap/tidb/util/set"
+	"github.com/pingcap/tidb/util/hack"
 )
 
 func TestMergePartialResult4Sum(t *testing.T) {
@@ -57,9 +57,9 @@ func TestMemSum(t *testing.T) {
 		buildAggMemTester(ast.AggFuncSum, mysql.TypeNewDecimal, 5,
 			aggfuncs.DefPartialResult4SumDecimalSize, defaultUpdateMemDeltaGens, false),
 		buildAggMemTester(ast.AggFuncSum, mysql.TypeDouble, 5,
-			aggfuncs.DefPartialResult4SumDistinctFloat64Size+set.DefFloat64SetBucketMemoryUsage, distinctUpdateMemDeltaGens, true),
+			aggfuncs.DefPartialResult4SumDistinctFloat64Size+hack.DefBucketMemoryUsageForSetFloat64, distinctUpdateMemDeltaGens, true),
 		buildAggMemTester(ast.AggFuncSum, mysql.TypeNewDecimal, 5,
-			aggfuncs.DefPartialResult4SumDistinctDecimalSize+set.DefStringSetBucketMemoryUsage, distinctUpdateMemDeltaGens, true),
+			aggfuncs.DefPartialResult4SumDistinctDecimalSize+hack.DefBucketMemoryUsageForSetString, distinctUpdateMemDeltaGens, true),
 	}
 
 	for i, test := range tests {

executor/aggfuncs/func_varpop_test.go

@@ -22,7 +22,7 @@ import (
 	"github.com/pingcap/tidb/parser/ast"
 	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/types"
-	"github.com/pingcap/tidb/util/set"
+	"github.com/pingcap/tidb/util/hack"
 )
 
 func TestMergePartialResult4Varpop(t *testing.T) {
@@ -48,7 +48,7 @@ func TestMemVarpop(t *testing.T) {
 		buildAggMemTester(ast.AggFuncVarPop, mysql.TypeDouble, 5,
 			aggfuncs.DefPartialResult4VarPopFloat64Size, defaultUpdateMemDeltaGens, false),
 		buildAggMemTester(ast.AggFuncVarPop, mysql.TypeDouble, 5,
-			aggfuncs.DefPartialResult4VarPopDistinctFloat64Size+set.DefFloat64SetBucketMemoryUsage, distinctUpdateMemDeltaGens, true),
+			aggfuncs.DefPartialResult4VarPopDistinctFloat64Size+hack.DefBucketMemoryUsageForSetFloat64, distinctUpdateMemDeltaGens, true),
 	}
 	for n, test := range tests {
 		test := test

executor/aggregate.go

@@ -498,7 +498,7 @@ func getGroupKeyMemUsage(groupKey [][]byte) int64 {
 	for _, key := range groupKey {
 		mem += int64(cap(key))
 	}
-	mem += 12 * int64(cap(groupKey))
+	mem += aggfuncs.DefSliceSize * int64(cap(groupKey))
 	return mem
 }
 
@@ -602,11 +602,11 @@ func (w *baseHashAggWorker) getPartialResult(sc *stmtctx.StatementContext, group
 		for _, af := range w.aggFuncs {
 			partialResult, memDelta := af.AllocPartialResult()
 			partialResults[i] = append(partialResults[i], partialResult)
-			allMemDelta += memDelta
+			allMemDelta += memDelta + 8 // the memory usage of PartialResult
 		}
 		mapper[string(groupKey[i])] = partialResults[i]
 		allMemDelta += int64(len(groupKey[i]))
-		// Map will expand when count > bucketNum * loadFactor. The memory usage will doubled.
+		// Map will expand when count > bucketNum * loadFactor. The memory usage will double.
 		if len(mapper) > (1<<w.BInMap)*hack.LoadFactorNum/hack.LoadFactorDen {
 			w.memTracker.Consume(hack.DefBucketMemoryUsageForMapStrToSlice * (1 << w.BInMap))
 			w.BInMap++

util/hack/hack.go

@@ -61,7 +61,18 @@ const (
 	// DefBucketMemoryUsageForMapStrToSlice = bucketSize*(1+unsafe.Sizeof(string) + unsafe.Sizeof(slice))+2*ptrSize
 	// ref https://github.com/golang/go/blob/go1.15.6/src/reflect/type.go#L2162.
 	// The bucket size may be changed by golang implement in the future.
-	DefBucketMemoryUsageForMapStrToSlice = 8*(1+16+24) + 16
+	// Golang Map needs to acquire double the memory when expanding, and the old buckets will be released after the data is migrated.
+	// Considering the worst case, the data in the old bucket cannot be migrated in time, and the old bucket cannot
+	// be GCed, we expand the bucket size to 1.5 times to estimate the memory usage of Golang Map.
+	DefBucketMemoryUsageForMapStrToSlice = (8*(1+16+24) + 16) / 2 * 3
 	// DefBucketMemoryUsageForMapIntToPtr = bucketSize*(1+unsafe.Sizeof(uint64) + unsafe.Sizeof(pointer))+2*ptrSize
-	DefBucketMemoryUsageForMapIntToPtr = 8*(1+8+8) + 16
+	DefBucketMemoryUsageForMapIntToPtr = (8*(1+8+8) + 16) / 2 * 3
+	// DefBucketMemoryUsageForMapStringToAny = bucketSize*(1+unsafe.Sizeof(string) + unsafe.Sizeof(interface{}))+2*ptrSize
+	DefBucketMemoryUsageForMapStringToAny = (8*(1+16+16) + 16) / 2 * 3
+	// DefBucketMemoryUsageForSetString = bucketSize*(1+unsafe.Sizeof(string) + unsafe.Sizeof(struct{}))+2*ptrSize
+	DefBucketMemoryUsageForSetString = (8*(1+16+0) + 16) / 2 * 3
+	// DefBucketMemoryUsageForSetFloat64 = bucketSize*(1+unsafe.Sizeof(float64) + unsafe.Sizeof(struct{}))+2*ptrSize
+	DefBucketMemoryUsageForSetFloat64 = (8*(1+8+0) + 16) / 2 * 3
+	// DefBucketMemoryUsageForSetInt64 = bucketSize*(1+unsafe.Sizeof(int64) + unsafe.Sizeof(struct{}))+2*ptrSize
+	DefBucketMemoryUsageForSetInt64 = (8*(1+8+0) + 16) / 2 * 3
 )