diff --git a/pkg/planner/BUILD.bazel b/pkg/planner/BUILD.bazel index a0adbf6a9ef97..a09de68b586a4 100644 --- a/pkg/planner/BUILD.bazel +++ b/pkg/planner/BUILD.bazel @@ -13,6 +13,7 @@ go_library( "//pkg/parser/ast", "//pkg/planner/core", "//pkg/planner/core/base", + "//pkg/planner/core/operator/physicalop", "//pkg/planner/core/resolve", "//pkg/planner/core/rule", "//pkg/planner/indexadvisor", diff --git a/pkg/planner/core/casetest/mpp/BUILD.bazel b/pkg/planner/core/casetest/mpp/BUILD.bazel index 4d5be0776fc59..9c92c23872d49 100644 --- a/pkg/planner/core/casetest/mpp/BUILD.bazel +++ b/pkg/planner/core/casetest/mpp/BUILD.bazel @@ -4,23 +4,34 @@ go_test( name = "mpp_test", timeout = "short", srcs = [ + "engine_rounds_bench_test.go", + "engine_rounds_test.go", "main_test.go", "mpp_test.go", ], data = glob(["testdata/**"]), flaky = True, - shard_count = 23, + shard_count = 24, deps = [ "//pkg/config", "//pkg/domain", "//pkg/meta/model", "//pkg/parser/ast", + "//pkg/planner", + "//pkg/planner/core", + "//pkg/planner/core/resolve", + "//pkg/session", + "//pkg/store/mockstore", + "//pkg/store/mockstore/unistore", "//pkg/testkit", + "//pkg/testkit/external", "//pkg/testkit/testdata", "//pkg/testkit/testmain", "//pkg/testkit/testsetup", "@com_github_pingcap_failpoint//:failpoint", + "@com_github_pingcap_kvproto//pkg/metapb", "@com_github_stretchr_testify//require", + "@com_github_tikv_client_go_v2//testutils", "@org_uber_go_goleak//:goleak", ], ) diff --git a/pkg/planner/core/casetest/mpp/engine_rounds_bench_test.go b/pkg/planner/core/casetest/mpp/engine_rounds_bench_test.go new file mode 100644 index 0000000000000..b73ebf71cdebb --- /dev/null +++ b/pkg/planner/core/casetest/mpp/engine_rounds_bench_test.go @@ -0,0 +1,114 @@ +// 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 mpp + +import ( + "context" + "testing" + + "github.com/pingcap/tidb/pkg/domain" + "github.com/pingcap/tidb/pkg/parser/ast" + "github.com/pingcap/tidb/pkg/planner" + "github.com/pingcap/tidb/pkg/planner/core" + "github.com/pingcap/tidb/pkg/planner/core/resolve" + "github.com/pingcap/tidb/pkg/session" + "github.com/pingcap/tidb/pkg/testkit" + "github.com/stretchr/testify/require" +) + +// setBenchTiFlashReplica mirrors setRealTiFlashReplica but accepts testing.TB +// so it can be used from a benchmark. It is kept local to avoid changing the +// shared *testing.T helper signature. +func setBenchTiFlashReplica(b *testing.B, tk *testkit.TestKit, tableName string) { + tk.MustExec("alter table " + tableName + " set tiflash replica 1") + dom := domain.GetDomain(tk.Session()) + require.NoError(b, dom.Reload()) + tb, err := dom.InfoSchema().TableByName(context.Background(), ast.NewCIStr("test"), ast.NewCIStr(tableName)) + require.NoError(b, err) + require.NoError(b, dom.DDLExecutor().UpdateTableReplicaInfo(tk.Session(), tb.Meta().ID, true)) +} + +// BenchmarkEngineRoundOptimize measures the per-statement optimization time of +// the alternative logical plan driver. Running it on both candidate branches +// shows the compile-time cost of each design's arming strategy: the same query +// is optimized with the feature off (baseline) and on. +// +// The two designs differ only on all-TiKV plans. This branch arms the +// tiflash-only round whenever round 1 read from TiKV and every table has a +// replica; the mixed-engine-only design (#70005) never arms there, so its "on" +// time tracks the baseline. The benchmark focuses on that differentiator with +// an all-TiKV join, plus a scan-only fast-path baseline that no design arms. +// +// Cases whose armed round builds an MPP plan (round-1 TiFlash reads, or an +// aggregation whose tiflash-only round shuffles) are intentionally omitted: +// they either arm identically in both designs (so they do not differentiate +// them) or probe the mock TiFlash store's liveness, which is not reliably +// routable when the planner runs in isolation in a unit test. +func BenchmarkEngineRoundOptimize(b *testing.B) { + store := testkit.CreateMockStore(b, withMockTiFlashNodes(2)) + tk := testkit.NewTestKit(b, store) + tk.MustExec("use test") + + tk.MustExec("create table bench_small_a(a int primary key, b int)") + tk.MustExec("create table bench_small_b(a int primary key, b int)") + + tk.MustExec("insert into bench_small_a values (1, 10), (2, 20), (3, 30), (4, 40), (5, 50)") + tk.MustExec("insert into bench_small_b values (1, 11), (2, 22), (3, 33), (4, 44), (5, 55)") + tk.MustExec("analyze table bench_small_a, bench_small_b") + + setBenchTiFlashReplica(b, tk, "bench_small_a") + setBenchTiFlashReplica(b, tk, "bench_small_b") + + queries := []struct { + name string + sql string + }{ + { + // Round 1 reads both small tables from TiKV. This branch arms the + // tiflash-only round here; a mixed-engine-only design does not. + name: "oltp-indexed-join", + sql: "select count(*) from bench_small_a join bench_small_b" + + " on bench_small_a.a = bench_small_b.a", + }, + { + // No join or aggregation: the join/agg gate keeps every design off + // the rounds, so this measures the untouched fast path. + name: "scan-only", + sql: "select b from bench_small_a where a = 3", + }, + } + + for _, q := range queries { + ctx := tk.Session() + stmts, err := session.Parse(ctx, q.sql) + require.NoError(b, err) + require.Len(b, stmts, 1) + ret := &core.PreprocessorReturn{} + nodeW := resolve.NewNodeW(stmts[0]) + require.NoError(b, core.Preprocess(context.Background(), ctx, nodeW, core.WithPreprocessorReturn(ret))) + + for _, feature := range []string{"off", "on"} { + tk.MustExec("set @@tidb_opt_enable_alternative_logical_plans=" + feature) + b.Run(q.name+"/"+feature, func(b *testing.B) { + b.ReportAllocs() + b.ResetTimer() + for range b.N { + _, _, err := planner.Optimize(context.TODO(), ctx, nodeW, ret.InfoSchema) + require.NoError(b, err) + } + }) + } + } +} diff --git a/pkg/planner/core/casetest/mpp/engine_rounds_test.go b/pkg/planner/core/casetest/mpp/engine_rounds_test.go new file mode 100644 index 0000000000000..3c18715da1e5b --- /dev/null +++ b/pkg/planner/core/casetest/mpp/engine_rounds_test.go @@ -0,0 +1,193 @@ +// 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 mpp + +import ( + "fmt" + "strings" + "testing" + + "github.com/pingcap/failpoint" + "github.com/pingcap/kvproto/pkg/metapb" + "github.com/pingcap/tidb/pkg/domain" + "github.com/pingcap/tidb/pkg/store/mockstore" + "github.com/pingcap/tidb/pkg/store/mockstore/unistore" + "github.com/pingcap/tidb/pkg/testkit" + "github.com/pingcap/tidb/pkg/testkit/external" + "github.com/stretchr/testify/require" + "github.com/tikv/client-go/v2/testutils" +) + +const engineRoundFailpoint = "github.com/pingcap/tidb/pkg/planner/failIfAlternativeLogicalPlanRoundTriggered" + +// withMockTiFlashNodes adds TiFlash stores to the mock cluster so plans that +// read from TiFlash can actually execute in this test. +func withMockTiFlashNodes(nodes int) mockstore.MockTiKVStoreOption { + return mockstore.WithMultipleOptions( + mockstore.WithClusterInspector(func(c testutils.Cluster) { + mockCluster := c.(*unistore.Cluster) + _, _, region1 := mockstore.BootstrapWithSingleStore(c) + for tiflashIdx := range nodes { + store2 := c.AllocID() + peer2 := c.AllocID() + addr2 := fmt.Sprintf("tiflash%d", tiflashIdx) + mockCluster.AddStore(store2, addr2, &metapb.StoreLabel{Key: "engine", Value: "tiflash"}) + mockCluster.AddPeer(region1, store2, peer2) + } + }), + mockstore.WithStoreType(mockstore.EmbedUnistore), + ) +} + +func setRealTiFlashReplica(t *testing.T, tk *testkit.TestKit, tableName string) { + tk.MustExec("alter table " + tableName + " set tiflash replica 1") + tb := external.GetTableByName(t, tk, "test", tableName) + require.NoError(t, domain.GetDomain(tk.Session()).DDLExecutor(). + UpdateTableReplicaInfo(tk.Session(), tb.Meta().ID, true)) +} + +// armEngineRound makes the named round fail loudly if it runs, so a test can +// assert that the round was or was not armed. The failpoint matches on the +// statement text, so callers must run plain SQL: EXPLAIN's inner statement has +// an empty OriginalText and would never match. +func armEngineRound(t *testing.T, round, sql string) { + require.NoError(t, failpoint.Enable(engineRoundFailpoint, fmt.Sprintf("return(%q)", round+":"+sql))) + t.Cleanup(func() { + _ = failpoint.Disable(engineRoundFailpoint) + }) +} + +func disarmEngineRound(t *testing.T) { + require.NoError(t, failpoint.Disable(engineRoundFailpoint)) +} + +// requireRoundRan asserts the named round was armed for this statement: the +// failpoint aborts planning, so the statement must report the sentinel error. +func requireRoundRan(t *testing.T, tk *testkit.TestKit, round, sql string) { + armEngineRound(t, round, sql) + defer disarmEngineRound(t) + require.ErrorContains(t, tk.ExecToErr(sql), "unexpected alternative logical plan round", + "expected the %s round to run for: %s", round, sql) +} + +// requireRoundSkipped asserts the named round was not armed: with the failpoint +// active the statement still plans and executes normally. +func requireRoundSkipped(t *testing.T, tk *testkit.TestKit, round, sql string, expected [][]any) { + armEngineRound(t, round, sql) + defer disarmEngineRound(t) + tk.MustQuery(sql).Sort().Check(expected) +} + +// TestAlternativeEngineRoundArming covers when the tikv-only / tiflash-only +// alternative logical plan rounds arm. Unlike a gate that requires round 1 to +// have mixed engines, these rounds arm from the engines round 1 actually read: +// an all-TiKV plan still gets a whole-statement TiFlash/MPP alternative costed +// against it. The join-or-agg, replica, hint and enforce-mpp gates keep the +// extra rounds off statements that cannot benefit. +func TestAlternativeEngineRoundArming(t *testing.T) { + store := testkit.CreateMockStore(t, withMockTiFlashNodes(2)) + tk := testkit.NewTestKit(t, store) + tk.MustExec("use test") + + // Two small indexed tables with replicas: round 1 reads both from TiKV. + tk.MustExec("create table alt_small_a(a int primary key, b int)") + tk.MustExec("create table alt_small_b(a int primary key, b int)") + // A large table with a replica: round 1 pushes its aggregation to TiFlash. + tk.MustExec("create table alt_large(a int, b int, c int)") + // A table with no TiFlash replica at all. + tk.MustExec("create table alt_no_replica(a int primary key, b int)") + + tk.MustExec("insert into alt_small_a values (1, 10), (2, 20), (3, 30), (4, 40), (5, 50)") + tk.MustExec("insert into alt_small_b values (1, 11), (2, 22), (3, 33), (4, 44), (5, 55)") + tk.MustExec("insert into alt_no_replica values (1, 1), (2, 2), (3, 3), (4, 4), (5, 5)") + valsLarge := make([]string, 0, 5000) + for i := range 5000 { + valsLarge = append(valsLarge, fmt.Sprintf("(%d, %d, %d)", i%50, i, i%10)) + } + tk.MustExec("insert into alt_large values " + strings.Join(valsLarge, ",")) + tk.MustExec("analyze table alt_small_a, alt_small_b, alt_large, alt_no_replica") + + setRealTiFlashReplica(t, tk, "alt_small_a") + setRealTiFlashReplica(t, tk, "alt_small_b") + setRealTiFlashReplica(t, tk, "alt_large") + + tk.MustExec("set @@tidb_opt_enable_alternative_logical_plans=on") + enginesBefore := tk.MustQuery("select @@tidb_isolation_read_engines").Rows() + + // The headline case: round 1 reads only TiKV, so no plan mixing engines ever + // appears, yet the whole-statement TiFlash/MPP alternative is still uncosted + // and worth building. + allTiKVSQL := "select count(*) from alt_small_a join alt_small_b" + + " on alt_small_a.a = alt_small_b.a" + tk.MustQuery(allTiKVSQL).Check(testkit.Rows("5")) + stmtCtx := tk.Session().GetSessionVars().StmtCtx + require.True(t, stmtCtx.AlternativeLogicalPlanReadsFromTiKV) + require.False(t, stmtCtx.AlternativeLogicalPlanReadsFromTiFlash) + require.True(t, stmtCtx.AlternativeLogicalPlanHasJoinOrAgg) + requireRoundRan(t, tk, "tiflash-only", allTiKVSQL) + // The tikv-only round stays disarmed: round 1's plan is already all TiKV, so + // rebuilding it under a TiKV-only restriction could only repeat that work. + requireRoundSkipped(t, tk, "tikv-only", allTiKVSQL, testkit.Rows("5")) + + // A missing TiFlash replica anywhere in the statement makes a fully-TiFlash + // plan impossible, so the tiflash-only round must not run. + noReplicaSQL := "select count(*) from alt_small_a join alt_no_replica" + + " on alt_small_a.a = alt_no_replica.a" + requireRoundSkipped(t, tk, "tiflash-only", noReplicaSQL, testkit.Rows("5")) + stmtCtx = tk.Session().GetSessionVars().StmtCtx + require.True(t, stmtCtx.AlternativeLogicalPlanMissingTiFlashPath) + + // Mirror direction: round 1's aggregation over the large table goes to + // TiFlash, so the all-TiKV alternative is the uncosted one. + tiFlashSQL := "select sum(b) from alt_large group by c" + tiFlashResult := testkit.Rows("1247500", "1248000", "1248500", "1249000", "1249500", + "1250000", "1250500", "1251000", "1251500", "1252000") + tk.MustQuery(tiFlashSQL).Sort().Check(tiFlashResult) + stmtCtx = tk.Session().GetSessionVars().StmtCtx + require.True(t, stmtCtx.AlternativeLogicalPlanReadsFromTiFlash) + requireRoundRan(t, tk, "tikv-only", tiFlashSQL) + + // No join and no aggregation: a wholesale engine switch has nothing to offer, + // and this gate is what keeps the extra rounds off the OLTP fast path. + scanOnlySQL := "select b from alt_small_a where b > 30" + requireRoundSkipped(t, tk, "tiflash-only", scanOnlySQL, testkit.Rows("40", "50")) + stmtCtx = tk.Session().GetSessionVars().StmtCtx + require.False(t, stmtCtx.AlternativeLogicalPlanHasJoinOrAgg) + + // An explicit READ_FROM_STORAGE hint pins the engine choice; the rounds must + // not run, or the cost comparison could override the hint. + hintSQL := "select /*+ read_from_storage(tikv[alt_small_a, alt_small_b]) */ count(*)" + + " from alt_small_a join alt_small_b on alt_small_a.a = alt_small_b.a" + requireRoundSkipped(t, tk, "tiflash-only", hintSQL, testkit.Rows("5")) + stmtCtx = tk.Session().GetSessionVars().StmtCtx + require.True(t, stmtCtx.AlternativeLogicalPlanHasStoreTypeHint) + + // Enforced MPP skips the rounds: its cost discount would make the cross-round + // comparison meaningless. + tk.MustExec("set @@tidb_enforce_mpp=1") + requireRoundSkipped(t, tk, "tiflash-only", allTiKVSQL, testkit.Rows("5")) + tk.MustExec("set @@tidb_enforce_mpp=0") + + // With the feature variable off, no engine round runs. + tk.MustExec("set @@tidb_opt_enable_alternative_logical_plans=off") + requireRoundSkipped(t, tk, "tiflash-only", allTiKVSQL, testkit.Rows("5")) + tk.MustExec("set @@tidb_opt_enable_alternative_logical_plans=on") + + // Let the rounds run to completion (no failpoint): planning and execution + // must succeed and the session's isolation read engines must be untouched. + tk.MustQuery(allTiKVSQL).Check(testkit.Rows("5")) + tk.MustQuery(tiFlashSQL).Sort().Check(tiFlashResult) + tk.MustQuery("select @@tidb_isolation_read_engines").Check(enginesBefore) +} diff --git a/pkg/planner/core/logical_plan_builder.go b/pkg/planner/core/logical_plan_builder.go index d37b0d418f2ee..df92a2f2d90b7 100644 --- a/pkg/planner/core/logical_plan_builder.go +++ b/pkg/planner/core/logical_plan_builder.go @@ -635,6 +635,9 @@ func setPreferredStoreType(ds *logicalop.DataSource, hintInfo *h.PlanHints) { ds.SCtx().GetSessionVars().StmtCtx.SetHintWarning(errMsg) } } + if ds.PreferStoreType != 0 { + ds.SCtx().GetSessionVars().StmtCtx.MarkAlternativeLogicalPlanHasStoreTypeHint() + } } // findJoinFullSchema walks through single-child wrapper operators (e.g., LogicalSelection @@ -5086,7 +5089,8 @@ func (b *PlanBuilder) buildDataSource(ctx context.Context, tn *ast.TableName, as // remain tikv access path to generate point get acceess path if existed // see detail in issue: https://github.com/pingcap/tidb/issues/39543 - if !(b.isForUpdateRead && b.ctx.GetSessionVars().TxnCtx.IsExplicit) { + isForUpdateReadInExplicitTxn := b.isForUpdateRead && b.ctx.GetSessionVars().TxnCtx.IsExplicit + if !isForUpdateReadInExplicitTxn { // Skip storage engine check for CreateView. if b.capFlag&canExpandAST == 0 { possiblePaths, err = util.FilterPathByIsolationRead(b.ctx, possiblePaths, tblName, dbName) @@ -5095,6 +5099,22 @@ func (b *PlanBuilder) buildDataSource(ctx context.Context, tn *ast.TableName, as } } } + if b.ctx.GetSessionVars().EnableAlternativeLogicalPlans { + // A table with no TiFlash access path rules out a fully-TiFlash plan for + // the whole statement, so the tiflash-only round is not worth building. + // FOR UPDATE reads in explicit transactions drop their TiFlash paths + // later during physical optimization, so treat them as missing one too. + hasTiFlashPath := false + for _, path := range possiblePaths { + if path.StoreType == kv.TiFlash { + hasTiFlashPath = true + break + } + } + if !hasTiFlashPath || isForUpdateReadInExplicitTxn { + b.ctx.GetSessionVars().StmtCtx.MarkAlternativeLogicalPlanMissingTiFlashPath() + } + } // Try to substitute generate column only if there is an index on generate column. for _, index := range tableInfo.Indices { diff --git a/pkg/planner/core/operator/physicalop/BUILD.bazel b/pkg/planner/core/operator/physicalop/BUILD.bazel index 8207d94885eeb..a01f96867f488 100644 --- a/pkg/planner/core/operator/physicalop/BUILD.bazel +++ b/pkg/planner/core/operator/physicalop/BUILD.bazel @@ -52,6 +52,7 @@ go_library( "physical_utils.go", "physical_window.go", "plan_clone_generated.go", + "plan_shape.go", "task.go", "task_base.go", "tiflash_predicate_push_down.go", diff --git a/pkg/planner/core/operator/physicalop/plan_shape.go b/pkg/planner/core/operator/physicalop/plan_shape.go new file mode 100644 index 0000000000000..290a1676f0c59 --- /dev/null +++ b/pkg/planner/core/operator/physicalop/plan_shape.go @@ -0,0 +1,102 @@ +// 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 physicalop + +import ( + "github.com/pingcap/tidb/pkg/kv" + "github.com/pingcap/tidb/pkg/planner/core/base" +) + +// PlanShape summarizes the properties of a physical plan that the alternative +// engine-restricted logical plan rounds arm on: which storage engines the plan +// reads from, and whether it contains a join or aggregation. +type PlanShape struct { + // ReadsFromTiKV is true when some access path in the plan reads from TiKV. + ReadsFromTiKV bool + // ReadsFromTiFlash is true when some access path in the plan reads from TiFlash. + ReadsFromTiFlash bool + // HasJoinOrAgg is true when the plan contains a join or an aggregation, + // including ones pushed down into a cop or MPP task. + HasJoinOrAgg bool +} + +// InspectPlanShape walks a physical plan tree and reports its PlanShape. +// +// The walk descends through storage readers into their pushed-down subtrees +// (TablePlan / IndexPlan / partial plans), because an aggregation or an MPP +// join frequently sits inside a cop or MPP task rather than above the reader. +func InspectPlanShape(plan base.PhysicalPlan) PlanShape { + var shape PlanShape + shape.inspect(plan) + return shape +} + +func (s *PlanShape) inspect(plan base.PhysicalPlan) { + if plan == nil { + return + } + if isJoinOrAgg(plan) { + s.HasJoinOrAgg = true + } + switch x := plan.(type) { + case *PhysicalTableReader: + // A table reader is the only reader that can read from TiFlash; every + // other reader below is an index access path, which TiFlash lacks. + if x.StoreType == kv.TiFlash { + s.ReadsFromTiFlash = true + } else { + s.ReadsFromTiKV = true + } + s.inspect(x.TablePlan) + return + case *PhysicalIndexReader: + s.ReadsFromTiKV = true + s.inspect(x.IndexPlan) + return + case *PhysicalIndexLookUpReader: + s.ReadsFromTiKV = true + s.inspect(x.IndexPlan) + s.inspect(x.TablePlan) + return + case *PhysicalIndexMergeReader: + s.ReadsFromTiKV = true + for _, partial := range x.PartialPlansRaw { + s.inspect(partial) + } + s.inspect(x.TablePlan) + return + case *PointGetPlan, *BatchPointGetPlan: + s.ReadsFromTiKV = true + return + case *PhysicalCTE: + s.inspect(x.SeedPlan) + s.inspect(x.RecurPlan) + } + for _, child := range plan.Children() { + s.inspect(child) + } +} + +// isJoinOrAgg reports whether the operator is a join or an aggregation. +// PhysicalApply is checked separately: it overrides PhysicalJoinImplement with +// a different signature and so deliberately does not satisfy base.PhysicalJoin. +func isJoinOrAgg(plan base.PhysicalPlan) bool { + switch plan.(type) { + case *PhysicalHashAgg, *PhysicalStreamAgg, *PhysicalApply: + return true + } + _, isJoin := plan.(base.PhysicalJoin) + return isJoin +} diff --git a/pkg/planner/optimize.go b/pkg/planner/optimize.go index ef5c08337b201..ee5b014b70e1a 100644 --- a/pkg/planner/optimize.go +++ b/pkg/planner/optimize.go @@ -32,6 +32,7 @@ import ( "github.com/pingcap/tidb/pkg/parser/ast" "github.com/pingcap/tidb/pkg/planner/core" "github.com/pingcap/tidb/pkg/planner/core/base" + "github.com/pingcap/tidb/pkg/planner/core/operator/physicalop" "github.com/pingcap/tidb/pkg/planner/core/resolve" "github.com/pingcap/tidb/pkg/planner/core/rule" "github.com/pingcap/tidb/pkg/planner/indexadvisor" @@ -626,29 +627,61 @@ func shouldTrySemiJoinRewriteRound(sessVars *variable.SessionVars) bool { !sessVars.EnableSemiJoinRewrite } +// shouldTryEngineRestrictedRounds holds the gates shared by the tikv-only and +// tiflash-only rounds. +// +// Round 1's bottom-up search picks a storage engine per DataSource through +// local cost comparisons, so a plan that is homogeneous across the whole +// statement is never costed as a unit. These rounds build that missing global +// alternative and let the normal strict-`<` cost comparison decide. +// +// They arm only for statements with a join or an aggregation: a wholesale +// engine switch has little to offer a plain scan, and this keeps the extra +// rounds off the OLTP fast path. They never run when a READ_FROM_STORAGE hint +// pins an engine (the cost comparison must not override an explicit user +// choice), nor under enforced MPP (its cost discount would make the +// cross-round comparison meaningless). +func shouldTryEngineRestrictedRounds(sessVars *variable.SessionVars) bool { + return sessVars.EnableAlternativeLogicalPlans && + sessVars.StmtCtx.AlternativeLogicalPlanHasJoinOrAgg && + !sessVars.StmtCtx.AlternativeLogicalPlanHasStoreTypeHint && + !sessVars.IsMPPEnforced() +} + +// shouldTryTiKVOnlyRound arms whenever round 1 read from TiFlash anywhere: only +// then does an all-TiKV plan differ from the plan already chosen. +func shouldTryTiKVOnlyRound(sessVars *variable.SessionVars) bool { + return shouldTryEngineRestrictedRounds(sessVars) && + sessVars.StmtCtx.AlternativeLogicalPlanReadsFromTiFlash +} + +// shouldTryTiFlashOnlyRound arms whenever round 1 read from TiKV anywhere, +// including the all-TiKV case: local per-table costing can prefer index access +// table by table even when a whole-statement MPP plan would win, and that +// alternative is never costed otherwise. It is skipped when any table lacks a +// TiFlash access path (no replica at all is the common case, and then a +// fully-TiFlash plan cannot be built) or when MPP execution is not allowed. +func shouldTryTiFlashOnlyRound(sessVars *variable.SessionVars) bool { + return shouldTryEngineRestrictedRounds(sessVars) && + sessVars.StmtCtx.AlternativeLogicalPlanReadsFromTiKV && + !sessVars.StmtCtx.AlternativeLogicalPlanMissingTiFlashPath && + sessVars.IsMPPAllowed() +} + // alternativeRound describes one alternative logical-plan round. // adjustFlag adjusts the optimization flags for the round. // enabled returns true when the round should be attempted. -// setup/cleanup optionally modify session state before/after plan building. +// setup optionally modifies session state before plan building and returns a +// cleanup that restores exactly what it changed. The saved values live in the +// returned closure, never in package state: optimize can run concurrently in +// different sessions, so each invocation must own its own saved state. type alternativeRound struct { name string adjustFlag func(uint64) uint64 enabled func(*variable.SessionVars) bool - setup func(*variable.SessionVars) - cleanup func(*variable.SessionVars) + setup func(*variable.SessionVars) (cleanup func()) } -// savedEnableCorrelateSubquery holds the pre-round value of -// EnableCorrelateSubquery so setup/cleanup can share it without a closure -// wrapper. Safe because optimize is single-threaded per session. -var savedEnableCorrelateSubquery bool - -// savedFTSLikeFallback holds the pre-round value of -// AlternativeLogicalPlanFTSLikeFallback so the fts-like-fallback round's -// setup/cleanup can restore it after running with the flag forced on. Safe -// because optimize is single-threaded per session. -var savedFTSLikeFallback bool - var alternativeRounds = [...]alternativeRound{ { name: "non-decorrelate", @@ -664,23 +697,19 @@ var alternativeRounds = [...]alternativeRound{ name: "correlate", adjustFlag: func(flag uint64) uint64 { return flag | rule.FlagCorrelate }, enabled: shouldTryCorrelateRound, - setup: func(sv *variable.SessionVars) { - savedEnableCorrelateSubquery = sv.EnableCorrelateSubquery + setup: func(sv *variable.SessionVars) func() { + saved := sv.EnableCorrelateSubquery sv.EnableCorrelateSubquery = true - }, - cleanup: func(sv *variable.SessionVars) { - sv.EnableCorrelateSubquery = savedEnableCorrelateSubquery + return func() { sv.EnableCorrelateSubquery = saved } }, }, { name: "semi-join-rewrite", enabled: shouldTrySemiJoinRewriteRound, - setup: func(sv *variable.SessionVars) { + setup: func(sv *variable.SessionVars) func() { sv.EnableSemiJoinRewrite = true - }, - cleanup: func(sv *variable.SessionVars) { // This round is enabled only when the original value is false. - sv.EnableSemiJoinRewrite = false + return func() { sv.EnableSemiJoinRewrite = false } }, }, { @@ -703,12 +732,32 @@ var alternativeRounds = [...]alternativeRound{ return sv.StmtCtx.AlternativeLogicalPlanFTSLikeFallback || sv.StmtCtx.AlternativeLogicalPlanHasPredicateContextMatch }, - setup: func(sv *variable.SessionVars) { - savedFTSLikeFallback = sv.StmtCtx.AlternativeLogicalPlanFTSLikeFallback + setup: func(sv *variable.SessionVars) func() { + saved := sv.StmtCtx.AlternativeLogicalPlanFTSLikeFallback sv.StmtCtx.AlternativeLogicalPlanFTSLikeFallback = true + return func() { sv.StmtCtx.AlternativeLogicalPlanFTSLikeFallback = saved } + }, + }, + { + // tikv-only: rebuild the plan with TiFlash removed from the isolation + // read engines so every table is read from TiKV. + name: "tikv-only", + enabled: shouldTryTiKVOnlyRound, + setup: func(sv *variable.SessionVars) func() { + saved := sv.IsolationReadEngines + sv.IsolationReadEngines = map[kv.StoreType]struct{}{kv.TiKV: {}, kv.TiDB: {}} + return func() { sv.IsolationReadEngines = saved } }, - cleanup: func(sv *variable.SessionVars) { - sv.StmtCtx.AlternativeLogicalPlanFTSLikeFallback = savedFTSLikeFallback + }, + { + // tiflash-only: mirror of tikv-only with only TiFlash readable, which + // lets the physical search build the MPP plan for the whole statement. + name: "tiflash-only", + enabled: shouldTryTiFlashOnlyRound, + setup: func(sv *variable.SessionVars) func() { + saved := sv.IsolationReadEngines + sv.IsolationReadEngines = map[kv.StoreType]struct{}{kv.TiFlash: {}, kv.TiDB: {}} + return func() { sv.IsolationReadEngines = saved } }, }, } @@ -794,6 +843,24 @@ func optimize(ctx context.Context, sctx planctx.PlanContext, node *resolve.NodeW return p, names, 0, nil } + // The engine-restricted rounds (tikv-only / tiflash-only) arm on the shape + // of round 1's chosen plan rather than on a build-time signal: which engines + // it actually read from decides which homogeneous alternative is still + // uncosted, and whether it joins or aggregates decides whether the rebuild + // is worth the extra optimization. + if needRestoreLogicalPlanCtx && bestPlan != nil { + shape := physicalop.InspectPlanShape(bestPlan) + if shape.ReadsFromTiKV { + sessVars.StmtCtx.MarkAlternativeLogicalPlanReadsFromTiKV() + } + if shape.ReadsFromTiFlash { + sessVars.StmtCtx.MarkAlternativeLogicalPlanReadsFromTiFlash() + } + if shape.HasJoinOrAgg { + sessVars.StmtCtx.MarkAlternativeLogicalPlanHasJoinOrAgg() + } + } + // Pre-compute which rounds are enabled based on the signals from the first // (default) build. This prevents signal leakage: alternative rounds rebuild // the plan and may set AlternativeLogicalPlan* signals as a side effect, @@ -821,8 +888,8 @@ func optimize(ctx context.Context, sctx planctx.PlanContext, node *resolve.NodeW // EnableCorrelateSubquery) is restored even if the round panics. func() { if round.setup != nil { - round.setup(sessVars) - defer round.cleanup(sessVars) + cleanup := round.setup(sessVars) + defer cleanup() } p, names, nonLogical, err = buildAndOptimizeLogicalPlanRound( ctx, diff --git a/pkg/sessionctx/stmtctx/stmtctx.go b/pkg/sessionctx/stmtctx/stmtctx.go index 42982b6842ed1..9b0cc991d4eef 100644 --- a/pkg/sessionctx/stmtctx/stmtctx.go +++ b/pkg/sessionctx/stmtctx/stmtctx.go @@ -510,6 +510,31 @@ type StatementContext struct { // uses this to enable the fts-like-fallback round for cost competition even // when round 1's native plan is executable. AlternativeLogicalPlanHasPredicateContextMatch bool + // AlternativeLogicalPlanReadsFromTiKV indicates that round 1's chosen + // physical plan reads from TiKV somewhere. The round driver uses it to arm + // the tiflash-only round: only then does a fully-TiFlash plan differ from + // round 1's plan, so only then is it worth costing as a whole. + AlternativeLogicalPlanReadsFromTiKV bool + // AlternativeLogicalPlanReadsFromTiFlash indicates that round 1's chosen + // physical plan reads from TiFlash somewhere. The round driver uses it to + // arm the tikv-only round, mirroring AlternativeLogicalPlanReadsFromTiKV. + AlternativeLogicalPlanReadsFromTiFlash bool + // AlternativeLogicalPlanHasJoinOrAgg indicates that round 1's chosen + // physical plan contains a join or aggregation (including ones pushed into + // cop or MPP tasks). The engine-restricted rounds arm only then: a wholesale + // engine switch has little to offer a plain scan, and this gate keeps the + // extra optimization rounds off the OLTP fast path. + AlternativeLogicalPlanHasJoinOrAgg bool + // AlternativeLogicalPlanMissingTiFlashPath indicates that some DataSource in + // round 1's build ended up without any TiFlash access path (no replica, a + // system table, or a FOR UPDATE read). A fully-TiFlash plan is impossible, + // so the tiflash-only round is skipped. + AlternativeLogicalPlanMissingTiFlashPath bool + // AlternativeLogicalPlanHasStoreTypeHint indicates that some DataSource in + // round 1's build carries an explicit engine preference from a + // READ_FROM_STORAGE hint. The engine-restricted rounds are skipped so the + // cost comparison cannot override the user's explicit engine choice. + AlternativeLogicalPlanHasStoreTypeHint bool // FTSFunctionIsUsed indicates that FTS_MATCH_WORD() appears in the current // statement, allowing the optimizer to run FTS-specific validation and // rewrite rules only when needed. @@ -691,6 +716,11 @@ func (sc *StatementContext) ResetAlternativeLogicalPlanSignals() { sc.AlternativeLogicalPlanHasPredicateContextMatch = false sc.AlternativeLogicalPlanPreferCorrelate = false sc.AlternativeLogicalPlanSemiJoinRewrite = false + sc.AlternativeLogicalPlanReadsFromTiKV = false + sc.AlternativeLogicalPlanReadsFromTiFlash = false + sc.AlternativeLogicalPlanHasJoinOrAgg = false + sc.AlternativeLogicalPlanMissingTiFlashPath = false + sc.AlternativeLogicalPlanHasStoreTypeHint = false sc.FTSFunctionIsUsed = false } @@ -725,6 +755,36 @@ func (sc *StatementContext) MarkAlternativeLogicalPlanSemiJoinRewrite() { sc.AlternativeLogicalPlanSemiJoinRewrite = true } +// MarkAlternativeLogicalPlanReadsFromTiKV records that round 1's chosen physical +// plan reads from TiKV. +func (sc *StatementContext) MarkAlternativeLogicalPlanReadsFromTiKV() { + sc.AlternativeLogicalPlanReadsFromTiKV = true +} + +// MarkAlternativeLogicalPlanReadsFromTiFlash records that round 1's chosen +// physical plan reads from TiFlash. +func (sc *StatementContext) MarkAlternativeLogicalPlanReadsFromTiFlash() { + sc.AlternativeLogicalPlanReadsFromTiFlash = true +} + +// MarkAlternativeLogicalPlanHasJoinOrAgg records that round 1's chosen physical +// plan contains a join or an aggregation. +func (sc *StatementContext) MarkAlternativeLogicalPlanHasJoinOrAgg() { + sc.AlternativeLogicalPlanHasJoinOrAgg = true +} + +// MarkAlternativeLogicalPlanMissingTiFlashPath records that a DataSource in the +// current build has no TiFlash access path, making a fully-TiFlash plan impossible. +func (sc *StatementContext) MarkAlternativeLogicalPlanMissingTiFlashPath() { + sc.AlternativeLogicalPlanMissingTiFlashPath = true +} + +// MarkAlternativeLogicalPlanHasStoreTypeHint records that a DataSource in the +// current build carries an explicit READ_FROM_STORAGE engine preference. +func (sc *StatementContext) MarkAlternativeLogicalPlanHasStoreTypeHint() { + sc.AlternativeLogicalPlanHasStoreTypeHint = true +} + // CtxID returns the context id of the statement func (sc *StatementContext) CtxID() uint64 { return sc.ctxID