planner: Expression memory trace

expression/builtin.go

@@ -27,6 +27,7 @@ package expression
 import (
 	"strings"
 	"sync"
+	"unsafe"
 
 	"github.com/gogo/protobuf/proto"
 	"github.com/pingcap/errors"
@@ -501,6 +502,8 @@ type builtinFunc interface {
 	// Clone returns a copy of itself.
 	Clone() builtinFunc
 
+	MemoryUsage() int64
+
 	CollationInfo
 }
 
@@ -961,3 +964,20 @@ func (b *baseBuiltinFunc) setDecimalAndFlenForTime(fsp int) {
 		b.tp.SetFlenUnderLimit(b.tp.GetFlen() + 1)
 	}
 }
+
+const emptyBaseBuiltinFunc = int64(unsafe.Sizeof(baseBuiltinFunc{}))
+
+// MemoryUsage return the memory usage of baseBuiltinFunc
+func (b *baseBuiltinFunc) MemoryUsage() (sum int64) {
+	if b == nil {
+		return
+	}
+
+	sum = emptyBaseBuiltinFunc + b.bufAllocator.MemoryUsage() +
+		b.tp.MemoryUsage() + int64(unsafe.Sizeof(*b.childrenVectorizedOnce)) +
+		int64(unsafe.Sizeof(*b.childrenReversedOnce)) + int64(len(b.charset)+len(b.collation))
+	for _, e := range b.args {
+		sum += e.MemoryUsage()
+	}
+	return
+}

expression/builtin_vectorized.go

