Repository: matrixorigin/matrixone
Pull request #6758 (2022-11-17 11:40:35)

Fix the memory leak of insert, group by and compile
## What type of PR is this?

- [ ] API-change
- [X] BUG
- [ ] Improvement
- [ ] Documentation
- [ ] Feature
- [ ] Test and CI
- [ ] Code Refactoring

## Which issue(s) this PR fixes:

issue #6528

## What this PR does / why we need it:


Patch:
diff --git a/pkg/sql/colexec/agg/agg.go b/pkg/sql/colexec/agg/agg.go
index 4c88ded6e..51db8ea22 100644
--- a/pkg/sql/colexec/agg/agg.go
+++ b/pkg/sql/colexec/agg/agg.go
@@ -224,6 +224,7 @@ func (a *UnaryAgg[T1, T2]) BatchMerge(b Agg[any], start int64, os []uint8, vps [
 
 func (a *UnaryAgg[T1, T2]) Eval(m *mpool.MPool) (*vector.Vector, error) {
 	defer func() {
+		a.Free(m)
 		a.da = nil
 		a.vs = nil
 		a.es = nil
@@ -308,7 +309,7 @@ func getUnaryAggStrVs(strUnaryAgg any) []string {
 	return result
 }
 
-func (a *UnaryAgg[T1, T2]) UnmarshalBinary(data []byte) error {
+func (a *UnaryAgg[T1, T2]) UnmarshalBinary(data []byte, mp *mpool.MPool) error {
 	decoded := new(EncodeAgg)
 	if err := types.Decode(data, decoded); err != nil {
 		return err
@@ -319,7 +320,14 @@ func (a *UnaryAgg[T1, T2]) UnmarshalBinary(data []byte) error {
 	a.otyp = types.DecodeType(decoded.OutputType)
 	a.isCount = decoded.IsCount
 	a.es = decoded.Es
-	a.da = decoded.Da
+	//	a.da = decoded.Da
+	data, err := mp.Alloc(len(decoded.Da))
+	if err != nil {
+		return err
+	}
+	copy(data, decoded.Da)
+	a.da = data
+
 	setAggValues[T1, T2](a, a.otyp)
 
 	return a.priv.UnmarshalBinary(decoded.Private)
diff --git a/pkg/sql/colexec/agg/distagg.go b/pkg/sql/colexec/agg/distagg.go
index 6e64d8468..03b213fc3 100644
--- a/pkg/sql/colexec/agg/distagg.go
+++ b/pkg/sql/colexec/agg/distagg.go
@@ -312,6 +312,7 @@ func (a *UnaryDistAgg[T1, T2]) BatchMerge(b Agg[any], start int64, os []uint8, v
 
 func (a *UnaryDistAgg[T1, T2]) Eval(m *mpool.MPool) (*vector.Vector, error) {
 	defer func() {
+		a.Free(m)
 		a.da = nil
 		a.vs = nil
 		a.es = nil
@@ -400,7 +401,7 @@ func getDistAggStrVs(strUnaryDistAgg any) []string {
 	return result
 }
 
-func (a *UnaryDistAgg[T1, T2]) UnmarshalBinary(data []byte) error {
+func (a *UnaryDistAgg[T1, T2]) UnmarshalBinary(data []byte, m *mpool.MPool) error {
 	decode := new(EncodeAggDistinct[T1])
 	if err := types.Decode(data, decode); err != nil {
 		return err
@@ -411,10 +412,29 @@ func (a *UnaryDistAgg[T1, T2]) UnmarshalBinary(data []byte) error {
 	a.ityps = decode.InputType
 	a.otyp = decode.OutputType
 	a.es = decode.Es
-	a.da = decode.Da
+	data, err := m.Alloc(len(decode.Da))
+	if err != nil {
+		return err
+	}
+	copy(data, decode.Da)
+	a.da = data
 	setDistAggValues[T1, T2](a, a.otyp)
 	a.srcs = decode.Srcs
-
+	a.maps = make([]*hashmap.StrHashMap, len(a.srcs))
+	for i, src := range a.srcs {
+		mp, err := hashmap.NewStrMap(true, 0, 0, m)
+		if err != nil {
+			m.Free(data)
+			for j := 0; j < i; j++ {
+				a.maps[j].Free()
+			}
+			return err
+		}
+		a.maps[i] = mp
+		for _, v := range src {
+			mp.InsertValue(v)
+		}
+	}
 	return a.priv.UnmarshalBinary(decode.Private)
 }
 
diff --git a/pkg/sql/colexec/agg/types.go b/pkg/sql/colexec/agg/types.go
index 40d5b8e8e..85a8e1158 100644
--- a/pkg/sql/colexec/agg/types.go
+++ b/pkg/sql/colexec/agg/types.go
@@ -16,6 +16,7 @@ package agg
 
 import (
 	"encoding"
+
 	"github.com/matrixorigin/matrixone/pkg/common/hashmap"
 	"github.com/matrixorigin/matrixone/pkg/common/mpool"
 	"github.com/matrixorigin/matrixone/pkg/container/nulls"
@@ -66,7 +67,8 @@ type Aggregate struct {
 // Agg agg interface
 type Agg[T any] interface {
 	encoding.BinaryMarshaler
-	encoding.BinaryUnmarshaler
+	//encoding.BinaryUnmarshaler
+	UnmarshalBinary(data []byte, m *mpool.MPool) error
 
 	// Dup will duplicate a new agg with the same type.
 	Dup() Agg[any]
diff --git a/pkg/sql/compile/compile.go b/pkg/sql/compile/compile.go
index 0e154edbb..c33fea381 100644
--- a/pkg/sql/compile/compile.go
+++ b/pkg/sql/compile/compile.go
@@ -886,6 +886,7 @@ func (c *Compile) compileSort(n *plan.Node, ss []*Scope) []*Scope {
 		if err != nil {
 			panic(err)
 		}
+		defer vec.Free(c.proc.Mp())
 		return c.compileTop(n, vec.Col.([]int64)[0], ss)
 	case n.Limit == nil && n.Offset == nil && len(n.OrderBy) > 0: // top
 		return c.compileOrder(n, ss)
@@ -894,10 +895,12 @@ func (c *Compile) compileSort(n *plan.Node, ss []*Scope) []*Scope {
 		if err != nil {
 			panic(err)
 		}
+		defer vec1.Free(c.proc.Mp())
 		vec2, err := colexec.EvalExpr(constBat, c.proc, n.Offset)
 		if err != nil {
 			panic(err)
 		}
+		defer vec2.Free(c.proc.Mp())
 		limit, offset := vec1.Col.([]int64)[0], vec2.Col.([]int64)[0]
 		topN := limit + offset
 		if topN <= 8192*2 {
diff --git a/pkg/sql/compile/dml.go b/pkg/sql/compile/dml.go
index 8bd4fefc3..fe2893441 100644
--- a/pkg/sql/compile/dml.go
+++ b/pkg/sql/compile/dml.go
@@ -179,363 +179,26 @@ func (s *Scope) InsertValues(c *Compile, stmt *tree.Insert) (uint64, error) {
 
 // XXX: is this just fill batch with first vec.Col[0]?
 func fillBatch(bat *batch.Batch, p *plan.InsertValues, rows []tree.Exprs, proc *process.Process) error {
-	rowCount := len(p.Columns[0].Column)
-
 	tmpBat := batch.NewWithSize(0)
 	tmpBat.Zs = []int64{1}
 
 	for i, v := range bat.Vecs {
-		switch v.Typ.Oid {
-		case types.T_uuid:
-			vs := make([]types.Uuid, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vector.GetValueAt[types.Uuid](vec, 0)
-					}
-				}
-			}
-			if err := vector.AppendFixed(v, vs, proc.Mp()); err != nil {
-				return err
-			}
-		case types.T_bool:
-			vs := make([]bool, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vector.GetValueAt[bool](vec, 0)
-					}
-				}
-			}
-			if err := vector.AppendFixed(v, vs, proc.Mp()); err != nil {
-				return err
-			}
-		case types.T_int8:
-			vs := make([]int8, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vector.GetValueAt[int8](vec, 0)
-					}
-				}
-			}
-			if err := vector.AppendFixed(v, vs, proc.Mp()); err != nil {
-				return err
-			}
-		case types.T_int16:
-			vs := make([]int16, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vector.GetValueAt[int16](vec, 0)
-					}
-				}
-			}
-			if err := vector.AppendFixed(v, vs, proc.Mp()); err != nil {
-				return err
-			}
-		case types.T_int32:
-			vs := make([]int32, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vector.GetValueAt[int32](vec, 0)
-					}
-				}
-			}
-			if err := vector.AppendFixed(v, vs, proc.Mp()); err != nil {
-				return err
-			}
-		case types.T_int64:
-			vs := make([]int64, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vector.GetValueAt[int64](vec, 0)
-					}
-				}
-			}
-			if err := vector.AppendFixed(v, vs, proc.Mp()); err != nil {
-				return err
-			}
-		case types.T_uint8:
-			vs := make([]uint8, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vector.GetValueAt[uint8](vec, 0)
-					}
-				}
-			}
-			if err := vector.AppendFixed(v, vs, proc.Mp()); err != nil {
-				return err
-			}
-		case types.T_uint16:
-			vs := make([]uint16, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vector.GetValueAt[uint16](vec, 0)
-					}
-				}
-			}
-			if err := vector.AppendFixed(v, vs, proc.Mp()); err != nil {
-				return err
-			}
-		case types.T_uint32:
-			vs := make([]uint32, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vector.GetValueAt[uint32](vec, 0)
-					}
-				}
-			}
-			if err := vector.AppendFixed(v, vs, proc.Mp()); err != nil {
-				return err
-			}
-		case types.T_uint64:
-			vs := make([]uint64, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vector.GetValueAt[uint64](vec, 0)
-					}
-				}
-			}
-			if err := vector.AppendFixed(v, vs, proc.Mp()); err != nil {
-				return err
-			}
-		case types.T_float32:
-			vs := make([]float32, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vector.GetValueAt[float32](vec, 0)
-					}
-				}
-			}
-			if err := vector.AppendFixed(v, vs, proc.Mp()); err != nil {
-				return err
-			}
-		case types.T_float64:
-			vs := make([]float64, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vector.GetValueAt[float64](vec, 0)
-					}
-				}
-			}
-			if err := vector.AppendFixed(v, vs, proc.Mp()); err != nil {
-				return err
-			}
-		case types.T_char, types.T_varchar, types.T_json, types.T_blob, types.T_text:
-			vs := make([][]byte, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vec.GetBytes(0)
-					}
-				}
-			}
-			if err := vector.AppendBytes(v, vs, proc.Mp()); err != nil {
-				return err
-			}
-		case types.T_date:
-			vs := make([]types.Date, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vector.GetValueAt[types.Date](vec, 0)
-					}
-				}
-			}
-			if err := vector.AppendFixed(v, vs, proc.Mp()); err != nil {
-				return err
-			}
-		case types.T_time:
-			vs := make([]types.Time, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vector.GetValueAt[types.Time](vec, 0)
-					}
-				}
-			}
-			if err := vector.AppendFixed(v, vs, proc.Mp()); err != nil {
-				return err
-			}
-		case types.T_datetime:
-			vs := make([]types.Datetime, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vector.GetValueAt[types.Datetime](vec, 0)
-					}
-				}
-			}
-			if err := vector.AppendFixed(v, vs, proc.Mp()); err != nil {
-				return err
-			}
-		case types.T_timestamp:
-			vs := make([]types.Timestamp, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vector.GetValueAt[types.Timestamp](vec, 0)
-					}
-				}
-			}
-			if err := vector.AppendFixed(v, vs, proc.Mp()); err != nil {
-				return err
-			}
-		case types.T_decimal64:
-			vs := make([]types.Decimal64, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vector.GetValueAt[types.Decimal64](vec, 0)
-					}
-				}
+		for j, expr := range p.Columns[i].Column {
+			vec, err := colexec.EvalExpr(tmpBat, proc, expr)
+			if err != nil {
+				return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
 			}
-			if err := vector.AppendFixed(v, vs, proc.Mp()); err != nil {
-				return err
+			if vec.Size() == 0 {
+				vec = vec.ConstExpand(proc.Mp())
 			}
-		case types.T_decimal128:
-			vs := make([]types.Decimal128, rowCount)
-			{
-				for j, expr := range p.Columns[i].Column {
-					vec, err := colexec.EvalExpr(tmpBat, proc, expr)
-					if err != nil {
-						return y.MakeInsertError(v.Typ.Oid, p.ExplicitCols[i], rows, i, j)
-					}
-					if nulls.Any(vec.Nsp) {
-						nulls.Add(v.Nsp, uint64(j))
-					} else {
-						vs[j] = vector.GetValueAt[types.Decimal128](vec, 0)
-					}
-				}
-			}
-			if err := vector.AppendFixed(v, vs, proc.Mp()); err != nil {
+			if err := vector.UnionOne(v, vec, 0, proc.Mp()); err != nil {
+				vec.Free(proc.Mp())
 				return err
 			}
-		default:
-			return moerr.NewInternalError("data truncation: type of '%v' doesn't implement", v.Typ)
+			vec.Free(proc.Mp())
 		}
 	}
-	bat.Zs = make([]int64, len(rows))
-	for i := 0; i < len(rows); i++ {
-		bat.Zs[i] = 1
-	}
+	bat.SetZs(len(rows), proc.Mp())
 	return nil
 }
 
