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",
@@ -59,7 +58,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_decorrelate.go",
         "rule_derive_topn_from_window.go",

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

@@ -82,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)
@@ -94,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/operator/logicalop/logical_cte.go

@@ -100,6 +100,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

@@ -86,12 +86,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

@@ -6,14 +6,31 @@ go_library(
         "logical_rules.go",
         "rule_build_key_info.go",
         "rule_constant_propagation.go",
+        "collect_column_stats_usage.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
@@ -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:
@@ -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"