From b5678f0eda41d6afe84a5bfaff8f137d9ab9fa80 Mon Sep 17 00:00:00 2001 From: Weizhen Wang Date: Wed, 7 Aug 2024 22:51:02 +0800 Subject: [PATCH] planner: move rule_collect_plan_stats into rules package Signed-off-by: Weizhen Wang --- pkg/planner/core/BUILD.bazel | 3 +- .../core/casetest/planstats/BUILD.bazel | 1 + .../casetest/planstats/plan_stats_test.go | 3 +- .../core/collect_column_stats_usage_test.go | 5 +- pkg/planner/core/logical_cte.go | 10 ++++ pkg/planner/core/optimizer.go | 4 +- pkg/planner/core/rule/BUILD.bazel | 17 +++++++ .../{ => rule}/collect_column_stats_usage.go | 49 ++++++++++--------- .../{ => rule}/rule_collect_plan_stats.go | 2 +- 9 files changed, 62 insertions(+), 32 deletions(-) rename pkg/planner/core/{ => rule}/collect_column_stats_usage.go (94%) rename pkg/planner/core/{ => rule}/rule_collect_plan_stats.go (99%) diff --git a/pkg/planner/core/BUILD.bazel b/pkg/planner/core/BUILD.bazel index a91265ef909cb..bff37e4977882 100644 --- a/pkg/planner/core/BUILD.bazel +++ b/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", diff --git a/pkg/planner/core/casetest/planstats/BUILD.bazel b/pkg/planner/core/casetest/planstats/BUILD.bazel index 8d738edc3b502..b517bf7ed99ba 100644 --- a/pkg/planner/core/casetest/planstats/BUILD.bazel +++ b/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", diff --git a/pkg/planner/core/casetest/planstats/plan_stats_test.go b/pkg/planner/core/casetest/planstats/plan_stats_test.go index a7e8367fbb909..94890a6d25be2 100644 --- a/pkg/planner/core/casetest/planstats/plan_stats_test.go +++ b/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)) diff --git a/pkg/planner/core/collect_column_stats_usage_test.go b/pkg/planner/core/collect_column_stats_usage_test.go index 11aad00b3aa8e..a01735d7f7cd5 100644 --- a/pkg/planner/core/collect_column_stats_usage_test.go +++ b/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) diff --git a/pkg/planner/core/logical_cte.go b/pkg/planner/core/logical_cte.go index 9888eccf8e6be..3f6f9e960e88b 100644 --- a/pkg/planner/core/logical_cte.go +++ b/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. diff --git a/pkg/planner/core/optimizer.go b/pkg/planner/core/optimizer.go index ebdd4706dcd18..8cf730f9ab8f6 100644 --- a/pkg/planner/core/optimizer.go +++ b/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{}, diff --git a/pkg/planner/core/rule/BUILD.bazel b/pkg/planner/core/rule/BUILD.bazel index 7c2175db74f8e..09ac8a18442e7 100644 --- a/pkg/planner/core/rule/BUILD.bazel +++ b/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", ], ) diff --git a/pkg/planner/core/collect_column_stats_usage.go b/pkg/planner/core/rule/collect_column_stats_usage.go similarity index 94% rename from pkg/planner/core/collect_column_stats_usage.go rename to pkg/planner/core/rule/collect_column_stats_usage.go index 54c427c82289f..005f88d89acda 100644 --- a/pkg/planner/core/collect_column_stats_usage.go +++ b/pkg/planner/core/rule/collect_column_stats_usage.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package core +package rule import ( "github.com/pingcap/failpoint" @@ -20,6 +20,7 @@ import ( "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) } } diff --git a/pkg/planner/core/rule_collect_plan_stats.go b/pkg/planner/core/rule/rule_collect_plan_stats.go similarity index 99% rename from pkg/planner/core/rule_collect_plan_stats.go rename to pkg/planner/core/rule/rule_collect_plan_stats.go index e9047782c5095..9638e8e51933b 100644 --- a/pkg/planner/core/rule_collect_plan_stats.go +++ b/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"