planner: move projection elimination related ReplaceExpr/Column code to rule util.

pkg/planner/cascades/BUILD.bazel

@@ -21,6 +21,7 @@ go_library(
         "//pkg/planner/core",
         "//pkg/planner/core/base",
         "//pkg/planner/core/operator/logicalop",
+        "//pkg/planner/core/rule/util",
         "//pkg/planner/implementation",
         "//pkg/planner/memo",
         "//pkg/planner/pattern",

pkg/planner/cascades/implementation_rules.go

@@ -178,10 +178,10 @@ func (*ImplTiKVSingleReadGather) OnImplement(expr *memo.GroupExpr, reqProp *prop
 	logicProp := expr.Group.Prop
 	sg := expr.ExprNode.(*plannercore.TiKVSingleGather)
 	if sg.IsIndexGather {
-		reader := sg.GetPhysicalIndexReader(logicProp.Schema, logicProp.Stats.ScaleByExpectCnt(reqProp.ExpectedCnt), reqProp)
+		reader := plannercore.GetPhysicalIndexReader(sg, logicProp.Schema, logicProp.Stats.ScaleByExpectCnt(reqProp.ExpectedCnt), reqProp)
 		return []memo.Implementation{impl.NewIndexReaderImpl(reader, sg.Source)}, nil
 	}
-	reader := sg.GetPhysicalTableReader(logicProp.Schema, logicProp.Stats.ScaleByExpectCnt(reqProp.ExpectedCnt), reqProp)
+	reader := plannercore.GetPhysicalTableReader(sg, logicProp.Schema, logicProp.Stats.ScaleByExpectCnt(reqProp.ExpectedCnt), reqProp)
 	return []memo.Implementation{impl.NewTableReaderImpl(reader, sg.Source)}, nil
 }
 

pkg/planner/cascades/transformation_rules.go

@@ -25,6 +25,7 @@ import (
 	"github.com/pingcap/tidb/pkg/planner/context"
 	plannercore "github.com/pingcap/tidb/pkg/planner/core"
 	"github.com/pingcap/tidb/pkg/planner/core/base"
+	ruleutil "github.com/pingcap/tidb/pkg/planner/core/rule/util"
 	"github.com/pingcap/tidb/pkg/planner/memo"
 	"github.com/pingcap/tidb/pkg/planner/pattern"
 	"github.com/pingcap/tidb/pkg/planner/util"
@@ -1180,7 +1181,7 @@ func (*MergeAdjacentProjection) OnTransform(old *memo.ExprIter) (newExprs []*mem
 	newProj.SetSchema(old.GetExpr().Group.Prop.Schema)
 	for i, expr := range proj.Exprs {
 		newExpr := expr.Clone()
-		plannercore.ResolveExprAndReplace(newExpr, replace)
+		ruleutil.ResolveExprAndReplace(newExpr, replace)
 		newProj.Exprs[i] = plannercore.ReplaceColumnOfExpr(newExpr, child, childGroup.Prop.Schema)
 	}
 

pkg/planner/core/BUILD.bazel

@@ -147,6 +147,7 @@ go_library(
         "//pkg/planner/core/metrics",
         "//pkg/planner/core/operator/baseimpl",
         "//pkg/planner/core/operator/logicalop",
+        "//pkg/planner/core/rule/util",
         "//pkg/planner/funcdep",
         "//pkg/planner/property",
         "//pkg/planner/util",

pkg/planner/core/logical_aggregation.go

@@ -25,6 +25,7 @@ import (
 	"github.com/pingcap/tidb/pkg/planner/cardinality"
 	"github.com/pingcap/tidb/pkg/planner/core/base"
 	"github.com/pingcap/tidb/pkg/planner/core/operator/logicalop"
+	ruleutil "github.com/pingcap/tidb/pkg/planner/core/rule/util"
 	fd "github.com/pingcap/tidb/pkg/planner/funcdep"
 	"github.com/pingcap/tidb/pkg/planner/property"
 	"github.com/pingcap/tidb/pkg/planner/util"
@@ -86,14 +87,14 @@ func (la *LogicalAggregation) ExplainInfo() string {
 func (la *LogicalAggregation) ReplaceExprColumns(replace map[string]*expression.Column) {
 	for _, agg := range la.AggFuncs {
 		for _, aggExpr := range agg.Args {
-			ResolveExprAndReplace(aggExpr, replace)
+			ruleutil.ResolveExprAndReplace(aggExpr, replace)
 		}
 		for _, orderExpr := range agg.OrderByItems {
-			ResolveExprAndReplace(orderExpr.Expr, replace)
+			ruleutil.ResolveExprAndReplace(orderExpr.Expr, replace)
 		}
 	}
 	for _, gbyItem := range la.GroupByItems {
-		ResolveExprAndReplace(gbyItem, replace)
+		ruleutil.ResolveExprAndReplace(gbyItem, replace)
 	}
 }
 

pkg/planner/core/logical_cte.go

@@ -23,6 +23,7 @@ import (
 	"github.com/pingcap/tidb/pkg/planner/cardinality"
 	"github.com/pingcap/tidb/pkg/planner/core/base"
 	"github.com/pingcap/tidb/pkg/planner/core/operator/logicalop"
+	ruleutil "github.com/pingcap/tidb/pkg/planner/core/rule/util"
 	"github.com/pingcap/tidb/pkg/planner/property"
 	"github.com/pingcap/tidb/pkg/planner/util/coreusage"
 	"github.com/pingcap/tidb/pkg/planner/util/optimizetrace"
@@ -131,7 +132,7 @@ func (p *LogicalCTE) PredicatePushDown(predicates []expression.Expression, _ *op
 	newPred := make([]expression.Expression, 0, len(predicates))
 	for i := range pushedPredicates {
 		newPred = append(newPred, pushedPredicates[i].Clone())
-		ResolveExprAndReplace(newPred[i], p.Cte.ColumnMap)
+		ruleutil.ResolveExprAndReplace(newPred[i], p.Cte.ColumnMap)
 	}
 	p.Cte.pushDownPredicates = append(p.Cte.pushDownPredicates, expression.ComposeCNFCondition(p.SCtx().GetExprCtx(), newPred...))
 	return predicates, p.Self()

pkg/planner/core/logical_join.go

@@ -28,6 +28,7 @@ import (
 	"github.com/pingcap/tidb/pkg/planner/core/base"
 	"github.com/pingcap/tidb/pkg/planner/core/cost"
 	"github.com/pingcap/tidb/pkg/planner/core/operator/logicalop"
+	ruleutil "github.com/pingcap/tidb/pkg/planner/core/rule/util"
 	"github.com/pingcap/tidb/pkg/planner/funcdep"
 	"github.com/pingcap/tidb/pkg/planner/property"
 	"github.com/pingcap/tidb/pkg/planner/util"
@@ -177,16 +178,16 @@ func (p *LogicalJoin) ExplainInfo() string {
 // ReplaceExprColumns implements base.LogicalPlan interface.
 func (p *LogicalJoin) ReplaceExprColumns(replace map[string]*expression.Column) {
 	for _, equalExpr := range p.EqualConditions {
-		ResolveExprAndReplace(equalExpr, replace)
+		ruleutil.ResolveExprAndReplace(equalExpr, replace)
 	}
 	for _, leftExpr := range p.LeftConditions {
-		ResolveExprAndReplace(leftExpr, replace)
+		ruleutil.ResolveExprAndReplace(leftExpr, replace)
 	}
 	for _, rightExpr := range p.RightConditions {
-		ResolveExprAndReplace(rightExpr, replace)
+		ruleutil.ResolveExprAndReplace(rightExpr, replace)
 	}
 	for _, otherExpr := range p.OtherConditions {
-		ResolveExprAndReplace(otherExpr, replace)
+		ruleutil.ResolveExprAndReplace(otherExpr, replace)
 	}
 }
 

pkg/planner/core/logical_projection.go

@@ -22,6 +22,7 @@ import (
 	"github.com/pingcap/tidb/pkg/planner/cardinality"
 	"github.com/pingcap/tidb/pkg/planner/core/base"
 	"github.com/pingcap/tidb/pkg/planner/core/operator/logicalop"
+	ruleutil "github.com/pingcap/tidb/pkg/planner/core/rule/util"
 	fd "github.com/pingcap/tidb/pkg/planner/funcdep"
 	"github.com/pingcap/tidb/pkg/planner/property"
 	"github.com/pingcap/tidb/pkg/planner/util"
@@ -73,7 +74,7 @@ func (p *LogicalProjection) ExplainInfo() string {
 // ReplaceExprColumns implements base.LogicalPlan interface.
 func (p *LogicalProjection) ReplaceExprColumns(replace map[string]*expression.Column) {
 	for _, expr := range p.Exprs {
-		ResolveExprAndReplace(expr, replace)
+		ruleutil.ResolveExprAndReplace(expr, replace)
 	}
 }
 
@@ -250,7 +251,7 @@ func (p *LogicalProjection) PullUpConstantPredicates() []expression.Expression {
 			continue
 		}
 		clonePredicate := predicate.Clone()
-		ResolveExprAndReplace(clonePredicate, replace)
+		ruleutil.ResolveExprAndReplace(clonePredicate, replace)
 		result = append(result, clonePredicate)
 	}
 	return result

pkg/planner/core/logical_selection.go

@@ -24,6 +24,7 @@ import (
 	"github.com/pingcap/tidb/pkg/planner/core/base"
 	"github.com/pingcap/tidb/pkg/planner/core/cost"
 	"github.com/pingcap/tidb/pkg/planner/core/operator/logicalop"
+	ruleutil "github.com/pingcap/tidb/pkg/planner/core/rule/util"
 	fd "github.com/pingcap/tidb/pkg/planner/funcdep"
 	"github.com/pingcap/tidb/pkg/planner/property"
 	"github.com/pingcap/tidb/pkg/planner/util"
@@ -58,7 +59,7 @@ func (p *LogicalSelection) ExplainInfo() string {
 // ReplaceExprColumns implements base.LogicalPlan interface.
 func (p *LogicalSelection) ReplaceExprColumns(replace map[string]*expression.Column) {
 	for _, expr := range p.Conditions {
-		ResolveExprAndReplace(expr, replace)
+		ruleutil.ResolveExprAndReplace(expr, replace)
 	}
 }
 

pkg/planner/core/logical_sort.go

@@ -21,6 +21,7 @@ import (
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/planner/core/base"
 	"github.com/pingcap/tidb/pkg/planner/core/operator/logicalop"
+	ruleutil "github.com/pingcap/tidb/pkg/planner/core/rule/util"
 	"github.com/pingcap/tidb/pkg/planner/property"
 	"github.com/pingcap/tidb/pkg/planner/util"
 	"github.com/pingcap/tidb/pkg/planner/util/optimizetrace"
@@ -52,7 +53,7 @@ func (ls *LogicalSort) ExplainInfo() string {
 // ReplaceExprColumns implements base.LogicalPlan interface.
 func (ls *LogicalSort) ReplaceExprColumns(replace map[string]*expression.Column) {
 	for _, byItem := range ls.ByItems {
-		ResolveExprAndReplace(byItem.Expr, replace)
+		ruleutil.ResolveExprAndReplace(byItem.Expr, replace)
 	}
 }
 

pkg/planner/core/logical_tikv_single_gather.go

@@ -21,7 +21,6 @@ import (
 	"github.com/pingcap/tidb/pkg/parser/model"
 	"github.com/pingcap/tidb/pkg/planner/core/base"
 	"github.com/pingcap/tidb/pkg/planner/core/operator/logicalop"
-	"github.com/pingcap/tidb/pkg/planner/property"
 	"github.com/pingcap/tidb/pkg/util/plancodec"
 )
 
@@ -115,27 +114,3 @@ func (*TiKVSingleGather) PreparePossibleProperties(_ *expression.Schema, childre
 // ConvertOuterToInnerJoin inherits BaseLogicalPlan.LogicalPlan.<24th> implementation.
 
 // *************************** end implementation of logicalPlan interface ***************************
-
-// GetPhysicalIndexReader returns PhysicalIndexReader for logical TiKVSingleGather.
-func (sg *TiKVSingleGather) GetPhysicalIndexReader(schema *expression.Schema, stats *property.StatsInfo, props ...*property.PhysicalProperty) *PhysicalIndexReader {
-	reader := PhysicalIndexReader{}.Init(sg.SCtx(), sg.QueryBlockOffset())
-	reader.SetStats(stats)
-	reader.SetSchema(schema)
-	reader.childrenReqProps = props
-	return reader
-}
-
-// GetPhysicalTableReader returns PhysicalTableReader for logical TiKVSingleGather.
-func (sg *TiKVSingleGather) GetPhysicalTableReader(schema *expression.Schema, stats *property.StatsInfo, props ...*property.PhysicalProperty) *PhysicalTableReader {
-	reader := PhysicalTableReader{}.Init(sg.SCtx(), sg.QueryBlockOffset())
-	reader.PlanPartInfo = &PhysPlanPartInfo{
-		PruningConds:   sg.Source.AllConds,
-		PartitionNames: sg.Source.PartitionNames,
-		Columns:        sg.Source.TblCols,
-		ColumnNames:    sg.Source.OutputNames(),
-	}
-	reader.SetStats(stats)
-	reader.SetSchema(schema)
-	reader.childrenReqProps = props
-	return reader
-}

pkg/planner/core/logical_top_n.go

@@ -21,6 +21,7 @@ import (
 	"github.com/pingcap/tidb/pkg/expression"
 	"github.com/pingcap/tidb/pkg/planner/core/base"
 	"github.com/pingcap/tidb/pkg/planner/core/operator/logicalop"
+	ruleutil "github.com/pingcap/tidb/pkg/planner/core/rule/util"
 	"github.com/pingcap/tidb/pkg/planner/property"
 	"github.com/pingcap/tidb/pkg/planner/util"
 	"github.com/pingcap/tidb/pkg/planner/util/optimizetrace"
@@ -63,7 +64,7 @@ func (lt *LogicalTopN) ExplainInfo() string {
 // ReplaceExprColumns implements base.LogicalPlan interface.
 func (lt *LogicalTopN) ReplaceExprColumns(replace map[string]*expression.Column) {
 	for _, byItem := range lt.ByItems {
-		ResolveExprAndReplace(byItem.Expr, replace)
+		ruleutil.ResolveExprAndReplace(byItem.Expr, replace)
 	}
 }
 

pkg/planner/core/logical_window.go

@@ -20,6 +20,7 @@ import (
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/planner/core/base"
 	"github.com/pingcap/tidb/pkg/planner/core/operator/logicalop"
+	ruleutil "github.com/pingcap/tidb/pkg/planner/core/rule/util"
 	"github.com/pingcap/tidb/pkg/planner/property"
 	"github.com/pingcap/tidb/pkg/planner/util/optimizetrace"
 	"github.com/pingcap/tidb/pkg/sessionctx"
@@ -152,14 +153,14 @@ func (p LogicalWindow) Init(ctx base.PlanContext, offset int) *LogicalWindow {
 func (p *LogicalWindow) ReplaceExprColumns(replace map[string]*expression.Column) {
 	for _, desc := range p.WindowFuncDescs {
 		for _, arg := range desc.Args {
-			ResolveExprAndReplace(arg, replace)
+			ruleutil.ResolveExprAndReplace(arg, replace)
 		}
 	}
 	for _, item := range p.PartitionBy {
-		resolveColumnAndReplace(item.Col, replace)
+		ruleutil.ResolveColumnAndReplace(item.Col, replace)
 	}
 	for _, item := range p.OrderBy {
-		resolveColumnAndReplace(item.Col, replace)
+		ruleutil.ResolveColumnAndReplace(item.Col, replace)
 	}
 }
 

pkg/planner/core/physical_plans.go

@@ -248,6 +248,30 @@ func setMppOrBatchCopForTableScan(curPlan base.PhysicalPlan) {
 	}
 }
 
+// GetPhysicalIndexReader returns PhysicalIndexReader for logical TiKVSingleGather.
+func GetPhysicalIndexReader(sg *TiKVSingleGather, schema *expression.Schema, stats *property.StatsInfo, props ...*property.PhysicalProperty) *PhysicalIndexReader {
+	reader := PhysicalIndexReader{}.Init(sg.SCtx(), sg.QueryBlockOffset())
+	reader.SetStats(stats)
+	reader.SetSchema(schema)
+	reader.childrenReqProps = props
+	return reader
+}
+
+// GetPhysicalTableReader returns PhysicalTableReader for logical TiKVSingleGather.
+func GetPhysicalTableReader(sg *TiKVSingleGather, schema *expression.Schema, stats *property.StatsInfo, props ...*property.PhysicalProperty) *PhysicalTableReader {
+	reader := PhysicalTableReader{}.Init(sg.SCtx(), sg.QueryBlockOffset())
+	reader.PlanPartInfo = &PhysPlanPartInfo{
+		PruningConds:   sg.Source.AllConds,
+		PartitionNames: sg.Source.PartitionNames,
+		Columns:        sg.Source.TblCols,
+		ColumnNames:    sg.Source.OutputNames(),
+	}
+	reader.SetStats(stats)
+	reader.SetSchema(schema)
+	reader.childrenReqProps = props
+	return reader
+}
+
 // Clone implements op.PhysicalPlan interface.
 func (p *PhysicalTableReader) Clone(newCtx base.PlanContext) (base.PhysicalPlan, error) {
 	cloned := new(PhysicalTableReader)

pkg/planner/core/rule/util/BUILD.bazel

@@ -0,0 +1,9 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "util",
+    srcs = ["misc.go"],
+    importpath = "github.com/pingcap/tidb/pkg/planner/core/rule/util",
+    visibility = ["//visibility:public"],
+    deps = ["//pkg/expression"],
+)

pkg/planner/core/rule/util/misc.go

@@ -0,0 +1,41 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package util
+
+import "github.com/pingcap/tidb/pkg/expression"
+
+// ResolveExprAndReplace replaces columns fields of expressions by children logical plans.
+func ResolveExprAndReplace(origin expression.Expression, replace map[string]*expression.Column) {
+	switch expr := origin.(type) {
+	case *expression.Column:
+		ResolveColumnAndReplace(expr, replace)
+	case *expression.CorrelatedColumn:
+		ResolveColumnAndReplace(&expr.Column, replace)
+	case *expression.ScalarFunction:
+		for _, arg := range expr.GetArgs() {
+			ResolveExprAndReplace(arg, replace)
+		}
+	}
+}
+
+// ResolveColumnAndReplace replaces columns fields of expressions by children logical plans.
+func ResolveColumnAndReplace(origin *expression.Column, replace map[string]*expression.Column) {
+	dst := replace[string(origin.HashCode())]
+	if dst != nil {
+		retType, inOperand := origin.RetType, origin.InOperand
+		*origin = *dst
+		origin.RetType, origin.InOperand = retType, inOperand
+	}
+}

pkg/planner/core/rule_eliminate_projection.go

@@ -25,6 +25,7 @@ import (
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/planner/core/base"
+	ruleutil "github.com/pingcap/tidb/pkg/planner/core/rule/util"
 	"github.com/pingcap/tidb/pkg/planner/util/optimizetrace"
 )
 
@@ -92,29 +93,6 @@ func canProjectionBeEliminatedStrict(p *PhysicalProjection) bool {
 	return true
 }
 
-func resolveColumnAndReplace(origin *expression.Column, replace map[string]*expression.Column) {
-	dst := replace[string(origin.HashCode())]
-	if dst != nil {
-		retType, inOperand := origin.RetType, origin.InOperand
-		*origin = *dst
-		origin.RetType, origin.InOperand = retType, inOperand
-	}
-}
-
-// ResolveExprAndReplace replaces columns fields of expressions by children logical plans.
-func ResolveExprAndReplace(origin expression.Expression, replace map[string]*expression.Column) {
-	switch expr := origin.(type) {
-	case *expression.Column:
-		resolveColumnAndReplace(expr, replace)
-	case *expression.CorrelatedColumn:
-		resolveColumnAndReplace(&expr.Column, replace)
-	case *expression.ScalarFunction:
-		for _, arg := range expr.GetArgs() {
-			ResolveExprAndReplace(arg, replace)
-		}
-	}
-}
-
 func doPhysicalProjectionElimination(p base.PhysicalPlan) base.PhysicalPlan {
 	for i, child := range p.Children() {
 		p.Children()[i] = doPhysicalProjectionElimination(child)
@@ -205,7 +183,7 @@ func (pe *projectionEliminator) eliminate(p base.LogicalPlan, replace map[string
 		x.SetSchema(buildLogicalJoinSchema(x.JoinType, x))
 	default:
 		for _, dst := range p.Schema().Columns {
-			resolveColumnAndReplace(dst, replace)
+			ruleutil.ResolveColumnAndReplace(dst, replace)
 		}
 	}
 	// replace all of exprs in logical plan