*: support serialization and deserialization of aggregate function for spill

executor/aggfuncs/BUILD.bazel

@@ -27,6 +27,8 @@ go_library(
         "func_varpop.go",
         "func_varsamp.go",
         "row_number.go",
+        "spill_deserialize_helper.go",
+        "spill_serialize_helper.go",
     ],
     importpath = "github.com/pingcap/tidb/executor/aggfuncs",
     visibility = ["//visibility:public"],
@@ -49,6 +51,7 @@ go_library(
         "//util/mathutil",
         "//util/selection",
         "//util/set",
+        "//util/spill",
         "//util/stringutil",
         "@com_github_dgryski_go_farm//:go-farm",
         "@com_github_pingcap_errors//:errors",

executor/aggfuncs/aggfuncs.go

@@ -139,8 +139,21 @@ const (
 // to be any type.
 type PartialResult unsafe.Pointer
 
+// AggPartialResultMapper contains aggregate function results
+type AggPartialResultMapper map[string][]PartialResult
+
+type aggSpillHelper interface {
+	// SerializeForSpill will serialize meta data of aggregate function into bytes and put them into chunk.
+	SerializeForSpill(ctx sessionctx.Context, partialResult PartialResult, chk *chunk.Chunk, spillHelper *SpillSerializeHelper)
+
+	// DeserializeToPartialResultForSpill deserializes from bytes to PartialResult.
+	DeserializeToPartialResultForSpill(ctx sessionctx.Context, src *chunk.Chunk) ([]PartialResult, int64)
+}
+
 // AggFunc is the interface to evaluate the aggregate functions.
 type AggFunc interface {
+	aggSpillHelper
+
 	// AllocPartialResult allocates a specific data structure to store the
 	// partial result, initializes it, and converts it to PartialResult to
 	// return back. The second returned value is the memDelta used to trace
@@ -196,6 +209,8 @@ func (*baseAggFunc) MergePartialResult(sessionctx.Context, PartialResult, Partia
 	return 0, nil
 }
 
+
+
 // SlidingWindowAggFunc is the interface to evaluate the aggregate functions using sliding window.
 type SlidingWindowAggFunc interface {
 	// Slide evaluates the aggregate functions using a sliding window. The input

executor/aggfuncs/func_avg.go

@@ -47,6 +47,40 @@ type baseAvgDecimal struct {
 	baseAggFunc
 }
 
+func (c *baseAvgDecimal) SerializeForSpill(_ sessionctx.Context, partialResult PartialResult, chk *chunk.Chunk, spillHelper *SpillSerializeHelper) {
+	pr := (*partialResult4AvgDecimal)(partialResult)
+	resBuf := spillHelper.serializePartialResult4AvgDecimal(*pr)
+	chk.AppendBytes(c.ordinal, resBuf)
+}
+
+func (c *baseAvgDecimal) DeserializeToPartialResultForSpill(_ sessionctx.Context, src *chunk.Chunk) ([]PartialResult, int64) {
+	dataCol := src.Column(c.ordinal)
+	totalMemDelta := int64(0)
+	spillHelper := newDeserializeHelper(dataCol, src.NumRows())
+	partialResults := make([]PartialResult, 0, src.NumRows())
+
+	for {
+		pr, memDelta := c.deserializeForSpill(&spillHelper)
+		if pr == nil {
+			break
+		}
+		partialResults = append(partialResults, pr)
+		totalMemDelta += memDelta
+	}
+
+	return partialResults, totalMemDelta
+}
+
+func (c *baseAvgDecimal) deserializeForSpill(helper *spillDeserializeHelper) (PartialResult, int64) {
+	pr, memDelta := c.AllocPartialResult()
+	result := (*partialResult4AvgDecimal)(pr)
+	success := helper.deserializePartialResult4AvgDecimal(result)
+	if !success {
+		return nil, 0
+	}
+	return pr, memDelta
+}
+
 type partialResult4AvgDecimal struct {
 	sum   types.MyDecimal
 	count int64
@@ -285,6 +319,15 @@ func (e *avgOriginal4DistinctDecimal) AppendFinalResult2Chunk(_ sessionctx.Conte
 	return nil
 }
 
+// TODO implement it
+func (c *avgOriginal4DistinctDecimal) SerializeForSpill(_ sessionctx.Context, partialResult PartialResult, chk *chunk.Chunk, spillHelper *SpillSerializeHelper) {
+}
+
+// TODO implement it
+func (c *avgOriginal4DistinctDecimal) DeserializeToPartialResultForSpill(_ sessionctx.Context, src *chunk.Chunk) ([]PartialResult, int64) {
+	return nil, 0
+}
+
 // All the following avg function implementations return the float64 result,
 // which store the partial results in "partialResult4AvgFloat64".
 //
@@ -320,6 +363,40 @@ func (e *baseAvgFloat64) AppendFinalResult2Chunk(_ sessionctx.Context, pr Partia
 	return nil
 }
 
+func (c *baseAvgFloat64) SerializeForSpill(_ sessionctx.Context, partialResult PartialResult, chk *chunk.Chunk, spillHelper *SpillSerializeHelper) {
+	pr := (*partialResult4AvgFloat64)(partialResult)
+	resBuf := spillHelper.serializePartialResult4AvgFloat64(*pr)
+	chk.AppendBytes(c.ordinal, resBuf)
+}
+
+func (c *baseAvgFloat64) DeserializeToPartialResultForSpill(_ sessionctx.Context, src *chunk.Chunk) ([]PartialResult, int64) {
+	dataCol := src.Column(c.ordinal)
+	totalMemDelta := int64(0)
+	spillHelper := newDeserializeHelper(dataCol, src.NumRows())
+	partialResults := make([]PartialResult, 0, src.NumRows())
+
+	for {
+		pr, memDelta := c.deserializeForSpill(&spillHelper)
+		if pr == nil {
+			break
+		}
+		partialResults = append(partialResults, pr)
+		totalMemDelta += memDelta
+	}
+
+	return partialResults, totalMemDelta
+}
+
+func (c *baseAvgFloat64) deserializeForSpill(helper *spillDeserializeHelper) (PartialResult, int64) {
+	pr, memDelta := c.AllocPartialResult()
+	result := (*partialResult4AvgFloat64)(pr)
+	success := helper.deserializePartialResult4AvgFloat64(result)
+	if !success {
+		return nil, 0
+	}
+	return pr, memDelta
+}
+
 type avgOriginal4Float64HighPrecision struct {
 	baseAvgFloat64
 }
@@ -460,3 +537,12 @@ func (e *avgOriginal4DistinctFloat64) AppendFinalResult2Chunk(_ sessionctx.Conte
 	chk.AppendFloat64(e.ordinal, p.sum/float64(p.count))
 	return nil
 }
+
+// TODO implement it
+func (c *avgOriginal4DistinctFloat64) SerializeForSpill(_ sessionctx.Context, partialResult PartialResult, chk *chunk.Chunk, spillHelper *SpillSerializeHelper) {
+}
+
+// TODO implement it
+func (c *avgOriginal4DistinctFloat64) DeserializeToPartialResultForSpill(_ sessionctx.Context, src *chunk.Chunk) ([]PartialResult, int64) {
+	return nil, 0
+}

executor/aggfuncs/func_bitfuncs.go

@@ -48,6 +48,40 @@ func (e *baseBitAggFunc) AppendFinalResult2Chunk(_ sessionctx.Context, pr Partia
 	return nil
 }
 
+func (c *baseBitAggFunc) SerializeForSpill(_ sessionctx.Context, partialResult PartialResult, chk *chunk.Chunk, spillHelper *SpillSerializeHelper) {
+	pr := (*partialResult4BitFunc)(partialResult)
+	resBuf := spillHelper.serializePartialResult4BitFunc(*pr)
+	chk.AppendBytes(c.ordinal, resBuf)
+}
+
+func (c *baseBitAggFunc) DeserializeToPartialResultForSpill(_ sessionctx.Context, src *chunk.Chunk) ([]PartialResult, int64) {
+	dataCol := src.Column(c.ordinal)
+	totalMemDelta := int64(0)
+	spillHelper := newDeserializeHelper(dataCol, src.NumRows())
+	partialResults := make([]PartialResult, 0, src.NumRows())
+
+	for {
+		pr, memDelta := c.deserializeForSpill(&spillHelper)
+		if pr == nil {
+			break
+		}
+		partialResults = append(partialResults, pr)
+		totalMemDelta += memDelta
+	}
+
+	return partialResults, totalMemDelta
+}
+
+func (c *baseBitAggFunc) deserializeForSpill(helper *spillDeserializeHelper) (PartialResult, int64) {
+	pr, memDelta := c.AllocPartialResult()
+	result := (*partialResult4BitFunc)(pr)
+	success := helper.deserializePartialResult4BitFunc(result)
+	if !success {
+		return nil, 0
+	}
+	return pr, memDelta
+}
+
 type bitOrUint64 struct {
 	baseBitAggFunc
 }

executor/aggfuncs/func_count.go

@@ -32,6 +32,8 @@ type baseCount struct {
 
 type partialResult4Count = int64
 
+const partialResult4CountByteLen = int(unsafe.Sizeof(partialResult4Count(0)))
+
 func (*baseCount) AllocPartialResult() (pr PartialResult, memDelta int64) {
 	return PartialResult(new(partialResult4Count)), DefPartialResult4CountSize
 }
@@ -47,6 +49,40 @@ func (e *baseCount) AppendFinalResult2Chunk(_ sessionctx.Context, pr PartialResu
 	return nil
 }
 
+func (c *baseCount) SerializeForSpill(_ sessionctx.Context, partialResult PartialResult, chk *chunk.Chunk, spillHelper *SpillSerializeHelper) {
+	pr := (*partialResult4Count)(partialResult)
+	resBuf := spillHelper.serializePartialResult4Count(*pr)
+	chk.AppendBytes(c.ordinal, resBuf)
+}
+
+func (c *baseCount) DeserializeToPartialResultForSpill(_ sessionctx.Context, src *chunk.Chunk) ([]PartialResult, int64) {
+	dataCol := src.Column(c.ordinal)
+	totalMemDelta := int64(0)
+	spillHelper := newDeserializeHelper(dataCol, src.NumRows())
+	partialResults := make([]PartialResult, 0, src.NumRows())
+
+	for {
+		pr, memDelta := c.deserializeForSpill(&spillHelper)
+		if pr == nil {
+			break
+		}
+		partialResults = append(partialResults, pr)
+		totalMemDelta += memDelta
+	}
+
+	return partialResults, totalMemDelta
+}
+
+func (c *baseCount) deserializeForSpill(helper *spillDeserializeHelper) (PartialResult, int64) {
+	pr, memDelta := c.AllocPartialResult()
+	result := int64(*(*partialResult4Count)(pr))
+	success := helper.deserializePartialResult4Count(&result)
+	if !success {
+		return nil, 0
+	}
+	return pr, memDelta
+}
+
 type countOriginal4Int struct {
 	baseCount
 }

executor/aggfuncs/func_count_distinct.go

@@ -142,6 +142,15 @@ func (e *countOriginalWithDistinct4Real) UpdatePartialResult(sctx sessionctx.Con
 	return memDelta, nil
 }
 
+// TODO implement it
+func (e *countOriginalWithDistinct4Real) SerializeForSpill(_ sessionctx.Context, partialResult PartialResult, chk *chunk.Chunk, spillHelper *SpillSerializeHelper) {
+}
+
+// TODO implement it
+func (e *countOriginalWithDistinct4Real) DeserializeToPartialResultForSpill(sctx sessionctx.Context, src *chunk.Chunk) ([]PartialResult, int64) {
+	return nil, 0
+}
+
 type partialResult4CountDistinctDecimal struct {
 	valSet set.StringSetWithMemoryUsage
 }
@@ -194,6 +203,15 @@ func (e *countOriginalWithDistinct4Decimal) UpdatePartialResult(sctx sessionctx.
 	return memDelta, nil
 }
 
+// TODO implement it
+func (e *countOriginalWithDistinct4Decimal) SerializeForSpill(_ sessionctx.Context, partialResult PartialResult, chk *chunk.Chunk, spillHelper *SpillSerializeHelper) {
+}
+
+// TODO implement it
+func (e *countOriginalWithDistinct4Decimal) DeserializeToPartialResultForSpill(sctx sessionctx.Context, src *chunk.Chunk) ([]PartialResult, int64) {
+	return nil, 0
+}
+
 type partialResult4CountDistinctDuration struct {
 	valSet set.Int64SetWithMemoryUsage
 }
@@ -241,6 +259,15 @@ func (e *countOriginalWithDistinct4Duration) UpdatePartialResult(sctx sessionctx
 	return memDelta, nil
 }
 
+// TODO implement it
+func (e *countOriginalWithDistinct4Duration) SerializeForSpill(_ sessionctx.Context, partialResult PartialResult, chk *chunk.Chunk, spillHelper *SpillSerializeHelper) {
+}
+
+// TODO implement it
+func (e *countOriginalWithDistinct4Duration) DeserializeToPartialResultForSpill(sctx sessionctx.Context, src *chunk.Chunk) ([]PartialResult, int64) {
+	return nil, 0
+}
+
 type partialResult4CountDistinctString struct {
 	valSet set.StringSetWithMemoryUsage
 }
@@ -292,6 +319,15 @@ func (e *countOriginalWithDistinct4String) UpdatePartialResult(sctx sessionctx.C
 	return memDelta, nil
 }
 
+// TODO implement it
+func (e *countOriginalWithDistinct4String) SerializeForSpill(_ sessionctx.Context, partialResult PartialResult, chk *chunk.Chunk, spillHelper *SpillSerializeHelper) {
+}
+
+// TODO implement it
+func (e *countOriginalWithDistinct4String) DeserializeToPartialResultForSpill(sctx sessionctx.Context, src *chunk.Chunk) ([]PartialResult, int64) {
+	return nil, 0
+}
+
 type countOriginalWithDistinct struct {
 	baseCount
 }
@@ -355,6 +391,15 @@ func (e *countOriginalWithDistinct) UpdatePartialResult(sctx sessionctx.Context,
 	return memDelta, nil
 }
 
+// TODO implement it
+func (e *countOriginalWithDistinct) SerializeForSpill(_ sessionctx.Context, partialResult PartialResult, chk *chunk.Chunk, spillHelper *SpillSerializeHelper) {
+}
+
+// TODO implement it
+func (e *countOriginalWithDistinct) DeserializeToPartialResultForSpill(sctx sessionctx.Context, src *chunk.Chunk) ([]PartialResult, int64) {
+	return nil, 0
+}
+
 // evalAndEncode eval one row with an expression and encode value to bytes.
 func evalAndEncode(
 	sctx sessionctx.Context, arg expression.Expression, collator collate.Collator,
@@ -769,6 +814,15 @@ func (*baseApproxCountDistinct) MergePartialResult(_ sessionctx.Context, src, ds
 	return 0, nil
 }
 
+// TODO implement it
+func (e *baseApproxCountDistinct) SerializeForSpill(_ sessionctx.Context, partialResult PartialResult, chk *chunk.Chunk, spillHelper *SpillSerializeHelper) {
+}
+
+// TODO implement it
+func (e *baseApproxCountDistinct) DeserializeToPartialResultForSpill(sctx sessionctx.Context, src *chunk.Chunk) ([]PartialResult, int64) {
+	return nil, 0
+}
+
 type approxCountDistinctOriginal struct {
 	baseApproxCountDistinct
 }

executor/aggfuncs/func_group_concat.go

@@ -265,6 +265,15 @@ func (e *groupConcatDistinct) UpdatePartialResult(sctx sessionctx.Context, rowsI
 	return memDelta, nil
 }
 
+// TODO implement it
+func (c *groupConcatDistinct) SerializeForSpill(_ sessionctx.Context, partialResult PartialResult, chk *chunk.Chunk, spillHelper *SpillSerializeHelper) {
+}
+
+// TODO implement it
+func (c *groupConcatDistinct) DeserializeToPartialResultForSpill(_ sessionctx.Context, src *chunk.Chunk) ([]PartialResult, int64) {
+	return nil, 0
+}
+
 // SetTruncated will be called in `executorBuilder#buildHashAgg` with duck-type.
 func (e *groupConcatDistinct) SetTruncated(t *int32) {
 	e.truncated = t
@@ -611,6 +620,15 @@ func (*groupConcatDistinctOrder) MergePartialResult(sessionctx.Context, PartialR
 	return 0, plannercore.ErrInternal.GenWithStack("groupConcatDistinctOrder.MergePartialResult should not be called")
 }
 
+// TODO implement it
+func (c *groupConcatDistinctOrder) SerializeForSpill(_ sessionctx.Context, partialResult PartialResult, chk *chunk.Chunk, spillHelper *SpillSerializeHelper) {
+}
+
+// TODO implement it
+func (c *groupConcatDistinctOrder) DeserializeToPartialResultForSpill(_ sessionctx.Context, src *chunk.Chunk) ([]PartialResult, int64) {
+	return nil, 0
+}
+
 // GetDatumMemSize calculates the memory size of each types.Datum in sortRow.byItems.
 // types.Datum memory size = variable type's memory size + variable value's memory size.
 func GetDatumMemSize(d *types.Datum) int64 {