@@ -549,7 +212,8 @@ func makeInsertBatch(p *plan.InsertValues) *batch.Batch {
 		attrs = append(attrs, col.Name)
 	}
 
-	bat := batch.New(true, attrs)
+	bat := batch.NewWithSize(len(attrs))
+	bat.SetAttributes(attrs)
 	idx := 0
 	for _, col := range p.ExplicitCols {
 		bat.Vecs[idx] = vector.New(types.Type{Oid: types.T(col.Typ.GetId()), Scale: col.Typ.Scale, Width: col.Typ.Width})
diff --git a/pkg/sql/compile/operator.go b/pkg/sql/compile/operator.go
index 2823644a2..9a336a1c9 100644
--- a/pkg/sql/compile/operator.go
+++ b/pkg/sql/compile/operator.go
@@ -577,6 +577,7 @@ func constructLimit(n *plan.Node, proc *process.Process) *limit.Argument {
 	if err != nil {
 		panic(err)
 	}
+	defer vec.Free(proc.Mp())
 	return &limit.Argument{
 		Limit: uint64(vec.Col.([]int64)[0]),
 	}
@@ -666,6 +667,7 @@ func constructMergeOffset(n *plan.Node, proc *process.Process) *mergeoffset.Argu
 	if err != nil {
 		panic(err)
 	}
+	defer vec.Free(proc.Mp())
 	return &mergeoffset.Argument{
 		Offset: uint64(vec.Col.([]int64)[0]),
 	}
@@ -676,6 +678,7 @@ func constructMergeLimit(n *plan.Node, proc *process.Process) *mergelimit.Argume
 	if err != nil {
 		panic(err)
 	}
+	defer vec.Free(proc.Mp())
 	return &mergelimit.Argument{
 		Limit: uint64(vec.Col.([]int64)[0]),
 	}
diff --git a/pkg/vm/engine/tae/tables/jobs/compactblk.go b/pkg/vm/engine/tae/tables/jobs/compactblk.go
index 1aacafd46..30b84de4d 100644
--- a/pkg/vm/engine/tae/tables/jobs/compactblk.go
+++ b/pkg/vm/engine/tae/tables/jobs/compactblk.go
@@ -217,6 +217,13 @@ func (task *compactBlockTask) Execute() (err error) {
 		// 	return err
 		// }
 	}
+	if !table.GetSchema().HasSortKey() && task.created != nil {
+		n := task.created.Rows()
+		task.mapping = make([]uint32, n)
+		for i := 0; i < n; i++ {
+			task.mapping[i] = uint32(i)
+		}
+	}
 	txnEntry := txnentries.NewCompactBlockEntry(
 		task.txn,
 		task.compacted,


Test patch:
diff --git a/pkg/sql/colexec/agg/aggUt/agg_testutil.go b/pkg/sql/colexec/agg/aggUt/agg_testutil.go
index aa278a4a2..1ad7afbbf 100644
--- a/pkg/sql/colexec/agg/aggUt/agg_testutil.go
+++ b/pkg/sql/colexec/agg/aggUt/agg_testutil.go
@@ -171,7 +171,7 @@ func RunBaseMarshalTest(t *testing.T, c *testCase) {
 		d, marshalErr := agg0.MarshalBinary()
 		require.NoError(t, marshalErr)
 		agg1, _ := agg.New(c.op, c.isDistinct, c.inputTyp)
-		unmarshalErr := agg1.UnmarshalBinary(d)
+		unmarshalErr := agg1.UnmarshalBinary(d, m)
 		require.NoError(t, unmarshalErr)
 
 		// Fill() after marshal and unmarshal
@@ -217,7 +217,7 @@ func RunBaseMarshalTest(t *testing.T, c *testCase) {
 		d, marshalErr := agg0.MarshalBinary()
 		require.NoError(t, marshalErr)
 		mAgg, _ := agg.New(c.op, c.isDistinct, c.inputTyp)
-		unmarshalErr := mAgg.UnmarshalBinary(d)
+		unmarshalErr := mAgg.UnmarshalBinary(d, m)
 		require.NoError(t, unmarshalErr)
 
 		// Merge()
diff --git a/pkg/sql/compile/compile_test.go b/pkg/sql/compile/compile_test.go
index 901def68a..42ff869b8 100644
--- a/pkg/sql/compile/compile_test.go
+++ b/pkg/sql/compile/compile_test.go
@@ -23,6 +23,7 @@ import (
 	mock_frontend "github.com/matrixorigin/matrixone/pkg/frontend/test"
 	"github.com/matrixorigin/matrixone/pkg/pb/plan"
 	"github.com/matrixorigin/matrixone/pkg/sql/parsers/dialect/mysql"
+	"github.com/matrixorigin/matrixone/pkg/sql/parsers/tree"
 	plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan"
 	"github.com/matrixorigin/matrixone/pkg/testutil"
 	"github.com/matrixorigin/matrixone/pkg/testutil/testengine"
@@ -36,6 +37,7 @@ type compileTestCase struct {
 	sql  string
 	pn   *plan.Plan
 	e    engine.Engine
+	stmt tree.Statement
 	proc *process.Process
 }
 
@@ -57,8 +59,11 @@ func init() {
 		newTestCase("select * from R left join S on R.uid = S.uid", new(testing.T)),
 		newTestCase("select * from R right join S on R.uid = S.uid", new(testing.T)),
 		newTestCase("select * from R join S on R.uid > S.uid", new(testing.T)),
+		newTestCase("select * from R limit 10", new(testing.T)),
+		newTestCase("insert into R values('1', '2', '3')", new(testing.T)),
 		newTestCase("insert into R select * from R", new(testing.T)),
-		newTestCase(`select * from unnest('{"a":1}') as f`, new(testing.T)),
+		newTestCase("select count(*) from R group by uid", new(testing.T)),
+		newTestCase("select count(distinct uid) from R", new(testing.T)),
 	}
 }
 
@@ -78,12 +83,14 @@ func TestCompile(t *testing.T) {
 
 	for _, tc := range tcs {
 		tc.proc.TxnClient = txnClient
-		c := New("", "test", tc.sql, "", context.TODO(), tc.e, tc.proc, nil)
+		c := New("", "test", tc.sql, "", context.TODO(), tc.e, tc.proc, tc.stmt)
 		err := c.Compile(tc.pn, nil, testPrint)
 		require.NoError(t, err)
 		c.GetAffectedRows()
 		err = c.Run(0)
 		require.NoError(t, err)
+		// Enable memory check
+		require.Equal(t, int64(0), tc.proc.Mp().CurrNB())
 	}
 }
 
@@ -103,19 +110,15 @@ func TestCompileWithFaults(t *testing.T) {
 func newTestCase(sql string, t *testing.T) compileTestCase {
 	proc := testutil.NewProcess()
 	e, _, compilerCtx := testengine.New(context.Background())
-	opt := plan2.NewBaseOptimizer(compilerCtx)
 	stmts, err := mysql.Parse(sql)
 	require.NoError(t, err)
-	qry, err := opt.Optimize(stmts[0])
+	pn, err := plan2.BuildPlan(compilerCtx, stmts[0])
 	require.NoError(t, err)
 	return compileTestCase{
 		e:    e,
 		sql:  sql,
 		proc: proc,
-		pn: &plan.Plan{
-			Plan: &plan.Plan_Query{
-				Query: qry,
-			},
-		},
+		pn:   pn,
+		stmt: stmts[0],
 	}
 }