planner: move rule_collect_plan_stats into rules package

Signed-off-by: Weizhen Wang <[email protected]>

pkg/planner/core/BUILD.bazel

@@ -4,7 +4,6 @@ go_library(
     name = "core",
     srcs = [
         "access_object.go",
-        "collect_column_stats_usage.go",
         "common_plans.go",
         "core_init.go",
         "debugtrace.go",
@@ -68,7 +67,6 @@ go_library(
         "rule_aggregation_elimination.go",
         "rule_aggregation_push_down.go",
         "rule_aggregation_skew_rewrite.go",
-        "rule_collect_plan_stats.go",
         "rule_column_pruning.go",
         "rule_constant_propagation.go",
         "rule_decorrelate.go",
@@ -291,6 +289,7 @@ go_test(
         "//pkg/planner",
         "//pkg/planner/core/base",
         "//pkg/planner/core/operator/logicalop",
+        "//pkg/planner/core/rule",
         "//pkg/planner/property",
         "//pkg/planner/util",
         "//pkg/planner/util/coretestsdk",

pkg/planner/core/casetest/planstats/BUILD.bazel

@@ -19,6 +19,7 @@ go_test(
         "//pkg/planner",
         "//pkg/planner/core",
         "//pkg/planner/core/base",
+        "//pkg/planner/core/rule",
         "//pkg/sessionctx",
         "//pkg/sessionctx/stmtctx",
         "//pkg/statistics",

pkg/planner/core/casetest/planstats/plan_stats_test.go

@@ -30,6 +30,7 @@ import (
 	"github.com/pingcap/tidb/pkg/planner"
 	plannercore "github.com/pingcap/tidb/pkg/planner/core"
 	"github.com/pingcap/tidb/pkg/planner/core/base"
+	"github.com/pingcap/tidb/pkg/planner/core/rule"
 	"github.com/pingcap/tidb/pkg/sessionctx"
 	"github.com/pingcap/tidb/pkg/sessionctx/stmtctx"
 	"github.com/pingcap/tidb/pkg/statistics"
@@ -392,7 +393,7 @@ func TestCollectDependingVirtualCols(t *testing.T) {
 		}
 
 		// call the function
-		res := plannercore.CollectDependingVirtualCols(tblID2Tbl, neededItems)
+		res := rule.CollectDependingVirtualCols(tblID2Tbl, neededItems)
 
 		// record and check the output
 		cols := make([]string, 0, len(res))

pkg/planner/core/collect_column_stats_usage_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/pingcap/tidb/pkg/infoschema"
 	"github.com/pingcap/tidb/pkg/parser/model"
 	"github.com/pingcap/tidb/pkg/planner/core/base"
+	"github.com/pingcap/tidb/pkg/planner/core/rule"
 	"github.com/pingcap/tidb/pkg/util/hint"
 	"github.com/stretchr/testify/require"
 )
@@ -81,7 +82,7 @@ func getStatsLoadItem(t *testing.T, is infoschema.InfoSchema, item model.StatsLo
 
 func checkColumnStatsUsageForPredicates(t *testing.T, is infoschema.InfoSchema, lp base.LogicalPlan, expected []string, comment string) {
 	var tblColIDs []model.TableItemID
-	tblColIDs, _, _ = CollectColumnStatsUsage(lp, false)
+	tblColIDs, _, _ = rule.CollectColumnStatsUsage(lp, false)
 	cols := make([]string, 0, len(tblColIDs))
 	for _, tblColID := range tblColIDs {
 		col := getColumnName(t, is, tblColID, comment)
@@ -93,7 +94,7 @@ func checkColumnStatsUsageForPredicates(t *testing.T, is infoschema.InfoSchema,
 
 func checkColumnStatsUsageForStatsLoad(t *testing.T, is infoschema.InfoSchema, lp base.LogicalPlan, expected []string, comment string) {
 	var loadItems []model.StatsLoadItem
-	_, loadItems, _ = CollectColumnStatsUsage(lp, true)
+	_, loadItems, _ = rule.CollectColumnStatsUsage(lp, true)
 	cols := make([]string, 0, len(loadItems))
 	for _, item := range loadItems {
 		col := getStatsLoadItem(t, is, item, comment)

pkg/planner/core/logical_cte.go

@@ -99,6 +99,16 @@ func (cc *CTEClass) MemoryUsage() (sum int64) {
 	return
 }
 
+// SeedPartLogicalPlan is to get the seed part logical plan.
+func (cc *CTEClass) SeedPartLogicalPlan() base.LogicalPlan {
+	return cc.seedPartLogicalPlan
+}
+
+// RecursivePartLogicalPlan is to get the recursive part logical plan.
+func (cc *CTEClass) RecursivePartLogicalPlan() base.LogicalPlan {
+	return cc.recursivePartLogicalPlan
+}
+
 // *************************** start implementation of logicalPlan interface ***************************
 
 // HashCode inherits the BaseLogicalPlan.<0th> implementation.

pkg/planner/core/optimizer.go

@@ -113,12 +113,12 @@ var optRuleList = []base.LogicalOptRule{
 	&PPDSolver{},
 	&OuterJoinEliminator{},
 	&PartitionProcessor{},
-	&CollectPredicateColumnsPoint{},
+	&rule.CollectPredicateColumnsPoint{},
 	&AggregationPushDownSolver{},
 	&DeriveTopNFromWindow{},
 	&PredicateSimplification{},
 	&PushDownTopNOptimizer{},
-	&SyncWaitStatsLoadPoint{},
+	&rule.SyncWaitStatsLoadPoint{},
 	&JoinReOrderSolver{},
 	&ColumnPruner{}, // column pruning again at last, note it will mess up the results of buildKeySolver
 	&PushDownSequenceSolver{},

pkg/planner/core/rule/BUILD.bazel

@@ -3,15 +3,32 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library")
 go_library(
     name = "rule",
     srcs = [
+        "collect_column_stats_usage.go",
         "rule_build_key_info.go",
+        "rule_collect_plan_stats.go",
         "rule_init.go",
     ],
     importpath = "github.com/pingcap/tidb/pkg/planner/core/rule",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/domain",
         "//pkg/expression",
+        "//pkg/infoschema",
+        "//pkg/parser/model",
+        "//pkg/planner/core",
         "//pkg/planner/core/base",
+        "//pkg/planner/core/operator/logicalop",
         "//pkg/planner/core/rule/util",
         "//pkg/planner/util/optimizetrace",
+        "//pkg/sessionctx/variable",
+        "//pkg/statistics",
+        "//pkg/statistics/asyncload",
+        "//pkg/table",
+        "//pkg/util/filter",
+        "//pkg/util/intset",
+        "//pkg/util/logutil",
+        "@com_github_pingcap_failpoint//:failpoint",
+        "@org_golang_x_exp//maps",
+        "@org_uber_go_zap//:zap",
     ],
 )

pkg/planner/core/collect_column_stats_usage.go → pkg/planner/core/rule/collect_column_stats_usage.go

@@ -12,14 +12,15 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package core
+package rule
 
 import (
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/pkg/domain"
 	"github.com/pingcap/tidb/pkg/expression"
 	"github.com/pingcap/tidb/pkg/infoschema"
 	"github.com/pingcap/tidb/pkg/parser/model"
+	"github.com/pingcap/tidb/pkg/planner/core"
 	"github.com/pingcap/tidb/pkg/planner/core/base"
 	"github.com/pingcap/tidb/pkg/planner/core/operator/logicalop"
 	"github.com/pingcap/tidb/pkg/sessionctx/variable"
@@ -124,7 +125,7 @@ func (c *columnStatsUsageCollector) updateColMapFromExpressions(col *expression.
 	c.updateColMap(col, expression.ExtractColumnsAndCorColumnsFromExpressions(c.cols[:0], list))
 }
 
-func (c *columnStatsUsageCollector) collectPredicateColumnsForDataSource(ds *DataSource) {
+func (c *columnStatsUsageCollector) collectPredicateColumnsForDataSource(ds *core.DataSource) {
 	// Skip all system tables.
 	if filter.IsSystemSchema(ds.DBName.L) {
 		return
@@ -143,7 +144,7 @@ func (c *columnStatsUsageCollector) collectPredicateColumnsForDataSource(ds *Dat
 	c.addPredicateColumnsFromExpressions(ds.PushedDownConds)
 }
 
-func (c *columnStatsUsageCollector) collectPredicateColumnsForJoin(p *LogicalJoin) {
+func (c *columnStatsUsageCollector) collectPredicateColumnsForJoin(p *core.LogicalJoin) {
 	// The only schema change is merging two schemas so there is no new column.
 	// Assume statistics of all the columns in EqualConditions/LeftConditions/RightConditions/OtherConditions are needed.
 	exprs := make([]expression.Expression, 0, len(p.EqualConditions)+len(p.LeftConditions)+len(p.RightConditions)+len(p.OtherConditions))
@@ -162,7 +163,7 @@ func (c *columnStatsUsageCollector) collectPredicateColumnsForJoin(p *LogicalJoi
 	c.addPredicateColumnsFromExpressions(exprs)
 }
 
-func (c *columnStatsUsageCollector) collectPredicateColumnsForUnionAll(p *LogicalUnionAll) {
+func (c *columnStatsUsageCollector) collectPredicateColumnsForUnionAll(p *core.LogicalUnionAll) {
 	// statistics of the ith column of UnionAll come from statistics of the ith column of each child.
 	schemas := make([]*expression.Schema, 0, len(p.Children()))
 	relatedCols := make([]*expression.Column, 0, len(p.Children()))
@@ -178,7 +179,7 @@ func (c *columnStatsUsageCollector) collectPredicateColumnsForUnionAll(p *Logica
 	}
 }
 
-func (c *columnStatsUsageCollector) addHistNeededColumns(ds *DataSource) {
+func (c *columnStatsUsageCollector) addHistNeededColumns(ds *core.DataSource) {
 	c.visitedPhysTblIDs.Insert(int(ds.PhysicalTableID))
 	if c.collectMode&collectHistNeededColumns == 0 {
 		return
@@ -232,12 +233,12 @@ func (c *columnStatsUsageCollector) collectFromPlan(lp base.LogicalPlan) {
 	}
 	if c.collectMode&collectPredicateColumns != 0 {
 		switch x := lp.(type) {
-		case *DataSource:
+		case *core.DataSource:
 			c.collectPredicateColumnsForDataSource(x)
-		case *LogicalIndexScan:
+		case *core.LogicalIndexScan:
 			c.collectPredicateColumnsForDataSource(x.Source)
 			c.addPredicateColumnsFromExpressions(x.AccessConds)
-		case *LogicalTableScan:
+		case *core.LogicalTableScan:
 			c.collectPredicateColumnsForDataSource(x.Source)
 			c.addPredicateColumnsFromExpressions(x.AccessConds)
 		case *logicalop.LogicalProjection:
@@ -246,11 +247,11 @@ func (c *columnStatsUsageCollector) collectFromPlan(lp base.LogicalPlan) {
 			for i, expr := range x.Exprs {
 				c.updateColMapFromExpressions(schema.Columns[i], []expression.Expression{expr})
 			}
-		case *LogicalSelection:
+		case *core.LogicalSelection:
 			// Though the conditions in LogicalSelection are complex conditions which cannot be pushed down to DataSource, we still
 			// regard statistics of the columns in the conditions as needed.
 			c.addPredicateColumnsFromExpressions(x.Conditions)
-		case *LogicalAggregation:
+		case *core.LogicalAggregation:
 			// Just assume statistics of all the columns in GroupByItems are needed.
 			c.addPredicateColumnsFromExpressions(x.GroupByItems)
 			// Schema change from children to self.
@@ -269,9 +270,9 @@ func (c *columnStatsUsageCollector) collectFromPlan(lp base.LogicalPlan) {
 			for i, col := range windowColumns {
 				c.updateColMapFromExpressions(col, x.WindowFuncDescs[i].Args)
 			}
-		case *LogicalJoin:
+		case *core.LogicalJoin:
 			c.collectPredicateColumnsForJoin(x)
-		case *LogicalApply:
+		case *core.LogicalApply:
 			c.collectPredicateColumnsForJoin(&x.LogicalJoin)
 			// Assume statistics of correlated columns are needed.
 			// Correlated columns can be found in LogicalApply.Children()[0].Schema(). Since we already visit LogicalApply.Children()[0],
@@ -289,22 +290,22 @@ func (c *columnStatsUsageCollector) collectFromPlan(lp base.LogicalPlan) {
 			for _, item := range x.ByItems {
 				c.addPredicateColumnsFromExpressions([]expression.Expression{item.Expr})
 			}
-		case *LogicalUnionAll:
+		case *core.LogicalUnionAll:
 			c.collectPredicateColumnsForUnionAll(x)
-		case *LogicalPartitionUnionAll:
+		case *core.LogicalPartitionUnionAll:
 			c.collectPredicateColumnsForUnionAll(&x.LogicalUnionAll)
-		case *LogicalCTE:
+		case *core.LogicalCTE:
 			// Visit seedPartLogicalPlan and recursivePartLogicalPlan first.
-			c.collectFromPlan(x.Cte.seedPartLogicalPlan)
-			if x.Cte.recursivePartLogicalPlan != nil {
-				c.collectFromPlan(x.Cte.recursivePartLogicalPlan)
+			c.collectFromPlan(x.Cte.SeedPartLogicalPlan())
+			if x.Cte.RecursivePartLogicalPlan() != nil {
+				c.collectFromPlan(x.Cte.RecursivePartLogicalPlan())
 			}
 			// Schema change from seedPlan/recursivePlan to self.
 			columns := x.Schema().Columns
-			seedColumns := x.Cte.seedPartLogicalPlan.Schema().Columns
+			seedColumns := x.Cte.SeedPartLogicalPlan().Schema().Columns
 			var recursiveColumns []*expression.Column
-			if x.Cte.recursivePartLogicalPlan != nil {
-				recursiveColumns = x.Cte.recursivePartLogicalPlan.Schema().Columns
+			if x.Cte.RecursivePartLogicalPlan() != nil {
+				recursiveColumns = x.Cte.RecursivePartLogicalPlan().Schema().Columns
 			}
 			relatedCols := make([]*expression.Column, 0, 2)
 			for i, col := range columns {
@@ -336,11 +337,11 @@ func (c *columnStatsUsageCollector) collectFromPlan(lp base.LogicalPlan) {
 	// Since c.visitedPhysTblIDs is also collected here and needs to be collected even collectHistNeededColumns is not set,
 	// so we do the c.collectMode check in addHistNeededColumns() after collecting c.visitedPhysTblIDs.
 	switch x := lp.(type) {
-	case *DataSource:
+	case *core.DataSource:
 		c.addHistNeededColumns(x)
-	case *LogicalIndexScan:
+	case *core.LogicalIndexScan:
 		c.addHistNeededColumns(x.Source)
-	case *LogicalTableScan:
+	case *core.LogicalTableScan:
 		c.addHistNeededColumns(x.Source)
 	}
 }

pkg/planner/core/rule_collect_plan_stats.go → pkg/planner/core/rule/rule_collect_plan_stats.go

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package core
+package rule
 
 import (
 	"context"