@@ -16,6 +16,7 @@ package expression
 
 import (
 	"sync"
+	"unsafe"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/parser/mysql"
@@ -29,6 +30,8 @@ type columnBufferAllocator interface {
 	get() (*chunk.Column, error)
 	// put releases a column buffer.
 	put(buf *chunk.Column)
+	// MemoryUsage return the memory usage of columnBufferAllocator
+	MemoryUsage() int64
 }
 
 // localColumnPool implements columnBufferAllocator interface.
@@ -72,6 +75,12 @@ func (r *localColumnPool) put(col *chunk.Column) {
 	r.Pool.Put(col)
 }
 
+const emptyLocalColumnPoolSize = int64(unsafe.Sizeof(localColumnPool{}))
+
+func (r *localColumnPool) MemoryUsage() (sum int64) {
+	return emptyLocalColumnPoolSize
+}
+
 // vecEvalIntByRows uses the non-vectorized(row-based) interface `evalInt` to eval the expression.
 func vecEvalIntByRows(sig builtinFunc, input *chunk.Chunk, result *chunk.Column) error {
 	n := input.NumRows()

expression/column.go

@@ -17,6 +17,7 @@ package expression
 import (
 	"fmt"
 	"strings"
+	"unsafe"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/parser/ast"
@@ -191,6 +192,16 @@ func (col *CorrelatedColumn) resolveIndicesByVirtualExpr(_ *Schema) bool {
 	return true
 }
 
+// MemoryUsage return the memory usage of CorrelatedColumn
+func (col *CorrelatedColumn) MemoryUsage() (sum int64) {
+	if col == nil {
+		return
+	}
+
+	sum = col.Column.MemoryUsage() + int64(unsafe.Sizeof(col.Data))
+	return sum
+}
+
 // Column represents a column.
 type Column struct {
 	RetType *types.FieldType
@@ -722,3 +733,20 @@ func GcColumnExprIsTidbShard(virtualExpr Expression) bool {
 
 	return true
 }
+
+const emptyColumnSize = int64(unsafe.Sizeof(Column{}))
+
+// MemoryUsage return the memory usage of Column
+func (col *Column) MemoryUsage() (sum int64) {
+	if col == nil {
+		return
+	}
+
+	sum = emptyColumnSize + col.RetType.MemoryUsage() + int64(cap(col.hashcode))*int64(unsafe.Sizeof(*new(byte))) +
+		int64(len(col.OrigName)+len(col.charset)+len(col.collation))
+
+	if col.VirtualExpr != nil {
+		sum += col.VirtualExpr.MemoryUsage()
+	}
+	return
+}

expression/column_test.go

@@ -263,3 +263,13 @@ func TestGcColumnExprIsTidbShard(t *testing.T) {
 	shardExpr := NewFunctionInternal(ctx, ast.TiDBShard, ft, col)
 	require.True(t, GcColumnExprIsTidbShard(shardExpr))
 }
+
+func TestColumnMemoryUsage(t *testing.T) {
+	c1 := &Column{OrigName: "Origin"}
+	c2 := Column{OrigName: "OriginName"}
+
+	require.Greater(t, c2.MemoryUsage(), c1.MemoryUsage())
+
+	c1 = nil
+	require.Equal(t, c1.MemoryUsage(), int64(0))
+}

expression/constant.go

@@ -16,6 +16,7 @@ package expression
 
 import (
 	"fmt"
+	"unsafe"
 
 	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/parser/terror"
@@ -449,3 +450,16 @@ func (c *Constant) Coercibility() Coercibility {
 	}
 	return c.collationInfo.Coercibility()
 }
+
+const emptyConstantSize = int64(unsafe.Sizeof(Constant{}))
+
+// MemoryUsage return the memory usage of Constant
+func (c *Constant) MemoryUsage() (sum int64) {
+	if c == nil {
+		return
+	}
+
+	sum = emptyConstantSize + c.RetType.MemoryUsage() + c.Value.MemUsage() +
+		int64(cap(c.hashcode))*int64(unsafe.Sizeof(*new(byte)))
+	return
+}

expression/constant_test.go

@@ -506,3 +506,10 @@ func TestSpecificConstant(t *testing.T) {
 	require.Equal(t, null.RetType.GetFlen(), 1)
 	require.Equal(t, null.RetType.GetDecimal(), 0)
 }
+
+func TestConstantMemoryUsage(t *testing.T) {
+	c1 := Constant{Value: types.NewIntDatum(1)}
+	c2 := Constant{Value: types.NewStringDatum("11")}
+
+	require.Greater(t, c2.MemoryUsage(), c1.MemoryUsage())
+}

expression/expression.go

@@ -176,6 +176,9 @@ type Expression interface {
 	// Column: ColumnFlag+encoded value
 	// ScalarFunction: SFFlag+encoded function name + encoded arg_1 + encoded arg_2 + ...
 	HashCode(sc *stmtctx.StatementContext) []byte
+
+	// MemoryUsage return the memory usage of Expression
+	MemoryUsage() int64
 }
 
 // CNFExprs stands for a CNF expression.

expression/scalar_function.go

@@ -17,6 +17,7 @@ package expression
 import (
 	"bytes"
 	"fmt"
+	"unsafe"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/parser/ast"
@@ -590,3 +591,16 @@ func (sf *ScalarFunction) Repertoire() Repertoire {
 func (sf *ScalarFunction) SetRepertoire(r Repertoire) {
 	sf.Function.SetRepertoire(r)
 }
+
+const emptyScalarFunctionSize = int64(unsafe.Sizeof(ScalarFunction{}))
+
+// MemoryUsage return the memory usage of ScalarFunction
+func (sf *ScalarFunction) MemoryUsage() (sum int64) {
+	if sf == nil {
+		return
+	}
+
+	sum = emptyScalarFunctionSize + int64(len(sf.FuncName.L)+len(sf.FuncName.O)) + sf.RetType.MemoryUsage() +
+		int64(cap(sf.hashcode))*int64(unsafe.Sizeof(*new(byte))) + sf.Function.MemoryUsage()
+	return sum
+}

expression/util_test.go

@@ -594,3 +594,7 @@ func (m *MockExpr) CharsetAndCollation() (string, string) {
 	return "", ""
 }
 func (m *MockExpr) SetCharsetAndCollation(chs, coll string) {}
+
+func (m *MockExpr) MemoryUsage() (sum int64) {
+	return
+}

parser/types/field_type.go

@@ -18,6 +18,7 @@ import (
 	"fmt"
 	"io"
 	"strings"
+	"unsafe"
 
 	"github.com/cznic/mathutil"
 	"github.com/pingcap/tidb/parser/charset"
@@ -601,3 +602,21 @@ func (ft *FieldType) MarshalJSON() ([]byte, error) {
 	r.ElemsIsBinaryLit = ft.elemsIsBinaryLit
 	return json.Marshal(r)
 }
+
+const emptyFieldTypeSize = int64(unsafe.Sizeof(FieldType{}))
+
+// MemoryUsage return the memory usage of FieldType
+func (ft *FieldType) MemoryUsage() (sum int64) {
+	if ft == nil {
+		return
+	}
+	sum = emptyFieldTypeSize +
+		int64(len(ft.charset)+len(ft.collate))
+
+	for _, s := range ft.elems {
+		sum += int64(len(s))
+	}
+	sum += int64(cap(ft.elems)) * int64(unsafe.Sizeof(*new(string)))
+	sum += int64(cap(ft.elemsIsBinaryLit)) * int64(unsafe.Sizeof(unsafe.Sizeof(*new(bool))))
+	return
+}