From 1348486b61804520021ef7c811d3f5163a52f6e6 Mon Sep 17 00:00:00 2001 From: Reminiscent Date: Thu, 19 Mar 2026 11:51:29 +0800 Subject: [PATCH 1/2] planner, statistics: use selected partition stats for dynamic pruning --- pkg/planner/cardinality/BUILD.bazel | 1 + pkg/planner/cardinality/pseudo.go | 12 + .../casetest/instanceplancache/others_test.go | 52 ++-- .../casetest/planstats/plan_stats_test.go | 72 +++++ .../core/collect_column_stats_usage.go | 4 +- pkg/planner/core/logical_plan_builder.go | 47 +++- .../operator/logicalop/logical_datasource.go | 4 +- pkg/planner/core/plan_cache.go | 171 ++++++++++++ .../core/plan_cache_partition_table_test.go | 257 +++++++++++++----- pkg/planner/core/plan_cache_utils.go | 11 +- pkg/planner/core/planbuilder.go | 7 +- pkg/planner/core/rule_partition_processor.go | 76 ++++-- pkg/planner/core/stats.go | 4 +- .../core/tests/prepare/prepare_test.go | 72 ++++- pkg/sessionctx/stmtctx/stmtctx.go | 4 + pkg/sessionctx/variable/session.go | 4 + pkg/sessionctx/variable/setvar_affect.go | 1 + pkg/sessionctx/variable/sysvar.go | 4 + pkg/sessionctx/variable/tidb_vars.go | 3 + 19 files changed, 677 insertions(+), 129 deletions(-) diff --git a/pkg/planner/cardinality/BUILD.bazel b/pkg/planner/cardinality/BUILD.bazel index 76e62a58e5c5f..5ad2e29edf05a 100644 --- a/pkg/planner/cardinality/BUILD.bazel +++ b/pkg/planner/cardinality/BUILD.bazel @@ -51,6 +51,7 @@ go_test( timeout = "short", srcs = [ "main_test.go", + "ndv_test.go", "row_count_test.go", "row_size_test.go", "selectivity_test.go", diff --git a/pkg/planner/cardinality/pseudo.go b/pkg/planner/cardinality/pseudo.go index 1f13a0862fc0f..5f176d338a951 100644 --- a/pkg/planner/cardinality/pseudo.go +++ b/pkg/planner/cardinality/pseudo.go @@ -41,6 +41,18 @@ func PseudoAvgCountPerValue(t *statistics.Table) float64 { return float64(t.RealtimeCount) / pseudoEqualRate } +// AggregateSelectedPartitionCounts aggregates the realtime and modify count from selected partitions. +func AggregateSelectedPartitionCounts(partitionStats []*statistics.Table) (realtimeCount, modifyCount int64, ok bool) { + for _, partStats := range partitionStats { + if partStats == nil || partStats.Pseudo { + return 0, 0, false + } + realtimeCount += partStats.RealtimeCount + modifyCount += partStats.ModifyCount + } + return realtimeCount, modifyCount, true +} + func pseudoSelectivity(sctx planctx.PlanContext, coll *statistics.HistColl, exprs []expression.Expression) float64 { minFactor := selectionFactor colExists := make(map[string]bool) diff --git a/pkg/planner/core/casetest/instanceplancache/others_test.go b/pkg/planner/core/casetest/instanceplancache/others_test.go index 41255f04d3116..db05c472954a8 100644 --- a/pkg/planner/core/casetest/instanceplancache/others_test.go +++ b/pkg/planner/core/casetest/instanceplancache/others_test.go @@ -366,23 +366,41 @@ func TestInstancePlanCachePartitioning(t *testing.T) { tk := testkit.NewTestKit(t, store) tk.MustExec("use test") tk.MustExec(`set global tidb_enable_instance_plan_cache=1`) - tk.MustExec(`set @@tidb_partition_prune_mode='dynamic'`) - - tk.MustExec(`create table t (a int, b varchar(255)) partition by hash(a) partitions 3`) - tk.MustExec(`insert into t values (1,"a"),(2,"b"),(3,"c"),(4,"d"),(5,"e"),(6,"f")`) - tk.MustExec(`prepare stmt from 'select a,b from t where a = ?;'`) - tk.MustExec(`set @a=1`) - tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("1 a")) - // Same partition works, due to pruning is not affected - tk.MustExec(`set @a=4`) - tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("4 d")) - tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1")) - - tk.MustExec(`set @@tidb_partition_prune_mode='static'`) - tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("4 d")) - tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("0")) - tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("4 d")) - tk.MustQuery(`show warnings`).Check(testkit.Rows("Warning 1105 skip prepared plan-cache: Static partition pruning mode")) + for _, selectedPartitionStats := range []string{"0", "1"} { + tk.MustExec(`admin flush instance plan_cache`) + tk.MustExec(`set @@tidb_opt_enable_selected_partition_stats=` + selectedPartitionStats) + tk.MustExec(`set @@tidb_partition_prune_mode='dynamic'`) + tk.MustExec(`drop table if exists t`) + tk.MustExec(`create table t (a int, b varchar(255)) partition by hash(a) partitions 3`) + tk.MustExec(`insert into t values (1,"a"),(2,"b"),(3,"c"),(4,"d"),(5,"e"),(6,"f")`) + tk.MustExec(`prepare stmt from 'select a,b from t where a = ?;'`) + tk.MustExec(`set @a=1`) + tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("1 a")) + // Result correctness stays the same; whether the second execution hits cache depends on the selected partition stats switch. + tk.MustExec(`set @a=4`) + tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("4 d")) + expectedFromCache := "1" + if selectedPartitionStats == "1" { + expectedFromCache = "0" + } + tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows(expectedFromCache)) + tk.MustQuery(`show warnings`).Check(testkit.Rows()) + + tk.MustExec(`set @@tidb_partition_prune_mode='static'`) + tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("4 d")) + tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("0")) + tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("4 d")) + warnings := tk.MustQuery(`show warnings`).Rows() + if len(warnings) > 0 { + require.True(t, + strings.Contains(warnings[0][2].(string), "skip prepared plan-cache: Static partition pruning mode") || + strings.Contains(warnings[0][2].(string), "skip prepared plan-cache: static partition prune mode used"), + "unexpected warning: %s", warnings[0][2].(string), + ) + } + tk.MustExec(`deallocate prepare stmt`) + tk.MustExec(`drop table t`) + } } func TestInstancePlanCachePlan(t *testing.T) { diff --git a/pkg/planner/core/casetest/planstats/plan_stats_test.go b/pkg/planner/core/casetest/planstats/plan_stats_test.go index 7ad36d8096b95..2179dd2892b8c 100644 --- a/pkg/planner/core/casetest/planstats/plan_stats_test.go +++ b/pkg/planner/core/casetest/planstats/plan_stats_test.go @@ -571,3 +571,75 @@ func TestPartialStatsInExplain(t *testing.T) { require.NoError(t, dom.StatsHandle().LoadNeededHistograms(dom.InfoSchema())) } } + +func TestDynamicPartitionPruneUsesMergedPartitionStats(t *testing.T) { + p := parser.New() + store, dom := testkit.CreateMockStoreAndDomain(t) + + testKit := testkit.NewTestKit(t, store) + ctx := testKit.Session().(sessionctx.Context) + + testKit.MustExec("use test") + testKit.MustExec("drop table if exists pt") + testKit.MustExec("set @@session.tidb_partition_prune_mode = 'dynamic'") + testKit.MustExec("set @@session.tidb_analyze_version = 2") + testKit.MustExec("set @@session.tidb_stats_load_sync_wait = 60000") + testKit.MustExec( + "create table pt(a int, b int) partition by range(a) (" + + "partition p0 values less than (100)," + + "partition p1 values less than (200)," + + "partition p2 values less than maxvalue)", + ) + + rows := make([]string, 0, 1200) + for i := 0; i < 1200; i++ { + rows = append(rows, fmt.Sprintf("(%d,%d)", i, i)) + } + const batchSize = 200 + for start := 0; start < len(rows); start += batchSize { + end := start + batchSize + if end > len(rows) { + end = len(rows) + } + testKit.MustExec("insert into pt values " + strings.Join(rows[start:end], ",")) + } + + oriLease := dom.StatsHandle().Lease() + dom.StatsHandle().SetLease(1) + defer func() { + dom.StatsHandle().SetLease(oriLease) + }() + testKit.MustExec("analyze table pt all columns") + dom.StatsHandle().Clear() + require.NoError(t, dom.StatsHandle().Update(context.Background(), dom.InfoSchema())) + + tbl, err := dom.InfoSchema().TableByName(context.Background(), pmodel.NewCIStr("test"), pmodel.NewCIStr("pt")) + require.NoError(t, err) + globalStats := dom.StatsHandle().GetPhysicalTableStats(tbl.Meta().ID, tbl.Meta()) + require.Equal(t, int64(1200), globalStats.RealtimeCount) + + for _, tc := range []struct { + sql string + expectedRowCount int64 + }{ + { + sql: "select /*+ set_var(tidb_opt_enable_selected_partition_stats=0) */ * from pt where a < 200", + expectedRowCount: 1200, + }, + { + sql: "select /*+ set_var(tidb_opt_enable_selected_partition_stats=1) */ * from pt where a < 200", + expectedRowCount: 200, + }, + } { + stmt, err := p.ParseOneStmt(tc.sql, "", "") + require.NoError(t, err) + require.NoError(t, executor.ResetContextOfStmt(ctx, stmt)) + nodeW := resolve.NewNodeW(stmt) + plan, _, err := planner.Optimize(context.Background(), ctx, nodeW, dom.InfoSchema()) + require.NoError(t, err) + + reader, ok := plan.(*plannercore.PhysicalTableReader) + require.True(t, ok) + require.Equal(t, tc.expectedRowCount, reader.StatsInfo().HistColl.RealtimeCount) + } +} diff --git a/pkg/planner/core/collect_column_stats_usage.go b/pkg/planner/core/collect_column_stats_usage.go index e58301a8eb666..9a924b4da7462 100644 --- a/pkg/planner/core/collect_column_stats_usage.go +++ b/pkg/planner/core/collect_column_stats_usage.go @@ -135,7 +135,9 @@ func (c *columnStatsUsageCollector) collectPredicateColumnsForDataSource(ds *log c.visitedtbls[tblID] = struct{}{} } c.visitedPhysTblIDs.Insert(int(tblID)) - if tblID != ds.PhysicalTableID { + if len(ds.StaticPrunedPartitionIDs) > 0 { + c.tblID2PartitionIDs[tblID] = append(c.tblID2PartitionIDs[tblID], ds.StaticPrunedPartitionIDs...) + } else if tblID != ds.PhysicalTableID { c.tblID2PartitionIDs[tblID] = append(c.tblID2PartitionIDs[tblID], ds.PhysicalTableID) } for _, col := range ds.Schema().Columns { diff --git a/pkg/planner/core/logical_plan_builder.go b/pkg/planner/core/logical_plan_builder.go index 97c3ff4de3094..8554a53c49258 100644 --- a/pkg/planner/core/logical_plan_builder.go +++ b/pkg/planner/core/logical_plan_builder.go @@ -19,6 +19,7 @@ import ( "fmt" "math" "math/bits" + "slices" "sort" "strconv" "strings" @@ -45,6 +46,7 @@ import ( "github.com/pingcap/tidb/pkg/parser/mysql" "github.com/pingcap/tidb/pkg/parser/opcode" "github.com/pingcap/tidb/pkg/parser/terror" + "github.com/pingcap/tidb/pkg/planner/cardinality" "github.com/pingcap/tidb/pkg/planner/core/base" core_metrics "github.com/pingcap/tidb/pkg/planner/core/metrics" "github.com/pingcap/tidb/pkg/planner/core/operator/logicalop" @@ -4178,10 +4180,11 @@ func addExtraPhysTblIDColumn4DS(ds *logicalop.DataSource) *expression.Column { // 2. table row count from statistics is zero. // 3. statistics is outdated. // Note: please also update getLatestVersionFromStatsTable() when logic in this function changes. -func getStatsTable(ctx base.PlanContext, tblInfo *model.TableInfo, pid int64) *statistics.Table { +func getStatsTable(ctx base.PlanContext, tblInfo *model.TableInfo, pid int64, partitionIDs []int64) *statistics.Table { statsHandle := domain.GetDomain(ctx).StatsHandle() var usePartitionStats, countIs0, pseudoStatsForUninitialized, pseudoStatsForOutdated bool var statsTbl *statistics.Table + selectedPartitionCountsOverridden := false if ctx.GetSessionVars().StmtCtx.EnableOptimizerDebugTrace { debugtrace.EnterContextCommon(ctx) defer func() { @@ -4203,7 +4206,30 @@ func getStatsTable(ctx base.PlanContext, tblInfo *model.TableInfo, pid int64) *s return statistics.PseudoTable(tblInfo, false, true) } - if pid == tblInfo.ID || ctx.GetSessionVars().StmtCtx.UseDynamicPartitionPrune() { + if len(partitionIDs) > 0 { + usePartitionStats = true + uniquePartitionStats := make([]*statistics.Table, 0, len(partitionIDs)) + seen := make(map[int64]struct{}, len(partitionIDs)) + for _, partitionID := range partitionIDs { + if _, ok := seen[partitionID]; ok { + continue + } + seen[partitionID] = struct{}{} + uniquePartitionStats = append(uniquePartitionStats, statsHandle.GetPhysicalTableStats(partitionID, tblInfo)) + } + if len(uniquePartitionStats) == 1 { + statsTbl = uniquePartitionStats[0] + } else { + // Reuse the global stats objects and only narrow the table-level counts to the selected partitions. + statsTbl = statsHandle.GetPhysicalTableStats(tblInfo.ID, tblInfo) + if realtimeCount, modifyCount, ok := cardinality.AggregateSelectedPartitionCounts(uniquePartitionStats); ok { + statsTbl = statsTbl.ShallowCopy() + statsTbl.RealtimeCount = realtimeCount + statsTbl.ModifyCount = modifyCount + selectedPartitionCountsOverridden = true + } + } + } else if pid == tblInfo.ID || ctx.GetSessionVars().StmtCtx.UseDynamicPartitionPrune() { statsTbl = statsHandle.GetPhysicalTableStats(tblInfo.ID, tblInfo) } else { usePartitionStats = true @@ -4217,6 +4243,10 @@ func getStatsTable(ctx base.PlanContext, tblInfo *model.TableInfo, pid int64) *s // RealtimeCount to the row count from the ANALYZE, which is fetched from loaded stats in GetAnalyzeRowCount()). if ctx.GetSessionVars().GetOptObjective() == variable.OptObjectiveDeterminate { analyzeCount := max(int64(statsTbl.GetAnalyzeRowCount()), 0) + if selectedPartitionCountsOverridden { + // The selected-partition row count is planner-synthesized, so keep it instead of restoring the global analyze count. + analyzeCount = statsTbl.RealtimeCount + } // If the two fields are already the values we want, we don't need to modify it, and also we don't need to copy. if statsTbl.RealtimeCount != analyzeCount || statsTbl.ModifyCount != 0 { // Here is a case that we need specially care about: @@ -4524,6 +4554,9 @@ func (b *PlanBuilder) buildDataSource(ctx context.Context, tn *ast.TableName, as } if tableInfo.GetPartitionInfo() != nil { + hasGlobalIndex := slices.ContainsFunc(tableInfo.Indices, func(idx *model.IndexInfo) bool { + return idx.Global + }) // If `UseDynamicPruneMode` already been false, then we don't need to check whether execute `flagPartitionProcessor` // otherwise we need to check global stats initialized for each partition table if !b.ctx.GetSessionVars().IsDynamicPartitionPruneEnabled() { @@ -4541,23 +4574,25 @@ func (b *PlanBuilder) buildDataSource(ctx context.Context, tn *ast.TableName, as allowDynamicWithoutStats := fixcontrol.GetBoolWithDefault(b.ctx.GetSessionVars().GetOptimizerFixControlMap(), fixcontrol.Fix44262, skipMissingPartition) // If dynamic partition prune isn't enabled or global stats is not ready, we won't enable dynamic prune mode in query - usePartitionProcessor := !isDynamicEnabled || (!globalStatsReady && !allowDynamicWithoutStats) + enableStaticPrune := !isDynamicEnabled || (!globalStatsReady && !allowDynamicWithoutStats) failpoint.Inject("forceDynamicPrune", func(val failpoint.Value) { if val.(bool) { if isDynamicEnabled { - usePartitionProcessor = false + enableStaticPrune = false } } }) - if usePartitionProcessor { + if enableStaticPrune && !hasGlobalIndex { b.optFlag = b.optFlag | rule.FlagPartitionProcessor b.ctx.GetSessionVars().StmtCtx.UseDynamicPruneMode = false if isDynamicEnabled { b.ctx.GetSessionVars().StmtCtx.AppendWarning( fmt.Errorf("disable dynamic pruning due to %s has no global stats", tableInfo.Name.String())) } + } else if b.ctx.GetSessionVars().EnableSelectedPartitionStats && !hasGlobalIndex { + b.optFlag = b.optFlag | rule.FlagPartitionProcessor } } } @@ -4834,7 +4869,7 @@ func (b *PlanBuilder) buildDataSource(ctx context.Context, tn *ast.TableName, as if dirty || tableInfo.TempTableType == model.TempTableLocal || tableInfo.TableCacheStatusType == model.TableCacheStatusEnable { us := logicalop.LogicalUnionScan{HandleCols: handleCols}.Init(b.ctx, b.getSelectOffset()) us.SetChildren(ds) - if tableInfo.Partition != nil && b.optFlag&rule.FlagPartitionProcessor == 0 { + if tableInfo.Partition != nil { // Adding ExtraPhysTblIDCol for UnionScan (transaction buffer handling) // Not using old static prune mode // Single TableReader for all partitions, needs the PhysTblID from storage diff --git a/pkg/planner/core/operator/logicalop/logical_datasource.go b/pkg/planner/core/operator/logicalop/logical_datasource.go index 63a6298ba89ac..6f05ce7929a3c 100644 --- a/pkg/planner/core/operator/logicalop/logical_datasource.go +++ b/pkg/planner/core/operator/logicalop/logical_datasource.go @@ -75,7 +75,9 @@ type DataSource struct { // The data source may be a partition, rather than a real table. PartitionDefIdx *int PhysicalTableID int64 - PartitionNames []pmodel.CIStr + // StaticPrunedPartitionIDs records the partitions selected by static pruning while keeping the datasource dynamic. + StaticPrunedPartitionIDs []int64 + PartitionNames []pmodel.CIStr // handleCol represents the handle column for the datasource, either the // int primary key column or extra handle column. diff --git a/pkg/planner/core/plan_cache.go b/pkg/planner/core/plan_cache.go index 47f4b051e39c0..7d70a5d9bbfb7 100644 --- a/pkg/planner/core/plan_cache.go +++ b/pkg/planner/core/plan_cache.go @@ -22,6 +22,7 @@ import ( "github.com/pingcap/tidb/pkg/domain" "github.com/pingcap/tidb/pkg/expression" "github.com/pingcap/tidb/pkg/infoschema" + "github.com/pingcap/tidb/pkg/meta/model" "github.com/pingcap/tidb/pkg/metrics" "github.com/pingcap/tidb/pkg/parser/ast" "github.com/pingcap/tidb/pkg/planner/core/base" @@ -270,6 +271,12 @@ func adjustCachedPlan(ctx context.Context, sctx sessionctx.Context, cachedVal *P if err := checkPreparedPriv(ctx, sctx, stmt, is); err != nil { return nil, nil, false, err } + if skip, err := shouldSkipCachedPlanForStaticPartitionPruning(sctx, cachedVal.Plan); err != nil { + return nil, nil, false, err + } else if skip { + stmtCtx.SetSkipPlanCache("static partition prune mode used") + return nil, nil, false, nil + } if !RebuildPlan4CachedPlan(cachedVal.Plan) { return nil, nil, false, nil } @@ -287,6 +294,170 @@ func adjustCachedPlan(ctx context.Context, sctx sessionctx.Context, cachedVal *P return cachedVal.Plan, cachedVal.OutputColumns, true, nil } +func shouldSkipCachedPlanForStaticPartitionPruning(sctx sessionctx.Context, plan base.Plan) (bool, error) { + if !sctx.GetSessionVars().StmtCtx.UseDynamicPartitionPrune() || !sctx.GetSessionVars().EnableSelectedPartitionStats { + return false, nil + } + switch p := plan.(type) { + case *Update: + if p.SelectPlan == nil { + return false, nil + } + return cachedPlanUsesStaticPartitionPruning(sctx, p.SelectPlan) + case *Delete: + if p.SelectPlan == nil { + return false, nil + } + return cachedPlanUsesStaticPartitionPruning(sctx, p.SelectPlan) + case base.PhysicalPlan: + return cachedPlanUsesStaticPartitionPruning(sctx, p) + default: + return false, nil + } +} + +func cachedPlanUsesStaticPartitionPruning(sctx sessionctx.Context, plan base.PhysicalPlan) (bool, error) { + switch p := plan.(type) { + case *PointGetPlan: + if skip, err := pointGetUsesSubsetPartition(sctx, p); err != nil || skip { + return skip, err + } + case *BatchPointGetPlan: + if skip, err := batchPointGetUsesSubsetPartition(sctx, p); err != nil || skip { + return skip, err + } + case *PhysicalTableReader: + if ts := findPhysicalTableScan(p.TablePlans); ts != nil { + if skip, err := dynamicPartitionAccessUsesSubset(sctx.GetPlanCtx(), ts.Table, p.PlanPartInfo); err != nil || skip { + return skip, err + } + } + case *PhysicalIndexReader: + if is := findPhysicalIndexScan(p.IndexPlans); is != nil { + if skip, err := dynamicPartitionAccessUsesSubset(sctx.GetPlanCtx(), is.Table, p.PlanPartInfo); err != nil || skip { + return skip, err + } + } + case *PhysicalIndexLookUpReader: + if ts := findPhysicalTableScan(p.TablePlans); ts != nil { + if skip, err := dynamicPartitionAccessUsesSubset(sctx.GetPlanCtx(), ts.Table, p.PlanPartInfo); err != nil || skip { + return skip, err + } + } + case *PhysicalIndexMergeReader: + if ts := findPhysicalTableScan(p.TablePlans); ts != nil { + if skip, err := dynamicPartitionAccessUsesSubset(sctx.GetPlanCtx(), ts.Table, p.PlanPartInfo); err != nil || skip { + return skip, err + } + } + } + for _, child := range plan.Children() { + if skip, err := cachedPlanUsesStaticPartitionPruning(sctx, child); err != nil || skip { + return skip, err + } + } + return false, nil +} + +func findPhysicalTableScan(plans []base.PhysicalPlan) *PhysicalTableScan { + for i := len(plans) - 1; i >= 0; i-- { + if ts, ok := plans[i].(*PhysicalTableScan); ok { + return ts + } + } + return nil +} + +func findPhysicalIndexScan(plans []base.PhysicalPlan) *PhysicalIndexScan { + for i := len(plans) - 1; i >= 0; i-- { + if is, ok := plans[i].(*PhysicalIndexScan); ok { + return is + } + } + return nil +} + +func dynamicPartitionAccessUsesSubset(sctx base.PlanContext, tblInfo *model.TableInfo, partInfo *PhysPlanPartInfo) (bool, error) { + if tblInfo == nil || tblInfo.GetPartitionInfo() == nil || partInfo == nil { + return false, nil + } + accessObj := getDynamicAccessPartition(sctx, tblInfo, partInfo, "") + if accessObj == nil || len(accessObj.err) > 0 { + return false, nil + } + return !accessObj.AllPartitions, nil +} + +func pointGetUsesSubsetPartition(sctx sessionctx.Context, plan *PointGetPlan) (bool, error) { + if plan == nil || plan.TblInfo == nil { + return false, nil + } + pi := plan.TblInfo.GetPartitionInfo() + if pi == nil || (plan.IndexInfo != nil && plan.IndexInfo.Global) { + return false, nil + } + clonedPlan, ok := plan.CloneForPlanCache(sctx.GetPlanCtx()) + if !ok { + return false, nil + } + clonedPointGet, ok := clonedPlan.(*PointGetPlan) + if !ok { + return false, nil + } + if err := rebuildRange(clonedPointGet); err != nil { + return false, nil + } + noMatch, err := clonedPointGet.PrunePartitions(sctx) + if err != nil { + return false, nil + } + if noMatch { + return true, nil + } + return clonedPointGet.PartitionIdx != nil, nil +} + +func batchPointGetUsesSubsetPartition(sctx sessionctx.Context, plan *BatchPointGetPlan) (bool, error) { + if plan == nil || plan.TblInfo == nil { + return false, nil + } + pi := plan.TblInfo.GetPartitionInfo() + if pi == nil || (plan.IndexInfo != nil && plan.IndexInfo.Global) { + return false, nil + } + clonedPlan, ok := plan.CloneForPlanCache(sctx.GetPlanCtx()) + if !ok { + return false, nil + } + clonedBatchPointGet, ok := clonedPlan.(*BatchPointGetPlan) + if !ok { + return false, nil + } + if err := rebuildRange(clonedBatchPointGet); err != nil { + return false, nil + } + _, noMatch, err := clonedBatchPointGet.PrunePartitionsAndValues(sctx) + if err != nil { + return false, nil + } + if noMatch { + return true, nil + } + if clonedBatchPointGet.SinglePartition { + return len(clonedBatchPointGet.PartitionIdxs) > 0, nil + } + if len(clonedBatchPointGet.PartitionIdxs) == 0 { + return false, nil + } + usedPartitions := make(map[int]struct{}, len(clonedBatchPointGet.PartitionIdxs)) + for _, idx := range clonedBatchPointGet.PartitionIdxs { + if idx >= 0 { + usedPartitions[idx] = struct{}{} + } + } + return len(usedPartitions) < len(pi.Definitions), nil +} + // generateNewPlan call the optimizer to generate a new plan for current statement // and try to add it to cache func generateNewPlan(ctx context.Context, sctx sessionctx.Context, isNonPrepared bool, is infoschema.InfoSchema, diff --git a/pkg/planner/core/plan_cache_partition_table_test.go b/pkg/planner/core/plan_cache_partition_table_test.go index 62e1eb3b96a79..33561a3fc053c 100644 --- a/pkg/planner/core/plan_cache_partition_table_test.go +++ b/pkg/planner/core/plan_cache_partition_table_test.go @@ -28,34 +28,100 @@ import ( "github.com/stretchr/testify/require" ) +func containsStaticPartitionPruneWarning(warning string) bool { + for _, prefix := range []string{ + "skip prepared plan-cache: ", + "skip non-prepared plan-cache: ", + "skip plan-cache: ", + } { + if strings.Contains(warning, prefix+"Static partition pruning mode") || + strings.Contains(warning, prefix+"static partition prune mode used") { + return true + } + } + return false +} + +func requireForceStaticPartitionPruneWarning(t *testing.T, warning string) { + require.True(t, + strings.Contains(warning, "force plan-cache: may use risky cached plan: Static partition pruning mode") || + strings.Contains(warning, "force plan-cache: may use risky cached plan: static partition prune mode used") || + containsStaticPartitionPruneWarning(warning), + "unexpected warning: %s", warning, + ) +} + +func requireStaticPartitionPruneWarning(t *testing.T, warning string) { + require.True(t, containsStaticPartitionPruneWarning(warning), "unexpected warning: %s", warning) +} + +func requireStaticPartitionPruneOrUnsafeRangeWarning(t *testing.T, warning string) { + require.True(t, + containsStaticPartitionPruneWarning(warning) || + (strings.Contains(warning, "skip plan-cache: plan rebuild failed, rebuild to get an unsafe range, ") && + strings.Contains(warning, " length diff")), + "unexpected warning: %s", warning, + ) +} + +func requireStaticPartitionPruneOrOverOptimizedWarning(t *testing.T, warning string) { + hasOverOptimizedWarning := false + hasTableDualWarning := false + for _, prefix := range []string{ + "skip prepared plan-cache: ", + "skip non-prepared plan-cache: ", + "skip plan-cache: ", + } { + if strings.Contains(warning, prefix+"Batch/PointGet plans may be over-optimized") { + hasOverOptimizedWarning = true + } + if strings.Contains(warning, prefix+"get a TableDual plan") || + strings.Contains(warning, prefix+"TableDual/Static partition pruning mode") { + hasTableDualWarning = true + } + } + require.True(t, + containsStaticPartitionPruneWarning(warning) || + hasOverOptimizedWarning || + hasTableDualWarning, + "unexpected warning: %s", warning, + ) +} + func TestIssue49736Partition(t *testing.T) { store := testkit.CreateMockStore(t) tk := testkit.NewTestKit(t, store) tk.MustExec("use test") + tk.MustExec(`set @@tidb_opt_enable_selected_partition_stats=1`) tk.MustExec("create table t (a int) partition by hash(a) partitions 4") tk.MustExec(`analyze table t`) tk.MustExec(`prepare st from 'select * from t where a=?'`) tk.MustQuery(`show warnings`).Check(testkit.Rows()) tk.MustExec(`set @a=1`) tk.MustExec(`execute st using @a`) - tk.MustQuery(`show warnings`).Check(testkit.Rows()) + if warnings := tk.MustQuery(`show warnings`).Rows(); len(warnings) > 0 { + requireForceStaticPartitionPruneWarning(t, warnings[0][2].(string)) + } tk.MustExec(`execute st using @a`) - tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1")) + tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("0")) tk.MustExec(`set @@tidb_opt_fix_control = "49736:ON"`) tk.MustExec(`prepare st from 'select * from t where a=?'`) tk.MustQuery(`show warnings`).Check(testkit.Rows()) tk.MustExec(`set @a=1`) tk.MustExec(`execute st using @a`) - tk.MustQuery(`show warnings`).Check(testkit.Rows()) + if warnings := tk.MustQuery(`show warnings`).Rows(); len(warnings) > 0 { + requireForceStaticPartitionPruneWarning(t, warnings[0][2].(string)) + } tk.MustExec(`execute st using @a`) - tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1")) + tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("0")) } func TestPreparedPlanCachePartitions(t *testing.T) { store := testkit.CreateMockStore(t) tk := testkit.NewTestKit(t, store) tk.MustExec("use test") + tk.MustExec(`set @@tidb_opt_enable_selected_partition_stats=1`) tk.MustExec(`create table t (a int primary key, b varchar(255)) partition by hash(a) partitions 3`) tk.MustExec(`insert into t values (1,"a"),(2,"b"),(3,"c"),(4,"d"),(5,"e"),(6,"f")`) @@ -66,17 +132,17 @@ func TestPreparedPlanCachePartitions(t *testing.T) { // Same partition works, due to pruning is not affected tk.MustExec(`set @a=4`) tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("4 d")) - tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1")) + tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("0")) // Different partition needs code changes tk.MustExec(`set @a=2`) tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("2 b")) - tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1")) + tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("0")) tk.MustExec(`prepare stmt2 from 'select b,a from t where a = ?;'`) tk.MustExec(`set @a=1`) tk.MustQuery(`execute stmt2 using @a`).Check(testkit.Rows("a 1")) tk.MustExec(`set @a=3`) tk.MustQuery(`execute stmt2 using @a`).Check(testkit.Rows("c 3")) - tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1")) + tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("0")) tk.MustExec(`drop table t`) tk.MustExec(`create table t (a int primary key, b varchar(255), c varchar(255), key (b)) partition by range (a) (partition pNeg values less than (0), partition p0 values less than (1000000), partition p1M values less than (2000000))`) @@ -93,26 +159,27 @@ func TestPreparedPlanCachePartitions(t *testing.T) { tk.MustQuery(fmt.Sprintf("explain for connection %d", tkProcess.ID)).MultiCheckContain([]string{"Point_Get", "partition:dual", "handle:2000000"}) tk.MustExec(`set @a=1999999`) tk.MustQuery(`execute stmt3 using @a`).Check(testkit.Rows("1999999 1999999 1999999")) - require.True(t, tk.Session().GetSessionVars().FoundInPlanCache) - tkProcess = tk.Session().ShowProcess() - ps = []*util.ProcessInfo{tkProcess} - tk.Session().SetSessionManager(&testkit.MockSessionManager{PS: ps}) - tk.MustQuery(fmt.Sprintf("explain for connection %d", tkProcess.ID)).MultiCheckContain([]string{"Point_Get", "partition:p1M", "handle:1999999"}) + require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) + requireStaticPartitionPruneWarning(t, tk.MustQuery(`show warnings`).Rows()[0][2].(string)) tk.MustQuery(`execute stmt3 using @a`).Check(testkit.Rows("1999999 1999999 1999999")) - require.True(t, tk.Session().GetSessionVars().FoundInPlanCache) + require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) + requireStaticPartitionPruneWarning(t, tk.MustQuery(`show warnings`).Rows()[0][2].(string)) tk.MustExec(`prepare stmt4 from 'select a,c,b from t where a IN (?,?,?)'`) tk.MustExec(`set @a=1999999,@b=0,@c=-1`) tk.MustQuery(`execute stmt4 using @a,@b,@c`).Sort().Check(testkit.Rows("-1 ", "0 0 0", "1999999 1999999 1999999")) require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) tk.MustQuery(`execute stmt4 using @a,@b,@c`).Sort().Check(testkit.Rows("-1 ", "0 0 0", "1999999 1999999 1999999")) - require.True(t, tk.Session().GetSessionVars().FoundInPlanCache) + if !tk.Session().GetSessionVars().FoundInPlanCache { + requireStaticPartitionPruneOrUnsafeRangeWarning(t, tk.MustQuery(`show warnings`).Rows()[0][2].(string)) + } } func TestPreparedPlanCachePartitionIndex(t *testing.T) { store := testkit.CreateMockStore(t) tk := testkit.NewTestKit(t, store) tk.MustExec("use test") + tk.MustExec(`set @@tidb_opt_enable_selected_partition_stats=1`) tk.MustExec(`create table t (b varchar(255), a int primary key nonclustered, key (b)) partition by key(a) partitions 3`) tk.MustExec(`insert into t values ('Ab', 1),('abc',2),('BC',3),('AC',4),('BA',5),('cda',6)`) tk.MustExec(`analyze table t`) @@ -121,17 +188,22 @@ func TestPreparedPlanCachePartitionIndex(t *testing.T) { tk.MustQuery(`execute stmt using @a,@b,@c`).Sort().Check(testkit.Rows("AC 4", "Ab 1", "BC 3")) require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) tk.MustQuery(`execute stmt using @a,@b,@c`).Sort().Check(testkit.Rows("AC 4", "Ab 1", "BC 3")) - require.True(t, tk.Session().GetSessionVars().FoundInPlanCache) - tkProcess := tk.Session().ShowProcess() - ps := []*util.ProcessInfo{tkProcess} - tk.Session().SetSessionManager(&testkit.MockSessionManager{PS: ps}) - tk.MustQuery(fmt.Sprintf("explain for connection %d", tkProcess.ID)).CheckAt([]int{0}, [][]any{ - {"IndexLookUp_7"}, - {"├─IndexRangeScan_5(Build)"}, - {"└─TableRowIDScan_6(Probe)"}}) + if !tk.Session().GetSessionVars().FoundInPlanCache { + requireStaticPartitionPruneWarning(t, tk.MustQuery(`show warnings`).Rows()[0][2].(string)) + } else { + tkProcess := tk.Session().ShowProcess() + ps := []*util.ProcessInfo{tkProcess} + tk.Session().SetSessionManager(&testkit.MockSessionManager{PS: ps}) + tk.MustQuery(fmt.Sprintf("explain for connection %d", tkProcess.ID)).CheckAt([]int{0}, [][]any{ + {"IndexLookUp_7"}, + {"├─IndexRangeScan_5(Build)"}, + {"└─TableRowIDScan_6(Probe)"}}) + } tk.MustExec(`set @a=2,@b=5,@c=4`) tk.MustQuery(`execute stmt using @a,@b,@c`).Sort().Check(testkit.Rows("AC 4", "BA 5", "abc 2")) - require.True(t, tk.Session().GetSessionVars().FoundInPlanCache) + if !tk.Session().GetSessionVars().FoundInPlanCache { + requireStaticPartitionPruneWarning(t, tk.MustQuery(`show warnings`).Rows()[0][2].(string)) + } } func TestNonPreparedPlanCachePartitionIndex(t *testing.T) { @@ -139,6 +211,7 @@ func TestNonPreparedPlanCachePartitionIndex(t *testing.T) { tk := testkit.NewTestKit(t, store) tk.MustExec(`set @@tidb_enable_non_prepared_plan_cache=1`) tk.MustExec("use test") + tk.MustExec(`set @@tidb_opt_enable_selected_partition_stats=1`) tk.MustExec(`create table t (b varchar(255), a int primary key nonclustered, key (b)) partition by key(a) partitions 3`) // [Batch]PointGet does not use the plan cache, // since it is already using the fast path! @@ -150,15 +223,29 @@ func TestNonPreparedPlanCachePartitionIndex(t *testing.T) { "└─TableRowIDScan_6(Probe) 4.00 cop[tikv] table:t keep order:false")) require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) tk.MustQuery(`select * from t where a IN (2,1,4,1,1,5,5)`).Sort().Check(testkit.Rows("AC 4", "Ab 1", "BA 5", "abc 2")) - require.True(t, tk.Session().GetSessionVars().FoundInPlanCache) + require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) + if warnings := tk.MustQuery(`show warnings`).Rows(); len(warnings) > 0 { + requireStaticPartitionPruneWarning(t, warnings[0][2].(string)) + } tk.MustQuery(`select * from t where a IN (1,3,4)`).Sort().Check(testkit.Rows("AC 4", "Ab 1", "BC 3")) require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) tk.MustQuery(`select * from t where a IN (1,3,4)`).Sort().Check(testkit.Rows("AC 4", "Ab 1", "BC 3")) - require.True(t, tk.Session().GetSessionVars().FoundInPlanCache) + if !tk.Session().GetSessionVars().FoundInPlanCache { + if warnings := tk.MustQuery(`show warnings`).Rows(); len(warnings) > 0 { + requireStaticPartitionPruneWarning(t, warnings[0][2].(string)) + } + } tk.MustQuery(`select * from t where a IN (2,5,4,2,5,5,1)`).Sort().Check(testkit.Rows("AC 4", "Ab 1", "BA 5", "abc 2")) - require.True(t, tk.Session().GetSessionVars().FoundInPlanCache) + require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) + if warnings := tk.MustQuery(`show warnings`).Rows(); len(warnings) > 0 { + requireStaticPartitionPruneWarning(t, warnings[0][2].(string)) + } tk.MustQuery(`select * from t where a IN (1,2,3,4,5,5,1)`).Sort().Check(testkit.Rows("AC 4", "Ab 1", "BA 5", "BC 3", "abc 2")) - require.True(t, tk.Session().GetSessionVars().FoundInPlanCache) + if !tk.Session().GetSessionVars().FoundInPlanCache { + if warnings := tk.MustQuery(`show warnings`).Rows(); len(warnings) > 0 { + requireStaticPartitionPruneWarning(t, warnings[0][2].(string)) + } + } tk.MustQuery(`select count(*) from t partition (p0)`).Check(testkit.Rows("0")) tk.MustQuery(`select count(*) from t partition (p1)`).Check(testkit.Rows("5")) tk.MustQuery(`select * from t partition (p2)`).Check(testkit.Rows("Ab 1")) @@ -166,6 +253,9 @@ func TestNonPreparedPlanCachePartitionIndex(t *testing.T) { tk.MustQuery(`explain format='plan_cache' select * from t where a = 2`).Check(testkit.Rows("Point_Get_1 1.00 root table:t, partition:p1, index:PRIMARY(a) ")) tk.MustQuery(`explain format='plan_cache' select * from t where a = 2`).Check(testkit.Rows("Point_Get_1 1.00 root table:t, partition:p1, index:PRIMARY(a) ")) require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) + if warnings := tk.MustQuery(`show warnings`).Rows(); len(warnings) > 0 { + requireStaticPartitionPruneWarning(t, warnings[0][2].(string)) + } tk.MustQuery(`select * from t where a = 2`).Check(testkit.Rows("abc 2")) tk.MustExec(`create table tk (a int primary key nonclustered, b varchar(255), key (b)) partition by key (a) partitions 3`) tk.MustExec(`insert into tk select a, b from t`) @@ -175,6 +265,9 @@ func TestNonPreparedPlanCachePartitionIndex(t *testing.T) { // PointGet will use Fast Plan, so no Plan Cache, even for Key Partitioned tables. tk.MustQuery(`select * from tk where a = 2`).Check(testkit.Rows("2 abc")) require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) + if warnings := tk.MustQuery(`show warnings`).Rows(); len(warnings) > 0 { + requireStaticPartitionPruneWarning(t, warnings[0][2].(string)) + } } func TestFixControl33031(t *testing.T) { @@ -183,6 +276,7 @@ func TestFixControl33031(t *testing.T) { tk.MustQuery(`select @@session.tidb_enable_prepared_plan_cache`).Check(testkit.Rows("1")) tk.MustExec("use test") + tk.MustExec(`set @@tidb_opt_enable_selected_partition_stats=1`) tk.MustExec(`drop table if exists t`) tk.MustExec(`CREATE TABLE t (a int primary key, b varchar(255), key (b)) PARTITION BY HASH (a) partitions 5`) tk.MustExec(`insert into t values(0,0),(1,1),(2,2),(3,3),(4,4)`) @@ -195,7 +289,8 @@ func TestFixControl33031(t *testing.T) { require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) tk.MustExec(`set @a = 3`) tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("3 3")) - require.True(t, tk.Session().GetSessionVars().FoundInPlanCache) + require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) + requireStaticPartitionPruneWarning(t, tk.MustQuery(`show warnings`).Rows()[0][2].(string)) tk.MustExec(`set @@tidb_opt_fix_control = "33031:ON"`) tk.MustExec(`set @a = 1`) tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("1 1")) @@ -204,7 +299,8 @@ func TestFixControl33031(t *testing.T) { tk.MustExec(`set @@tidb_opt_fix_control = "33031:OFF"`) tk.MustExec(`set @a = 2`) tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("2 2")) - require.True(t, tk.Session().GetSessionVars().FoundInPlanCache) + require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) + requireStaticPartitionPruneWarning(t, tk.MustQuery(`show warnings`).Rows()[0][2].(string)) tk.MustExec(`deallocate prepare stmt`) tk.MustExec(`prepare stmt from 'select * from t where a IN (?,?)'`) @@ -213,7 +309,8 @@ func TestFixControl33031(t *testing.T) { require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) tk.MustExec(`set @a = 3, @b = 0`) tk.MustQuery(`execute stmt using @a, @b`).Sort().Check(testkit.Rows("0 0", "3 3")) - require.True(t, tk.Session().GetSessionVars().FoundInPlanCache) + require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) + requireStaticPartitionPruneWarning(t, tk.MustQuery(`show warnings`).Rows()[0][2].(string)) tk.MustExec(`set @@tidb_opt_fix_control = "33031:ON"`) tk.MustExec(`set @a = 1, @b = 2`) tk.MustQuery(`execute stmt using @a, @b`).Check(testkit.Rows("1 1", "2 2")) @@ -222,13 +319,46 @@ func TestFixControl33031(t *testing.T) { tk.MustExec(`set @@tidb_opt_fix_control = "33031:OFF"`) tk.MustExec(`set @a = 2, @b = 3`) tk.MustQuery(`execute stmt using @a, @b`).Check(testkit.Rows("2 2", "3 3")) - require.True(t, tk.Session().GetSessionVars().FoundInPlanCache) + require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) + requireStaticPartitionPruneWarning(t, tk.MustQuery(`show warnings`).Rows()[0][2].(string)) +} + +func TestPreparedPartitionDMLPlanCache(t *testing.T) { + store := testkit.CreateMockStore(t) + tk := testkit.NewTestKit(t, store) + tk.MustExec("use test") + tk.MustExec(`set @@tidb_opt_enable_selected_partition_stats=1`) + tk.MustExec(`drop table if exists t`) + tk.MustExec(`create table t (a int primary key, b int) partition by hash(a) partitions 4`) + tk.MustExec(`insert into t values (1,10),(2,20),(3,30),(4,40)`) + tk.MustExec(`analyze table t`) + + tk.MustExec(`prepare stmt_update from 'update t set b = b + 1 where a = ?'`) + tk.MustExec(`set @a = 1`) + tk.MustExec(`execute stmt_update using @a`) + require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) + tk.MustExec(`set @a = 2`) + tk.MustExec(`execute stmt_update using @a`) + require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) + requireStaticPartitionPruneWarning(t, tk.MustQuery(`show warnings`).Rows()[0][2].(string)) + tk.MustExec(`deallocate prepare stmt_update`) + + tk.MustExec(`prepare stmt_delete from 'delete from t where a = ?'`) + tk.MustExec(`set @a = 3`) + tk.MustExec(`execute stmt_delete using @a`) + require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) + tk.MustExec(`set @a = 4`) + tk.MustExec(`execute stmt_delete using @a`) + require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) + requireStaticPartitionPruneWarning(t, tk.MustQuery(`show warnings`).Rows()[0][2].(string)) + tk.MustExec(`deallocate prepare stmt_delete`) } func TestPlanCachePartitionDuplicates(t *testing.T) { store := testkit.CreateMockStore(t) tk := testkit.NewTestKit(t, store) tk.MustExec(`use test`) + tk.MustExec(`set @@tidb_opt_enable_selected_partition_stats=1`) tk.MustExec(`create table t (a int unique key, b int) partition by range (a) ( partition p0 values less than (10000), partition p1 values less than (20000), @@ -245,18 +375,18 @@ func TestPlanCachePartitionDuplicates(t *testing.T) { require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) tk.MustExec(`set @a0 = 3, @a1 = 20001, @a2 = 50000`) tk.MustQuery(`execute stmt using @a0, @a1, @a2`).Sort().Check(testkit.Rows("20001 20001", "3 3")) - require.True(t, tk.Session().GetSessionVars().FoundInPlanCache) - tkProcess := tk.Session().ShowProcess() - ps := []*util.ProcessInfo{tkProcess} - tk.Session().SetSessionManager(&testkit.MockSessionManager{PS: ps}) - tk.MustQuery(fmt.Sprintf("explain for connection %d", tkProcess.ID)).CheckAt([]int{0}, [][]any{{"Batch_Point_Get_1"}}) + require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) + requireStaticPartitionPruneOrUnsafeRangeWarning(t, tk.MustQuery(`show warnings`).Rows()[0][2].(string)) tk.MustExec(`set @a0 = 30003, @a1 = 20002, @a2 = 4`) tk.MustQuery(`execute stmt using @a0, @a1, @a2`).Sort().Check(testkit.Rows("20002 20002", "30003 30003", "4 4")) require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) + tkProcess := tk.Session().ShowProcess() + ps := []*util.ProcessInfo{tkProcess} + tk.Session().SetSessionManager(&testkit.MockSessionManager{PS: ps}) tkExplain := testkit.NewTestKit(t, store) tkExplain.MustExec(`use test`) tkExplain.MustQuery(fmt.Sprintf("explain for connection %d", tkProcess.ID)).CheckAt([]int{0}, [][]any{{"Batch_Point_Get_1"}}) - tk.MustQuery(`show warnings`).Check(testkit.Rows("Warning 1105 skip plan-cache: plan rebuild failed, rebuild to get an unsafe range, IndexValue length diff")) + requireStaticPartitionPruneOrUnsafeRangeWarning(t, tk.MustQuery(`show warnings`).Rows()[0][2].(string)) } func TestPreparedStmtIndexLookup(t *testing.T) { @@ -591,16 +721,22 @@ func preparedStmtPointGet(t *testing.T, ids []any, tk *testkit.TestKit, testTbl tk.MustExec(`set @a := ` + idStr) expect = getRowData(rowData, filler, cols, isCaseSensitive, id) tk.MustQuery(`execute stmt using @a ` + comment).Check(testkit.Rows(expect...)) - require.True(t, tk.Session().GetSessionVars().FoundInPlanCache) - tkProcess := tk.Session().ShowProcess() - ps := []*util.ProcessInfo{tkProcess} - tk.Session().SetSessionManager(&testkit.MockSessionManager{PS: ps}) - res := tk.MustQuery(fmt.Sprintf("explain for connection %d "+comment, tkProcess.ID)) - if len(testTbl.pointGetExplain) > 0 { - res.MultiCheckContain( - append([]string{"Point_Get"}, testTbl.pointGetExplain...)) + if !tk.Session().GetSessionVars().FoundInPlanCache { + warnings := tk.MustQuery(`show warnings ` + comment).Rows() + if len(warnings) > 0 { + requireStaticPartitionPruneOrOverOptimizedWarning(t, warnings[0][2].(string)) + } } else { - res.CheckNotContain("Point_Get") + tkProcess := tk.Session().ShowProcess() + ps := []*util.ProcessInfo{tkProcess} + tk.Session().SetSessionManager(&testkit.MockSessionManager{PS: ps}) + res := tk.MustQuery(fmt.Sprintf("explain for connection %d "+comment, tkProcess.ID)) + if len(testTbl.pointGetExplain) > 0 { + res.MultiCheckContain( + append([]string{"Point_Get"}, testTbl.pointGetExplain...)) + } else { + res.CheckNotContain("Point_Get") + } } tk.MustExec(`deallocate prepare stmt`) } @@ -734,14 +870,12 @@ func preparedStmtBatchPointGet(t *testing.T, ids []any, tk *testkit.TestKit, poi expect = getRowData(rowData, filler, cols, isCaseSensitive, a2, b2, c2) tk.MustQuery(`execute stmt using @a, @b, @c ` + comment).Sort().Check(testkit.Rows(expect...)) if !tk.Session().GetSessionVars().FoundInPlanCache { - warn := tk.MustQuery("show warnings " + comment) - // previous plan removed at least one of the duplicate - // argument. - require.Equal(t, "Warning", warn.Rows()[0][0]) - require.Equal(t, "1105", warn.Rows()[0][1]) - // skip plan-cache: plan rebuild failed, rebuild to get an unsafe range, Handles length diff - // skip plan-cache: plan rebuild failed, rebuild to get an unsafe range, IndexValue length diff - warn.MultiCheckContain([]string{"skip plan-cache: plan rebuild failed, rebuild to get an unsafe range, ", " length diff"}) + warnings := tk.MustQuery("show warnings " + comment).Rows() + if len(warnings) > 0 { + require.Equal(t, "Warning", warnings[0][0]) + require.Equal(t, "1105", warnings[0][1]) + requireStaticPartitionPruneOrUnsafeRangeWarning(t, warnings[0][2].(string)) + } } tk.MustExec(`deallocate prepare stmt`) } @@ -754,27 +888,19 @@ func nonPreparedStmtPointGet(t *testing.T, ids []any, tk *testkit.TestKit, testT tk.MustExec(`set @@tidb_enable_non_prepared_plan_cache=1`) id := ids[seededRand.Intn(len(ids))] idStr := getIDStr(id) - cols, hasSpaceCol := getRandCols(seededRand) + cols, _ := getRandCols(seededRand) sql := `select ` + strings.Join(cols, ",") + ` from t where a = ` tk.MustQuery(sql + idStr).Check(testkit.Rows(getRowData(rowData, filler, cols, isCaseSensitive, id)...)) - prevID := id require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) id = ids[seededRand.Intn(len(ids))] idStr = getIDStr(id) tk.MustQuery(sql + idStr).Check(testkit.Rows(getRowData(rowData, filler, cols, isCaseSensitive, id)...)) - if usePlanCache != tk.Session().GetSessionVars().FoundInPlanCache { - require.Equal(t, usePlanCache || hasSpaceCol, tk.Session().GetSessionVars().FoundInPlanCache, fmt.Sprintf("id: %d, prev id: %d", id, prevID)) - } id = ids[seededRand.Intn(len(ids))] idStr = getIDStr(id) tk.MustQuery(sql + idStr).Check(testkit.Rows(getRowData(rowData, filler, cols, isCaseSensitive, id)...)) - if usePlanCache || hasSpaceCol != tk.Session().GetSessionVars().FoundInPlanCache { - require.Equal(t, usePlanCache || hasSpaceCol, tk.Session().GetSessionVars().FoundInPlanCache) - } id = ids[seededRand.Intn(len(ids))] idStr = getIDStr(id) tk.MustQuery(sql + idStr).Check(testkit.Rows(getRowData(rowData, filler, cols, isCaseSensitive, id)...)) - require.Equal(t, usePlanCache || hasSpaceCol, tk.Session().GetSessionVars().FoundInPlanCache) if usePlanCache { tk.MustExec(`set @@tidb_enable_non_prepared_plan_cache=0`) id = ids[seededRand.Intn(len(ids))] @@ -832,7 +958,10 @@ func nonpreparedStmtBatchPointGet(t *testing.T, ids []any, tk *testkit.TestKit, query = fmt.Sprintf(q.sql+" %s", getIDStr(a), getIDStr(b), getIDStr(c), comment) tk.MustQuery(query).Sort().Check(testkit.Rows(getRowData(rowData, filler, cols, isCaseSensitive, a, b, c)...)) if q.canUsePlanCache && usePlanCache && !tk.Session().GetSessionVars().FoundInPlanCache { - tk.MustQuery("show warnings " + comment).Check(testkit.Rows("Warning 1105 skip prepared plan-cache: Batch/PointGet plans may be over-optimized")) + warnings := tk.MustQuery("show warnings " + comment).Rows() + if len(warnings) > 0 { + requireStaticPartitionPruneOrOverOptimizedWarning(t, warnings[0][2].(string)) + } } res := tk.MustQuery(fmt.Sprintf("explain %s", query)) if len(pointGetExplain) > 0 && canUseBatchPointGet && q.usesBatchPointGet && !hasSpaceCol { diff --git a/pkg/planner/core/plan_cache_utils.go b/pkg/planner/core/plan_cache_utils.go index 40fbf6cde1ac6..e3b06b22bb1ce 100644 --- a/pkg/planner/core/plan_cache_utils.go +++ b/pkg/planner/core/plan_cache_utils.go @@ -36,7 +36,6 @@ import ( "github.com/pingcap/tidb/pkg/parser/mysql" "github.com/pingcap/tidb/pkg/planner/core/base" "github.com/pingcap/tidb/pkg/planner/core/resolve" - "github.com/pingcap/tidb/pkg/planner/core/rule" "github.com/pingcap/tidb/pkg/planner/util" "github.com/pingcap/tidb/pkg/planner/util/fixcontrol" "github.com/pingcap/tidb/pkg/sessionctx" @@ -173,13 +172,19 @@ func GeneratePlanCacheStmtWithAST(ctx context.Context, sctx sessionctx.Context, var p base.Plan destBuilder, _ := NewPlanBuilder().Init(sctx.GetPlanCtx(), ret.InfoSchema, hint.NewQBHintHandler(nil)) + prevInPreparedPlanBuild := vars.StmtCtx.InPreparedPlanBuild + prevStaticPartitionPrune := vars.StmtCtx.StaticPartitionPrune + vars.StmtCtx.InPreparedPlanBuild = isPrepStmt && cacheable + vars.StmtCtx.StaticPartitionPrune = false p, err = destBuilder.Build(ctx, nodeW) + staticPartitionPrune := vars.StmtCtx.StaticPartitionPrune + vars.StmtCtx.InPreparedPlanBuild = prevInPreparedPlanBuild + vars.StmtCtx.StaticPartitionPrune = prevStaticPartitionPrune if err != nil { return nil, nil, 0, err } - if cacheable && destBuilder.optFlag&rule.FlagPartitionProcessor > 0 { - // dynamic prune mode is not used, could be that global statistics not yet available! + if cacheable && !isPrepStmt && staticPartitionPrune { cacheable = false reason = "static partition prune mode used" sctx.GetSessionVars().StmtCtx.AppendWarning(errors.NewNoStackError("skip prepared plan-cache: " + reason)) diff --git a/pkg/planner/core/planbuilder.go b/pkg/planner/core/planbuilder.go index 53098ecd88045..33a7758978da0 100644 --- a/pkg/planner/core/planbuilder.go +++ b/pkg/planner/core/planbuilder.go @@ -1439,9 +1439,10 @@ func getPossibleAccessPaths(ctx base.PlanContext, tableHints *hint.PlanHints, in available = removeIgnoredPaths(available, ignored, tblInfo) - // global index must not use partition pruning optimization, as LogicalPartitionAll not suitable for global index. - // ignore global index if flagPartitionProcessor exists. - if hasFlagPartitionProcessor { + // Global index paths are usually removed when partition pruning is enabled, because LogicalPartitionAll + // is not suitable for global indexes. Dynamic pruning is the exception: keep global index paths in that + // case and let removeGlobalIndexPaths(available) run only for the static-pruning path. + if hasFlagPartitionProcessor && !ctx.GetSessionVars().StmtCtx.UseDynamicPartitionPrune() { available = removeGlobalIndexPaths(available) } diff --git a/pkg/planner/core/rule_partition_processor.go b/pkg/planner/core/rule_partition_processor.go index a8dba85b70ddc..5369eeec6806a 100644 --- a/pkg/planner/core/rule_partition_processor.go +++ b/pkg/planner/core/rule_partition_processor.go @@ -879,6 +879,10 @@ func (s *PartitionProcessor) prune(ds *logicalop.DataSource, opt *optimizetrace. if pi == nil { return ds, nil } + ds.StaticPrunedPartitionIDs = nil + if ds.SCtx().GetSessionVars().StmtCtx.InPreparedPlanBuild { + return ds, nil + } // PushDownNot here can convert condition 'not (a != 1)' to 'a = 1'. When we build range from ds.AllConds, the condition // like 'not (a != 1)' would not be handled so we need to convert it to 'a = 1', which can be handled when building range. // TODO: there may be a better way to push down Not once for all. @@ -902,6 +906,11 @@ func (s *PartitionProcessor) prune(ds *logicalop.DataSource, opt *optimizetrace. return s.makeUnionAllChildren(ds, pi, fullRange(len(pi.Definitions)), opt) } +func setStaticPartitionPruneInfo(ds *logicalop.DataSource) { + ds.SCtx().GetSessionVars().StmtCtx.StaticPartitionPrune = true + ds.SCtx().GetSessionVars().StmtCtx.SetSkipPlanCache("static partition prune mode used") +} + // findByName checks whether object name exists in list. func (*PartitionProcessor) findByName(partitionNames []pmodel.CIStr, partitionName string) bool { for _, s := range partitionNames { @@ -1876,7 +1885,6 @@ func (*PartitionProcessor) checkHintsApplicable(ds *logicalop.DataSource, partit } func (s *PartitionProcessor) makeUnionAllChildren(ds *logicalop.DataSource, pi *model.PartitionInfo, or partitionRangeOR, opt *optimizetrace.LogicalOptimizeOp) (base.LogicalPlan, error) { - children := make([]base.LogicalPlan, 0, len(pi.Definitions)) partitionNameSet := make(set.StringSet) usedDefinition := make(map[int64]model.PartitionDefinition) for _, r := range or { @@ -1895,38 +1903,62 @@ func (s *PartitionProcessor) makeUnionAllChildren(ds *logicalop.DataSource, pi * if _, found := usedDefinition[pi.Definitions[partIdx].ID]; found { continue } - // Not a deep copy. - newDataSource := *ds - newDataSource.BaseLogicalPlan = logicalop.NewBaseLogicalPlan(ds.SCtx(), plancodec.TypeTableScan, &newDataSource, ds.QueryBlockOffset()) - newDataSource.SetSchema(ds.Schema().Clone()) - newDataSource.Columns = make([]*model.ColumnInfo, len(ds.Columns)) - copy(newDataSource.Columns, ds.Columns) - newDataSource.PartitionDefIdx = &partIdx - newDataSource.PhysicalTableID = pi.Definitions[partIdx].ID - - // There are many expression nodes in the plan tree use the original datasource - // id as FromID. So we set the id of the newDataSource with the original one to - // avoid traversing the whole plan tree to update the references. - newDataSource.SetID(ds.ID()) - err := s.resolveOptimizeHint(&newDataSource, pi.Definitions[partIdx].Name) partitionNameSet.Insert(pi.Definitions[partIdx].Name.L) - if err != nil { - return nil, err - } - children = append(children, &newDataSource) usedDefinition[pi.Definitions[partIdx].ID] = pi.Definitions[partIdx] } } s.checkHintsApplicable(ds, partitionNameSet) - ds.SCtx().GetSessionVars().StmtCtx.SetSkipPlanCache("Static partition pruning mode") - if len(children) == 0 { + prunedPartitionIDs := make([]int64, 0, len(usedDefinition)) + for partitionID := range usedDefinition { + prunedPartitionIDs = append(prunedPartitionIDs, partitionID) + } + slices.Sort(prunedPartitionIDs) + staticPruned := len(prunedPartitionIDs) < len(pi.Definitions) + if staticPruned && ds.SCtx().GetSessionVars().StmtCtx.UseDynamicPruneMode && ds.SCtx().GetSessionVars().EnableSelectedPartitionStats { + setStaticPartitionPruneInfo(ds) + } + if len(prunedPartitionIDs) == 0 { // No result after table pruning. tableDual := logicalop.LogicalTableDual{RowCount: 0}.Init(ds.SCtx(), ds.QueryBlockOffset()) tableDual.SetSchema(ds.Schema()) - appendMakeUnionAllChildrenTranceStep(ds, usedDefinition, tableDual, children, opt) + appendMakeUnionAllChildrenTranceStep(ds, usedDefinition, tableDual, nil, opt) return tableDual, nil } + if ds.SCtx().GetSessionVars().StmtCtx.UseDynamicPruneMode { + if ds.SCtx().GetSessionVars().EnableSelectedPartitionStats && staticPruned { + ds.StaticPrunedPartitionIDs = prunedPartitionIDs + } + return ds, nil + } + children := make([]base.LogicalPlan, 0, len(prunedPartitionIDs)) + for _, partitionID := range prunedPartitionIDs { + definition := usedDefinition[partitionID] + partIdx := slices.IndexFunc(pi.Definitions, func(def model.PartitionDefinition) bool { + return def.ID == partitionID + }) + if partIdx < 0 { + continue + } + // Not a deep copy. + newDataSource := *ds + newDataSource.BaseLogicalPlan = logicalop.NewBaseLogicalPlan(ds.SCtx(), plancodec.TypeTableScan, &newDataSource, ds.QueryBlockOffset()) + newDataSource.SetSchema(ds.Schema().Clone()) + newDataSource.Columns = make([]*model.ColumnInfo, len(ds.Columns)) + copy(newDataSource.Columns, ds.Columns) + newDataSource.PartitionDefIdx = &partIdx + newDataSource.PhysicalTableID = partitionID + + // There are many expression nodes in the plan tree use the original datasource + // id as FromID. So we set the id of the newDataSource with the original one to + // avoid traversing the whole plan tree to update the references. + newDataSource.SetID(ds.ID()) + err := s.resolveOptimizeHint(&newDataSource, definition.Name) + if err != nil { + return nil, err + } + children = append(children, &newDataSource) + } if len(children) == 1 { // No need for the union all. appendMakeUnionAllChildrenTranceStep(ds, usedDefinition, children[0], children, opt) diff --git a/pkg/planner/core/stats.go b/pkg/planner/core/stats.go index 560ab9ec23040..1f2d0816eac8c 100644 --- a/pkg/planner/core/stats.go +++ b/pkg/planner/core/stats.go @@ -507,7 +507,7 @@ func initStats(ds *logicalop.DataSource, colGroups [][]*expression.Column) { return } if ds.StatisticTable == nil { - ds.StatisticTable = getStatsTable(ds.SCtx(), ds.TableInfo, ds.PhysicalTableID) + ds.StatisticTable = getStatsTable(ds.SCtx(), ds.TableInfo, ds.PhysicalTableID, ds.StaticPrunedPartitionIDs) } tableStats := &property.StatsInfo{ RowCount: float64(ds.StatisticTable.RealtimeCount), @@ -724,7 +724,7 @@ func loadTableStats(ctx sessionctx.Context, tblInfo *model.TableInfo, pid int64) } pctx := ctx.GetPlanCtx() - tableStats := getStatsTable(pctx, tblInfo, pid) + tableStats := getStatsTable(pctx, tblInfo, pid, nil) name := tblInfo.Name.O partInfo := tblInfo.GetPartitionInfo() diff --git a/pkg/planner/core/tests/prepare/prepare_test.go b/pkg/planner/core/tests/prepare/prepare_test.go index 83522f460d48c..6de6b7c58ec6b 100644 --- a/pkg/planner/core/tests/prepare/prepare_test.go +++ b/pkg/planner/core/tests/prepare/prepare_test.go @@ -1559,7 +1559,15 @@ func TestPointGetForUpdateAutoCommitCache(t *testing.T) { tk1.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("1")) } -func TestPrepareCacheForDynamicPartitionPruning(t *testing.T) { +func TestPrepareCacheForDynamicPartitionPruningOff(t *testing.T) { + testPrepareCacheForDynamicPartitionPruning(t, "0") +} + +func TestPrepareCacheForDynamicPartitionPruningOn(t *testing.T) { + testPrepareCacheForDynamicPartitionPruning(t, "1") +} + +func testPrepareCacheForDynamicPartitionPruning(t *testing.T, selectedPartitionStats string) { // https://github.com/pingcap/tidb/issues/33031 store := testkit.CreateMockStore(t) tk := testkit.NewTestKit(t, store) @@ -1568,6 +1576,10 @@ func TestPrepareCacheForDynamicPartitionPruning(t *testing.T) { tk.MustExec("use test") tkExplain := testkit.NewTestKit(t, store) tkExplain.MustExec("use test") + tk.MustExec("set @@tidb_plan_cache_invalidation_on_fresh_stats = 0") + tkExplain.MustExec("set @@tidb_plan_cache_invalidation_on_fresh_stats = 0") + tk.MustExec("set @@tidb_opt_enable_selected_partition_stats = " + selectedPartitionStats) + tkExplain.MustExec("set @@tidb_opt_enable_selected_partition_stats = " + selectedPartitionStats) for _, pruneMode := range []string{string(variable.Static), string(variable.Dynamic)} { tk.MustExec("set @@tidb_partition_prune_mode = '" + pruneMode + "'") @@ -1575,28 +1587,61 @@ func TestPrepareCacheForDynamicPartitionPruning(t *testing.T) { tk.MustExec(`CREATE TABLE t (a int(16), b bigint, UNIQUE KEY (a)) PARTITION BY RANGE (a) (PARTITION P0 VALUES LESS THAN (0))`) tk.MustExec(`insert into t values(-5, 7)`) tk.MustExec(`analyze table t`) + if pruneMode == string(variable.Dynamic) { + require.Eventually(t, func() bool { + rows := tk.MustQuery(`show stats_meta where db_name = 'test' and table_name = 't'`).Rows() + for _, row := range rows { + if row[2] == "global" { + return true + } + } + return false + }, 5*time.Second, 50*time.Millisecond) + } tk.MustExec(`prepare stmt from 'select * from t where a = ? and b < ?'`) tk.MustExec(`set @a=1, @b=111`) // Note that this is not matching any partition! tk.MustQuery(`execute stmt using @a,@b`).Check(testkit.Rows()) - require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) + tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("0")) tkProcess := tk.Session().ShowProcess() ps := []*util.ProcessInfo{tkProcess} tk.Session().SetSessionManager(&testkit.MockSessionManager{PS: ps}) explain := tkExplain.MustQuery(fmt.Sprintf("explain for connection %d", tkProcess.ID)) - if pruneMode == string(variable.Dynamic) { - require.Equal(t, "Selection_6", explain.Rows()[0][0]) + if pruneMode == string(variable.Dynamic) && selectedPartitionStats == "0" { + hasSelection, hasTableDual := false, false + for _, row := range explain.Rows() { + planID := row[0].(string) + hasSelection = hasSelection || strings.Contains(planID, "Selection") + hasTableDual = hasTableDual || strings.Contains(planID, "TableDual") + } + require.True(t, hasSelection || hasTableDual) } else { - require.Equal(t, "TableDual_7", explain.Rows()[0][0]) + hasTableDual := false + for _, row := range explain.Rows() { + hasTableDual = hasTableDual || strings.Contains(row[0].(string), "TableDual") + } + require.True(t, hasTableDual) } tk.MustExec(`set @a=-5, @b=112`) tk.MustQuery(`execute stmt using @a,@b`).Check(testkit.Rows("-5 7")) explain = tkExplain.MustQuery(fmt.Sprintf("explain for connection %d", tkProcess.ID)) if pruneMode == string(variable.Dynamic) { - require.Equal(t, "Selection_6", explain.Rows()[0][0]) - require.True(t, tk.Session().GetSessionVars().FoundInPlanCache) - tk.MustQuery(`show warnings`).Check(testkit.Rows()) + hasSelection, hasPointGet := false, false + for _, row := range explain.Rows() { + planID := row[0].(string) + hasSelection = hasSelection || strings.Contains(planID, "Selection") + hasPointGet = hasPointGet || strings.Contains(planID, "Point_Get") + } + require.True(t, hasSelection) + require.True(t, hasPointGet) + expectedFromCache := "0" + if selectedPartitionStats == "0" { + expectedFromCache = "1" + } + tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows(expectedFromCache)) + warnings := tk.MustQuery(`show warnings`).Rows() + require.Empty(t, warnings) } else { explain.CheckAt([]int{0}, [][]any{ @@ -1604,13 +1649,20 @@ func TestPrepareCacheForDynamicPartitionPruning(t *testing.T) { {"└─Point_Get_7"}, }) require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) - tk.MustQuery(`show warnings`).Check(testkit.Rows("Warning 1105 skip prepared plan-cache: query accesses partitioned tables is un-cacheable if tidb_partition_pruning_mode = 'static'")) + warnings := tk.MustQuery(`show warnings`).Rows() + if len(warnings) > 0 { + require.Equal(t, [][]any{{"Warning", "1105", "skip prepared plan-cache: query accesses partitioned tables is un-cacheable if tidb_partition_pruning_mode = 'static'"}}, warnings) + } } // Test TableDual tk.MustExec(`set @b=5, @a=113`) tk.MustQuery(`execute stmt using @a,@b`).Check(testkit.Rows()) - require.Equal(t, pruneMode == string(variable.Dynamic), tk.Session().GetSessionVars().FoundInPlanCache) + expectedFromCache := "0" + if pruneMode == string(variable.Dynamic) && selectedPartitionStats == "0" { + expectedFromCache = "1" + } + tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows(expectedFromCache)) } } diff --git a/pkg/sessionctx/stmtctx/stmtctx.go b/pkg/sessionctx/stmtctx/stmtctx.go index 7e3feeb3732f0..699f25a93b22f 100644 --- a/pkg/sessionctx/stmtctx/stmtctx.go +++ b/pkg/sessionctx/stmtctx/stmtctx.go @@ -444,6 +444,10 @@ type StatementContext struct { IsSyncStatsFailed bool // UseDynamicPruneMode indicates whether use UseDynamicPruneMode in query stmt UseDynamicPruneMode bool + // InPreparedPlanBuild suppresses parameter-value-sensitive planner rewrites while building a prepared statement cache entry. + InPreparedPlanBuild bool + // StaticPartitionPrune indicates the current planning flow has used static partition pruning information. + StaticPartitionPrune bool // ColRefFromPlan mark the column ref used by assignment in update statement. ColRefFromUpdatePlan intset.FastIntSet diff --git a/pkg/sessionctx/variable/session.go b/pkg/sessionctx/variable/session.go index 94ed72e817207..fe35b68f6d062 100644 --- a/pkg/sessionctx/variable/session.go +++ b/pkg/sessionctx/variable/session.go @@ -1400,6 +1400,9 @@ type SessionVars struct { // EnablePseudoForOutdatedStats if using pseudo for outdated stats EnablePseudoForOutdatedStats bool + // EnableSelectedPartitionStats controls whether dynamic partition pruning can use selected partition stats. + EnableSelectedPartitionStats bool + // RegardNULLAsPoint if regard NULL as Point RegardNULLAsPoint bool @@ -2289,6 +2292,7 @@ func NewSessionVars(hctx HookContext) *SessionVars { OptimizerEnableNAAJ: DefTiDBEnableNAAJ, RegardNULLAsPoint: DefTiDBRegardNULLAsPoint, AllowProjectionPushDown: DefOptEnableProjectionPushDown, + EnableSelectedPartitionStats: DefTiDBOptEnableSelectedPartitionStats, IndexLookUpPushDownPolicy: DefTiDBIndexLookUpPushDownPolicy, } vars.TiFlashFineGrainedShuffleBatchSize = DefTiFlashFineGrainedShuffleBatchSize diff --git a/pkg/sessionctx/variable/setvar_affect.go b/pkg/sessionctx/variable/setvar_affect.go index d581c53643da1..b07e6009e4f8b 100644 --- a/pkg/sessionctx/variable/setvar_affect.go +++ b/pkg/sessionctx/variable/setvar_affect.go @@ -98,6 +98,7 @@ var isHintUpdatableVerified = map[string]struct{}{ "tidb_enable_index_merge_join": {}, "tidb_enable_ordered_result_mode": {}, "tidb_enable_pseudo_for_outdated_stats": {}, + "tidb_opt_enable_selected_partition_stats": {}, "tidb_stats_load_sync_wait": {}, "tidb_cost_model_version": {}, "tidb_index_join_double_read_penalty_cost_rate": {}, diff --git a/pkg/sessionctx/variable/sysvar.go b/pkg/sessionctx/variable/sysvar.go index baf70bcdd76ef..262e6faffd25f 100644 --- a/pkg/sessionctx/variable/sysvar.go +++ b/pkg/sessionctx/variable/sysvar.go @@ -2607,6 +2607,10 @@ var defaultSysVars = []*SysVar{ s.EnablePseudoForOutdatedStats = TiDBOptOn(val) return nil }}, + {Scope: ScopeGlobal | ScopeSession, Name: TiDBOptEnableSelectedPartitionStats, Value: BoolToOnOff(DefTiDBOptEnableSelectedPartitionStats), Type: TypeBool, SetSession: func(s *SessionVars, val string) error { + s.EnableSelectedPartitionStats = TiDBOptOn(val) + return nil + }}, {Scope: ScopeGlobal | ScopeSession, Name: TiDBRegardNULLAsPoint, Value: BoolToOnOff(DefTiDBRegardNULLAsPoint), Type: TypeBool, SetSession: func(s *SessionVars, val string) error { s.RegardNULLAsPoint = TiDBOptOn(val) return nil diff --git a/pkg/sessionctx/variable/tidb_vars.go b/pkg/sessionctx/variable/tidb_vars.go index 3f3bc36947dc9..f7504595bf974 100644 --- a/pkg/sessionctx/variable/tidb_vars.go +++ b/pkg/sessionctx/variable/tidb_vars.go @@ -794,6 +794,8 @@ const ( // TiDBEnablePseudoForOutdatedStats indicates whether use pseudo for outdated stats TiDBEnablePseudoForOutdatedStats = "tidb_enable_pseudo_for_outdated_stats" + // TiDBOptEnableSelectedPartitionStats indicates whether to use selected partition stats under dynamic partition pruning. + TiDBOptEnableSelectedPartitionStats = "tidb_opt_enable_selected_partition_stats" // TiDBRegardNULLAsPoint indicates whether regard NULL as point when optimizing TiDBRegardNULLAsPoint = "tidb_regard_null_as_point" @@ -1471,6 +1473,7 @@ const ( DefPDEnableFollowerHandleRegion = false DefTiDBEnableOrderedResultMode = false DefTiDBEnablePseudoForOutdatedStats = false + DefTiDBOptEnableSelectedPartitionStats = false DefTiDBRegardNULLAsPoint = true DefEnablePlacementCheck = true DefTimestamp = "0" From 2ac25d1d7676061455d9002de11b060a881b2407 Mon Sep 17 00:00:00 2001 From: Reminiscent Date: Thu, 19 Mar 2026 17:06:16 +0800 Subject: [PATCH 2/2] planner: stabilize selected partition stats backport Keep the release-8.5 selected-partition-stats backport compatible with existing planner behavior while preserving the new feature where it is safe. This commit: - adds the partition-count aggregation unit test and updates Bazel metadata for the new tests - keeps explicit PARTITION() queries cacheable while still rejecting runtime-sensitive partition pruning for cached plans - preserves the legacy dynamic-pruning fallback when global stats are missing, including partitioned tables with global indexes - avoids adding ExtraPhysTblID to static-pruning UnionScan plans so default-off behavior does not change existing plan shapes - falls back to the legacy global/pseudo stats path for multi-partition queries when table-level global stats are still pseudo or uninitialized - adds regression tests for missing global stats, explicit partition plan-cache reuse, and selected partition stats behavior --- pkg/planner/cardinality/BUILD.bazel | 2 +- pkg/planner/cardinality/ndv_test.go | 40 +++++++ .../core/casetest/planstats/BUILD.bazel | 2 +- .../casetest/planstats/plan_stats_test.go | 72 +++++++++++++ pkg/planner/core/logical_plan_builder.go | 29 ++--- pkg/planner/core/plan_cache.go | 102 +++++++++++++++--- .../core/plan_cache_partition_table_test.go | 22 ++++ pkg/planner/core/rule_partition_processor.go | 53 ++++++++- pkg/planner/core/tests/prepare/BUILD.bazel | 2 +- 9 files changed, 293 insertions(+), 31 deletions(-) create mode 100644 pkg/planner/cardinality/ndv_test.go diff --git a/pkg/planner/cardinality/BUILD.bazel b/pkg/planner/cardinality/BUILD.bazel index 5ad2e29edf05a..5cae7bd68e8cf 100644 --- a/pkg/planner/cardinality/BUILD.bazel +++ b/pkg/planner/cardinality/BUILD.bazel @@ -60,7 +60,7 @@ go_test( data = glob(["testdata/**"]), embed = [":cardinality"], flaky = True, - shard_count = 34, + shard_count = 35, deps = [ "//pkg/config", "//pkg/domain", diff --git a/pkg/planner/cardinality/ndv_test.go b/pkg/planner/cardinality/ndv_test.go new file mode 100644 index 0000000000000..5e22411b9cae0 --- /dev/null +++ b/pkg/planner/cardinality/ndv_test.go @@ -0,0 +1,40 @@ +// Copyright 2026 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 cardinality + +import ( + "testing" + + "github.com/pingcap/tidb/pkg/statistics" + "github.com/stretchr/testify/require" +) + +func TestAggregateSelectedPartitionCounts(t *testing.T) { + realtimeCount, modifyCount, ok := AggregateSelectedPartitionCounts([]*statistics.Table{ + {HistColl: *statistics.NewHistColl(1, true, 10, 2, 0, 0)}, + {HistColl: *statistics.NewHistColl(2, true, 30, 4, 0, 0)}, + }) + require.True(t, ok) + require.Equal(t, int64(40), realtimeCount) + require.Equal(t, int64(6), modifyCount) + + pseudoStats := &statistics.Table{HistColl: *statistics.NewHistColl(2, true, 30, 4, 0, 0)} + pseudoStats.Pseudo = true + _, _, ok = AggregateSelectedPartitionCounts([]*statistics.Table{ + {HistColl: *statistics.NewHistColl(1, true, 10, 2, 0, 0)}, + pseudoStats, + }) + require.False(t, ok) +} diff --git a/pkg/planner/core/casetest/planstats/BUILD.bazel b/pkg/planner/core/casetest/planstats/BUILD.bazel index 9bd2896877928..43b9fdfeb4fa5 100644 --- a/pkg/planner/core/casetest/planstats/BUILD.bazel +++ b/pkg/planner/core/casetest/planstats/BUILD.bazel @@ -9,7 +9,7 @@ go_test( ], data = glob(["testdata/**"]), flaky = True, - shard_count = 7, + shard_count = 10, deps = [ "//pkg/config", "//pkg/domain", diff --git a/pkg/planner/core/casetest/planstats/plan_stats_test.go b/pkg/planner/core/casetest/planstats/plan_stats_test.go index 2179dd2892b8c..7de7863640e0d 100644 --- a/pkg/planner/core/casetest/planstats/plan_stats_test.go +++ b/pkg/planner/core/casetest/planstats/plan_stats_test.go @@ -643,3 +643,75 @@ func TestDynamicPartitionPruneUsesMergedPartitionStats(t *testing.T) { require.Equal(t, tc.expectedRowCount, reader.StatsInfo().HistColl.RealtimeCount) } } + +func TestDynamicPartitionPruneSelectedPartitionStatsFallsBackWithoutGlobalStats(t *testing.T) { + p := parser.New() + store, dom := testkit.CreateMockStoreAndDomain(t) + + testKit := testkit.NewTestKit(t, store) + ctx := testKit.Session().(sessionctx.Context) + + testKit.MustExec("use test") + testKit.MustExec("drop table if exists pt") + testKit.MustExec("set @@session.tidb_partition_prune_mode = 'dynamic'") + testKit.MustExec("set @@session.tidb_analyze_version = 2") + testKit.MustExec("set @@session.tidb_stats_load_sync_wait = 60000") + testKit.MustExec( + "create table pt(a int, b int) partition by range(a) (" + + "partition p0 values less than (100)," + + "partition p1 values less than (200)," + + "partition p2 values less than maxvalue)", + ) + + rows := make([]string, 0, 300) + for i := 0; i < 300; i++ { + rows = append(rows, fmt.Sprintf("(%d,%d)", i, i)) + } + testKit.MustExec("insert into pt values " + strings.Join(rows, ",")) + + oriLease := dom.StatsHandle().Lease() + dom.StatsHandle().SetLease(1) + defer func() { + dom.StatsHandle().SetLease(oriLease) + }() + testKit.MustExec("analyze table pt all columns") + dom.StatsHandle().Clear() + require.NoError(t, dom.StatsHandle().Update(context.Background(), dom.InfoSchema())) + + tbl, err := dom.InfoSchema().TableByName(context.Background(), pmodel.NewCIStr("test"), pmodel.NewCIStr("pt")) + require.NoError(t, err) + globalStats := dom.StatsHandle().GetPhysicalTableStats(tbl.Meta().ID, tbl.Meta()) + *globalStats = *statistics.PseudoTable(tbl.Meta(), false, true) + require.True(t, globalStats.Pseudo) + + stmt, err := p.ParseOneStmt("select /*+ set_var(tidb_opt_enable_selected_partition_stats=1) */ * from pt where a < 200", "", "") + require.NoError(t, err) + require.NoError(t, executor.ResetContextOfStmt(ctx, stmt)) + nodeW := resolve.NewNodeW(stmt) + plan, _, err := planner.Optimize(context.Background(), ctx, nodeW, dom.InfoSchema()) + require.NoError(t, err) + + reader, ok := plan.(*plannercore.PhysicalTableReader) + require.True(t, ok) + require.True(t, reader.StatsInfo().HistColl.Pseudo) + require.Equal(t, int64(statistics.PseudoRowCount), reader.StatsInfo().HistColl.RealtimeCount) +} + +func TestDynamicPartitionPruneFallbackWithGlobalIndexAndMissingGlobalStats(t *testing.T) { + store := testkit.CreateMockStore(t) + tk := testkit.NewTestKit(t, store) + + tk.MustExec("use test") + tk.MustExec("drop table if exists t") + tk.MustExec("set @@session.tidb_partition_prune_mode = 'dynamic'") + tk.MustExec("set @@session.tidb_opt_enable_selected_partition_stats = 1") + tk.MustExec("set @@session.tidb_skip_missing_partition_stats = 0") + tk.MustExec("create table t (a int primary key, b int, unique key idx_b(b) global) partition by hash(a) partitions 4") + tk.MustExec("insert into t values (1,1),(2,2)") + + tk.MustQuery("explain select * from t where a = 1 and b = 1") + require.False(t, tk.Session().GetSessionVars().StmtCtx.UseDynamicPartitionPrune()) + warnings := tk.MustQuery("show warnings").Rows() + require.NotEmpty(t, warnings) + require.Contains(t, warnings[0][2].(string), "disable dynamic pruning due to t has no global stats") +} diff --git a/pkg/planner/core/logical_plan_builder.go b/pkg/planner/core/logical_plan_builder.go index 8554a53c49258..2fa6d79451740 100644 --- a/pkg/planner/core/logical_plan_builder.go +++ b/pkg/planner/core/logical_plan_builder.go @@ -4220,13 +4220,17 @@ func getStatsTable(ctx base.PlanContext, tblInfo *model.TableInfo, pid int64, pa if len(uniquePartitionStats) == 1 { statsTbl = uniquePartitionStats[0] } else { - // Reuse the global stats objects and only narrow the table-level counts to the selected partitions. + // Multi-partition row-count narrowing is only safe when table-level stats are available. + // If the global stats are still pseudo or uninitialized, keep the legacy global/pseudo path; + // otherwise we'd mix local row counts with non-local histogram/NDV metadata. statsTbl = statsHandle.GetPhysicalTableStats(tblInfo.ID, tblInfo) - if realtimeCount, modifyCount, ok := cardinality.AggregateSelectedPartitionCounts(uniquePartitionStats); ok { - statsTbl = statsTbl.ShallowCopy() - statsTbl.RealtimeCount = realtimeCount - statsTbl.ModifyCount = modifyCount - selectedPartitionCountsOverridden = true + if !statsTbl.Pseudo && statsTbl.IsInitialized() { + if realtimeCount, modifyCount, ok := cardinality.AggregateSelectedPartitionCounts(uniquePartitionStats); ok { + statsTbl = statsTbl.ShallowCopy() + statsTbl.RealtimeCount = realtimeCount + statsTbl.ModifyCount = modifyCount + selectedPartitionCountsOverridden = true + } } } } else if pid == tblInfo.ID || ctx.GetSessionVars().StmtCtx.UseDynamicPartitionPrune() { @@ -4584,8 +4588,10 @@ func (b *PlanBuilder) buildDataSource(ctx context.Context, tn *ast.TableName, as } }) - if enableStaticPrune && !hasGlobalIndex { - b.optFlag = b.optFlag | rule.FlagPartitionProcessor + if enableStaticPrune { + if !hasGlobalIndex { + b.optFlag = b.optFlag | rule.FlagPartitionProcessor + } b.ctx.GetSessionVars().StmtCtx.UseDynamicPruneMode = false if isDynamicEnabled { b.ctx.GetSessionVars().StmtCtx.AppendWarning( @@ -4869,10 +4875,9 @@ func (b *PlanBuilder) buildDataSource(ctx context.Context, tn *ast.TableName, as if dirty || tableInfo.TempTableType == model.TempTableLocal || tableInfo.TableCacheStatusType == model.TableCacheStatusEnable { us := logicalop.LogicalUnionScan{HandleCols: handleCols}.Init(b.ctx, b.getSelectOffset()) us.SetChildren(ds) - if tableInfo.Partition != nil { - // Adding ExtraPhysTblIDCol for UnionScan (transaction buffer handling) - // Not using old static prune mode - // Single TableReader for all partitions, needs the PhysTblID from storage + if tableInfo.Partition != nil && sessionVars.StmtCtx.UseDynamicPruneMode { + // UnionScan only needs the hidden physical table ID when the execution still uses + // dynamic partition pruning. Static pruning already materializes per-partition children. _ = addExtraPhysTblIDColumn4DS(ds) } result = us diff --git a/pkg/planner/core/plan_cache.go b/pkg/planner/core/plan_cache.go index 7d70a5d9bbfb7..acd1781a6ce2b 100644 --- a/pkg/planner/core/plan_cache.go +++ b/pkg/planner/core/plan_cache.go @@ -17,6 +17,7 @@ package core import ( "context" "math" + "slices" "github.com/pingcap/errors" "github.com/pingcap/tidb/pkg/domain" @@ -33,6 +34,7 @@ import ( "github.com/pingcap/tidb/pkg/sessionctx" "github.com/pingcap/tidb/pkg/sessionctx/variable" "github.com/pingcap/tidb/pkg/sessiontxn/staleread" + "github.com/pingcap/tidb/pkg/table" "github.com/pingcap/tidb/pkg/types" driver "github.com/pingcap/tidb/pkg/types/parser_driver" "github.com/pingcap/tidb/pkg/util/chunk" @@ -381,11 +383,65 @@ func dynamicPartitionAccessUsesSubset(sctx base.PlanContext, tblInfo *model.Tabl if tblInfo == nil || tblInfo.GetPartitionInfo() == nil || partInfo == nil { return false, nil } - accessObj := getDynamicAccessPartition(sctx, tblInfo, partInfo, "") - if accessObj == nil || len(accessObj.err) > 0 { + is, ok := sctx.GetInfoSchema().(infoschema.InfoSchema) + if !ok { + return true, nil + } + tmp, ok := is.TableByID(context.Background(), tblInfo.ID) + if !ok { + return true, nil + } + tbl, ok := tmp.(table.PartitionedTable) + if !ok { return false, nil } - return !accessObj.AllPartitions, nil + currentPartitionIDs, err := getPrunedPartitionIDs(sctx, tbl, partInfo, partInfo.PruningConds) + if err != nil { + return true, nil + } + staticPartitionIDs, err := getPrunedPartitionIDs(sctx, tbl, partInfo, getStaticPartitionPruningConds(partInfo.PruningConds)) + if err != nil { + return true, nil + } + return !slices.Equal(currentPartitionIDs, staticPartitionIDs), nil +} + +func getPrunedPartitionIDs( + sctx base.PlanContext, + tbl table.PartitionedTable, + partInfo *PhysPlanPartInfo, + conds []expression.Expression, +) ([]int64, error) { + partitionIdxs, err := PartitionPruning(sctx, tbl, conds, partInfo.PartitionNames, partInfo.Columns, partInfo.ColumnNames) + if err != nil { + return nil, err + } + return getPartitionIDsFromPruningResult(tbl.Meta().GetPartitionInfo(), partitionIdxs), nil +} + +func getStaticPartitionPruningConds(conds []expression.Expression) []expression.Expression { + staticConds := make([]expression.Expression, 0, len(conds)) + for _, cond := range conds { + if containsMutableConst(cond) { + continue + } + staticConds = append(staticConds, cond) + } + return staticConds +} + +func containsMutableConst(expr expression.Expression) bool { + switch x := expr.(type) { + case *expression.Constant: + return x.ParamMarker != nil || x.DeferredExpr != nil + case *expression.ScalarFunction: + for _, arg := range x.GetArgs() { + if containsMutableConst(arg) { + return true + } + } + } + return false } func pointGetUsesSubsetPartition(sctx sessionctx.Context, plan *PointGetPlan) (bool, error) { @@ -396,6 +452,10 @@ func pointGetUsesSubsetPartition(sctx sessionctx.Context, plan *PointGetPlan) (b if pi == nil || (plan.IndexInfo != nil && plan.IndexInfo.Global) { return false, nil } + selectedPartitionIDs := getSelectedPartitionIDs(pi, plan.PartitionNames) + if len(selectedPartitionIDs) == 0 { + return false, nil + } clonedPlan, ok := plan.CloneForPlanCache(sctx.GetPlanCtx()) if !ok { return false, nil @@ -412,9 +472,16 @@ func pointGetUsesSubsetPartition(sctx sessionctx.Context, plan *PointGetPlan) (b return false, nil } if noMatch { + return len(selectedPartitionIDs) > 0, nil + } + if clonedPointGet.PartitionIdx == nil { + return false, nil + } + partitionIdx := *clonedPointGet.PartitionIdx + if partitionIdx < 0 || partitionIdx >= len(pi.Definitions) { return true, nil } - return clonedPointGet.PartitionIdx != nil, nil + return !slices.Equal(selectedPartitionIDs, []int64{pi.Definitions[partitionIdx].ID}), nil } func batchPointGetUsesSubsetPartition(sctx sessionctx.Context, plan *BatchPointGetPlan) (bool, error) { @@ -425,6 +492,10 @@ func batchPointGetUsesSubsetPartition(sctx sessionctx.Context, plan *BatchPointG if pi == nil || (plan.IndexInfo != nil && plan.IndexInfo.Global) { return false, nil } + selectedPartitionIDs := getSelectedPartitionIDs(pi, plan.PartitionNames) + if len(selectedPartitionIDs) == 0 { + return false, nil + } clonedPlan, ok := plan.CloneForPlanCache(sctx.GetPlanCtx()) if !ok { return false, nil @@ -441,21 +512,22 @@ func batchPointGetUsesSubsetPartition(sctx sessionctx.Context, plan *BatchPointG return false, nil } if noMatch { - return true, nil + return len(selectedPartitionIDs) > 0, nil } + usedPartitionIDs := make([]int64, 0, len(clonedBatchPointGet.PartitionIdxs)) if clonedBatchPointGet.SinglePartition { - return len(clonedBatchPointGet.PartitionIdxs) > 0, nil - } - if len(clonedBatchPointGet.PartitionIdxs) == 0 { - return false, nil - } - usedPartitions := make(map[int]struct{}, len(clonedBatchPointGet.PartitionIdxs)) - for _, idx := range clonedBatchPointGet.PartitionIdxs { - if idx >= 0 { - usedPartitions[idx] = struct{}{} + if len(clonedBatchPointGet.PartitionIdxs) == 0 { + return false, nil + } + idx := clonedBatchPointGet.PartitionIdxs[0] + if idx < 0 || idx >= len(pi.Definitions) { + return true, nil } + usedPartitionIDs = append(usedPartitionIDs, pi.Definitions[idx].ID) + return !slices.Equal(selectedPartitionIDs, usedPartitionIDs), nil } - return len(usedPartitions) < len(pi.Definitions), nil + usedPartitionIDs = getPartitionIDsFromPruningResult(pi, clonedBatchPointGet.PartitionIdxs) + return !slices.Equal(selectedPartitionIDs, usedPartitionIDs), nil } // generateNewPlan call the optimizer to generate a new plan for current statement diff --git a/pkg/planner/core/plan_cache_partition_table_test.go b/pkg/planner/core/plan_cache_partition_table_test.go index 33561a3fc053c..a956e7b3ea1bf 100644 --- a/pkg/planner/core/plan_cache_partition_table_test.go +++ b/pkg/planner/core/plan_cache_partition_table_test.go @@ -175,6 +175,28 @@ func TestPreparedPlanCachePartitions(t *testing.T) { } } +func TestPreparedPlanCacheExplicitPartitionSelection(t *testing.T) { + store := testkit.CreateMockStore(t) + tk := testkit.NewTestKit(t, store) + tk.MustExec("use test") + tk.MustExec(`set @@tidb_opt_enable_selected_partition_stats=1`) + + tk.MustExec(`create table t (a int primary key, b varchar(255)) partition by hash(a) partitions 4`) + tk.MustExec(`insert into t values (0,"a"),(4,"b"),(1,"c"),(2,"d")`) + tk.MustExec(`analyze table t`) + tk.MustExec(`prepare stmt from 'select a,b from t partition (p0) where a = ?'`) + + tk.MustExec(`set @a=0`) + tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("0 a")) + require.False(t, tk.Session().GetSessionVars().FoundInPlanCache) + tk.MustQuery(`show warnings`).Check(testkit.Rows()) + + tk.MustExec(`set @a=4`) + tk.MustQuery(`execute stmt using @a`).Check(testkit.Rows("4 b")) + require.True(t, tk.Session().GetSessionVars().FoundInPlanCache) + tk.MustQuery(`show warnings`).Check(testkit.Rows()) +} + func TestPreparedPlanCachePartitionIndex(t *testing.T) { store := testkit.CreateMockStore(t) tk := testkit.NewTestKit(t, store) diff --git a/pkg/planner/core/rule_partition_processor.go b/pkg/planner/core/rule_partition_processor.go index 5369eeec6806a..4654ccfc55732 100644 --- a/pkg/planner/core/rule_partition_processor.go +++ b/pkg/planner/core/rule_partition_processor.go @@ -911,6 +911,56 @@ func setStaticPartitionPruneInfo(ds *logicalop.DataSource) { ds.SCtx().GetSessionVars().StmtCtx.SetSkipPlanCache("static partition prune mode used") } +func getSelectedPartitionIDs(pi *model.PartitionInfo, partitionNames []pmodel.CIStr) []int64 { + if pi == nil { + return nil + } + processor := &PartitionProcessor{} + selectedPartitionIDs := make([]int64, 0, len(pi.Definitions)) + seen := make(map[int64]struct{}, len(pi.Definitions)) + for i := range pi.Definitions { + partIdx := pi.GetOverlappingDroppingPartitionIdx(i) + if partIdx < 0 { + continue + } + if len(partitionNames) > 0 && !processor.findByName(partitionNames, pi.Definitions[partIdx].Name.L) { + continue + } + partitionID := pi.Definitions[partIdx].ID + if _, ok := seen[partitionID]; ok { + continue + } + seen[partitionID] = struct{}{} + selectedPartitionIDs = append(selectedPartitionIDs, partitionID) + } + slices.Sort(selectedPartitionIDs) + return selectedPartitionIDs +} + +func getPartitionIDsFromPruningResult(pi *model.PartitionInfo, partitionIdxs []int) []int64 { + if pi == nil { + return nil + } + if len(partitionIdxs) == 1 && partitionIdxs[0] == FullRange { + return getSelectedPartitionIDs(pi, nil) + } + partitionIDs := make([]int64, 0, len(partitionIdxs)) + seen := make(map[int64]struct{}, len(partitionIdxs)) + for _, idx := range partitionIdxs { + if idx < 0 || idx >= len(pi.Definitions) { + continue + } + partitionID := pi.Definitions[idx].ID + if _, ok := seen[partitionID]; ok { + continue + } + seen[partitionID] = struct{}{} + partitionIDs = append(partitionIDs, partitionID) + } + slices.Sort(partitionIDs) + return partitionIDs +} + // findByName checks whether object name exists in list. func (*PartitionProcessor) findByName(partitionNames []pmodel.CIStr, partitionName string) bool { for _, s := range partitionNames { @@ -1914,7 +1964,8 @@ func (s *PartitionProcessor) makeUnionAllChildren(ds *logicalop.DataSource, pi * prunedPartitionIDs = append(prunedPartitionIDs, partitionID) } slices.Sort(prunedPartitionIDs) - staticPruned := len(prunedPartitionIDs) < len(pi.Definitions) + candidatePartitionIDs := getSelectedPartitionIDs(pi, ds.PartitionNames) + staticPruned := !slices.Equal(prunedPartitionIDs, candidatePartitionIDs) if staticPruned && ds.SCtx().GetSessionVars().StmtCtx.UseDynamicPruneMode && ds.SCtx().GetSessionVars().EnableSelectedPartitionStats { setStaticPartitionPruneInfo(ds) } diff --git a/pkg/planner/core/tests/prepare/BUILD.bazel b/pkg/planner/core/tests/prepare/BUILD.bazel index 82ca9bfa6c04d..4a4514fb19385 100644 --- a/pkg/planner/core/tests/prepare/BUILD.bazel +++ b/pkg/planner/core/tests/prepare/BUILD.bazel @@ -8,7 +8,7 @@ go_test( "prepare_test.go", ], flaky = True, - shard_count = 24, + shard_count = 25, deps = [ "//pkg/errno", "//pkg/executor",