planner: make group expression inherit logical plan & refine rule int…

…erface. (#57782)

ref #51664

pkg/planner/cascades/memo/BUILD.bazel

@@ -7,13 +7,13 @@ go_library(
         "group_expr.go",
         "group_id_generator.go",
         "memo.go",
-        "memo_expr.go",
     ],
     importpath = "github.com/pingcap/tidb/pkg/planner/cascades/memo",
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/planner/cascades/base",
         "//pkg/planner/cascades/pattern",
+        "//pkg/planner/cascades/util",
         "//pkg/planner/core/base",
         "//pkg/planner/property",
         "//pkg/sessionctx",

pkg/planner/cascades/memo/group.go

@@ -17,11 +17,11 @@ package memo
 import (
 	"container/list"
 	"fmt"
-	"io"
 	"strconv"
 
 	"github.com/pingcap/tidb/pkg/planner/cascades/base"
 	"github.com/pingcap/tidb/pkg/planner/cascades/pattern"
+	"github.com/pingcap/tidb/pkg/planner/cascades/util"
 	"github.com/pingcap/tidb/pkg/planner/property"
 )
 
@@ -75,9 +75,12 @@ func (g *Group) Equals(other any) bool {
 // ******************************************* end of HashEqual methods *******************************************
 
 // Exists checks whether a Group expression existed in a Group.
-func (g *Group) Exists(hash64u uint64) bool {
-	_, ok := g.hash2GroupExpr[hash64u]
-	return ok
+func (g *Group) Exists(e *GroupExpression) bool {
+	one, ok := g.hash2GroupExpr[e.GetHash64()]
+	if !ok {
+		return false
+	}
+	return one.Value.(*GroupExpression).Equals(e)
 }
 
 // Insert adds a GroupExpression to the Group.
@@ -86,11 +89,10 @@ func (g *Group) Insert(e *GroupExpression) bool {
 		return false
 	}
 	// GroupExpressions hash should be initialized within Init(xxx) method.
-	hash64 := e.Sum64()
-	if g.Exists(hash64) {
+	if g.Exists(e) {
 		return false
 	}
-	operand := pattern.GetOperand(e.logicalPlan)
+	operand := pattern.GetOperand(e.LogicalPlan)
 	var newEquiv *list.Element
 	mark, ok := g.Operand2FirstExpr[operand]
 	if ok {
@@ -101,7 +103,7 @@ func (g *Group) Insert(e *GroupExpression) bool {
 		newEquiv = g.logicalExpressions.PushBack(e)
 		g.Operand2FirstExpr[operand] = newEquiv
 	}
-	g.hash2GroupExpr[hash64] = newEquiv
+	g.hash2GroupExpr[e.GetHash64()] = newEquiv
 	e.group = g
 	return true
 }
@@ -126,8 +128,8 @@ func (g *Group) GetFirstElem(operand pattern.Operand) *list.Element {
 }
 
 // String implements fmt.Stringer interface.
-func (g *Group) String(w io.Writer) {
-	fmt.Fprintf(w, "inputs:%s", strconv.Itoa(int(g.groupID)))
+func (g *Group) String(w util.StrBufferWriter) {
+	w.WriteString(fmt.Sprintf("inputs:%s", strconv.Itoa(int(g.groupID))))
 }
 
 // NewGroup creates a new Group with given logical prop.

pkg/planner/cascades/memo/group_and_expr_test.go

@@ -52,13 +52,13 @@ func TestGroupExpressionHashEquals(t *testing.T) {
 	a := GroupExpression{
 		group:       &Group{groupID: 3},
 		Inputs:      []*Group{child1, child2},
-		logicalPlan: &logicalop.LogicalProjection{Exprs: []expression.Expression{expression.NewOne()}},
+		LogicalPlan: &logicalop.LogicalProjection{Exprs: []expression.Expression{expression.NewOne()}},
 	}
 	b := GroupExpression{
 		// root group should change the hash.
 		group:       &Group{groupID: 4},
 		Inputs:      []*Group{child1, child2},
-		logicalPlan: &logicalop.LogicalProjection{Exprs: []expression.Expression{expression.NewOne()}},
+		LogicalPlan: &logicalop.LogicalProjection{Exprs: []expression.Expression{expression.NewOne()}},
 	}
 	a.Hash64(hasher1)
 	b.Hash64(hasher2)

pkg/planner/cascades/memo/group_expr.go

@@ -15,10 +15,9 @@
 package memo
 
 import (
-	"io"
-
 	base2 "github.com/pingcap/tidb/pkg/planner/cascades/base"
 	"github.com/pingcap/tidb/pkg/planner/cascades/pattern"
+	"github.com/pingcap/tidb/pkg/planner/cascades/util"
 	"github.com/pingcap/tidb/pkg/planner/core/base"
 	"github.com/pingcap/tidb/pkg/util/intest"
 )
@@ -28,52 +27,48 @@ import (
 // coupling between Group and GroupExpression is the key to the success of the memory compact
 // of representing a forest.
 type GroupExpression struct {
+	// LogicalPlan is internal logical expression stands for this groupExpr.
+	// Define it in the header element can make GE as Logical Plan implementor.
+	base.LogicalPlan
+
 	// group is the Group that this GroupExpression belongs to.
 	group *Group
 
 	// inputs stores the Groups that this GroupExpression based on.
 	Inputs []*Group
 
-	// logicalPlan is internal logical expression stands for this groupExpr.
-	logicalPlan base.LogicalPlan
-
 	// hash64 is the unique fingerprint of the GroupExpression.
 	hash64 uint64
 }
 
-// GetLogicalPlan returns the logical plan of the GroupExpression.
-func (e *GroupExpression) GetLogicalPlan() base.LogicalPlan {
-	return e.logicalPlan
-}
-
 // GetGroup returns the Group that this GroupExpression belongs to.
 func (e *GroupExpression) GetGroup() *Group {
 	return e.group
 }
 
 // String implements the fmt.Stringer interface.
-func (e *GroupExpression) String(w io.Writer) {
-	e.GetLogicalPlan().ExplainID()
-	_, _ = w.Write([]byte("GE:" + e.GetLogicalPlan().ExplainID().String() + "{"))
+func (e *GroupExpression) String(w util.StrBufferWriter) {
+	e.LogicalPlan.ExplainID()
+	w.WriteString("GE:" + e.LogicalPlan.ExplainID().String() + "{")
 	for i, input := range e.Inputs {
 		if i != 0 {
-			_, _ = w.Write([]byte(", "))
+			w.WriteString(", ")
 		}
 		input.String(w)
 	}
-	_, _ = w.Write([]byte("}"))
+	w.WriteString("}")
 }
 
-// Sum64 returns the cached hash64 of the GroupExpression.
-func (e *GroupExpression) Sum64() uint64 {
+// GetHash64 returns the cached hash64 of the GroupExpression.
+func (e *GroupExpression) GetHash64() uint64 {
 	intest.Assert(e.hash64 != 0, "hash64 should not be 0")
 	return e.hash64
 }
 
 // Hash64 implements the Hash64 interface.
 func (e *GroupExpression) Hash64(h base2.Hasher) {
 	// logical plan hash.
-	e.logicalPlan.Hash64(h)
+	e.LogicalPlan.Hash64(h)
 	// children group hash.
 	for _, child := range e.Inputs {
 		child.Hash64(h)
@@ -95,12 +90,12 @@ func (e *GroupExpression) Equals(other any) bool {
 	if len(e.Inputs) != len(e2.Inputs) {
 		return false
 	}
-	if pattern.GetOperand(e.logicalPlan) != pattern.GetOperand(e2.logicalPlan) {
+	if pattern.GetOperand(e.LogicalPlan) != pattern.GetOperand(e2.LogicalPlan) {
 		return false
 	}
 	// current logical operator meta cmp, logical plan don't care logicalPlan's children.
 	// when we convert logicalPlan to GroupExpression, we will set children to nil.
-	if !e.logicalPlan.Equals(e2.logicalPlan) {
+	if !e.LogicalPlan.Equals(e2.LogicalPlan) {
 		return false
 	}
 	// if one of the children is different, then the two GroupExpressions are different.
@@ -117,7 +112,7 @@ func NewGroupExpression(lp base.LogicalPlan, inputs []*Group) *GroupExpression {
 	return &GroupExpression{
 		group:       nil,
 		Inputs:      inputs,
-		logicalPlan: lp,
+		LogicalPlan: lp,
 		hash64:      0,
 	}
 }

pkg/planner/cascades/rule/BUILD.bazel

@@ -12,7 +12,8 @@ go_library(
     deps = [
         "//pkg/planner/cascades/memo",
         "//pkg/planner/cascades/pattern",
-        "//pkg/sessionctx",
+        "//pkg/planner/cascades/util",
+        "//pkg/planner/core/base",
     ],
 )
 
@@ -27,8 +28,9 @@ go_test(
         "//pkg/parser/model",
         "//pkg/planner/cascades/memo",
         "//pkg/planner/cascades/pattern",
-        "//pkg/planner/core",
+        "//pkg/planner/cascades/util",
         "//pkg/planner/core/operator/logicalop",
+        "//pkg/util/mock",
         "@com_github_stretchr_testify//require",
     ],
 )