types: add a interface to eval the memory cost of datum array (pingcap#14185)

Author: XuHuaiyu

types/datum.go

@@ -21,6 +21,7 @@ import (
 	"strings"
 	"time"
 	"unicode/utf8"
+	"unsafe"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/parser/charset"
@@ -2057,3 +2058,30 @@ func ChangeReverseResultByUpperLowerBound(
 	}
 	return d, nil
 }
+
+const (
+	sizeOfEmptyDatum = int(unsafe.Sizeof(Datum{}))
+	sizeOfMysqlTime  = int(unsafe.Sizeof(Time{}))
+	sizeOfMyDecimal  = MyDecimalStructSize
+)
+
+// EstimatedMemUsage returns the estimated bytes consumed of a one-dimensional
+// or two-dimensional datum array.
+func EstimatedMemUsage(array []Datum, numOfRows int) int64 {
+	if numOfRows == 0 {
+		return 0
+	}
+	var bytesConsumed int
+	for _, d := range array {
+		switch d.Kind() {
+		case KindMysqlDecimal:
+			bytesConsumed += sizeOfMyDecimal
+		case KindMysqlTime:
+			bytesConsumed += sizeOfMysqlTime
+		default:
+			bytesConsumed += len(d.b)
+		}
+	}
+	bytesConsumed += len(array) * sizeOfEmptyDatum
+	return int64(bytesConsumed * numOfRows)
+}

types/datum_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/pingcap/parser/mysql"
 	"github.com/pingcap/tidb/sessionctx/stmtctx"
 	"github.com/pingcap/tidb/types/json"
+	"github.com/pingcap/tidb/util/hack"
 )
 
 var _ = Suite(&testDatumSuite{})
@@ -396,6 +397,25 @@ func newRetTypeWithFlenDecimal(tp byte, flen int, decimal int) *FieldType {
 	}
 }
 
+func (ts *testDatumSuite) TestEstimatedMemUsage(c *C) {
+	b := []byte{'a', 'b', 'c', 'd'}
+	enum := Enum{Name: "a", Value: 1}
+	datumArray := []Datum{
+		NewIntDatum(1),
+		NewFloat64Datum(1.0),
+		NewFloat32Datum(1.0),
+		NewStringDatum(string(b)),
+		NewBytesDatum(b),
+		NewDecimalDatum(newMyDecimal("1234.1234", c)),
+		NewMysqlEnumDatum(enum),
+	}
+	bytesConsumed := 10 * (len(datumArray)*sizeOfEmptyDatum +
+		sizeOfMyDecimal +
+		len(b)*2 +
+		len(hack.Slice(enum.Name)))
+	c.Assert(int(EstimatedMemUsage(datumArray, 10)), Equals, bytesConsumed)
+}
+
 func (ts *testDatumSuite) TestChangeReverseResultByUpperLowerBound(c *C) {
 	sc := new(stmtctx.StatementContext)
 	sc.IgnoreTruncate = true