From 07c13266c00273930170fefd81f43f9253e81b05 Mon Sep 17 00:00:00 2001
From: Ti Chi Robot <ti-community-prow-bot@tidb.io>
Date: Thu, 17 Aug 2023 15:32:03 +0800
Subject: [PATCH] planner: fix index heuristic rule will prune out hint
 preferred tiflash path (#46102) (#46144)

close pingcap/tidb#40146
---
 planner/core/enforce_mpp_test.go                 |  7 +++++++
 planner/core/stats.go                            |  8 +++++++-
 planner/core/testdata/enforce_mpp_suite_in.json  |  3 ++-
 planner/core/testdata/enforce_mpp_suite_out.json | 11 +++++++++++
 4 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/planner/core/enforce_mpp_test.go b/planner/core/enforce_mpp_test.go
index dfd0a4fadbe3a..c580594deb638 100644
--- a/planner/core/enforce_mpp_test.go
+++ b/planner/core/enforce_mpp_test.go
@@ -95,6 +95,7 @@ func TestEnforceMPP(t *testing.T) {
 	tk.MustExec("drop table if exists t")
 	tk.MustExec("create table t(a int, b int)")
 	tk.MustExec("create index idx on t(a)")
+	tk.MustExec("CREATE TABLE `s` (\n `a` int(11) DEFAULT NULL,\n `b` int(11) DEFAULT NULL,\n `c` int(11) DEFAULT NULL,\n `d` int(11) DEFAULT NULL,\n UNIQUE KEY `a` (`a`),\n KEY `ii` (`a`,`b`)\n)")
 
 	// Default RPC encoding may cause statistics explain result differ and then the test unstable.
 	tk.MustExec("set @@tidb_enable_chunk_rpc = on")
@@ -111,6 +112,12 @@ func TestEnforceMPP(t *testing.T) {
 				Available: true,
 			}
 		}
+		if tblInfo.Name.L == "s" {
+			tblInfo.TiFlashReplica = &model.TiFlashReplicaInfo{
+				Count:     1,
+				Available: true,
+			}
+		}
 	}
 
 	var input []string
diff --git a/planner/core/stats.go b/planner/core/stats.go
index be1c66d5528c4..b91a9266340ed 100644
--- a/planner/core/stats.go
+++ b/planner/core/stats.go
@@ -363,8 +363,14 @@ func (ds *DataSource) derivePathStatsAndTryHeuristics() error {
 			selected = uniqueBest
 		}
 	}
-	// If some path matches a heuristic rule, just remove other possible paths
+	// heuristic rule pruning other path should consider hint prefer.
+	// If no hints and some path matches a heuristic rule, just remove other possible paths.
 	if selected != nil {
+		// if user wanna tiFlash read, while current heuristic choose a TiKV path. so we shouldn't prune other paths.
+		keep := ds.preferStoreType&preferTiFlash != 0 && selected.StoreType != kv.TiFlash
+		if keep {
+			return nil
+		}
 		ds.possibleAccessPaths[0] = selected
 		ds.possibleAccessPaths = ds.possibleAccessPaths[:1]
 		if ds.ctx.GetSessionVars().StmtCtx.InVerboseExplain {
diff --git a/planner/core/testdata/enforce_mpp_suite_in.json b/planner/core/testdata/enforce_mpp_suite_in.json
index fc6d31dec9da9..3b4d09f0c7e65 100644
--- a/planner/core/testdata/enforce_mpp_suite_in.json
+++ b/planner/core/testdata/enforce_mpp_suite_in.json
@@ -21,7 +21,8 @@
       "set @@tidb_enforce_mpp=1;",
       "explain format='verbose' select count(*) from t where a=1",
       "explain format='verbose' select /*+ read_from_storage(tikv[t]) */ count(*) from t where a=1",
-      "explain format='verbose' select /*+ read_from_storage(tiflash[t]) */ count(*) from t where a=1"
+      "explain format='verbose' select /*+ read_from_storage(tiflash[t]) */ count(*) from t where a=1",
+      "explain select /*+ READ_FROM_STORAGE(TIFLASH[s]) */ a from s where a = 10 and b is null; -- index path huristic rule will prune tiflash path"
     ]
   },
   {
diff --git a/planner/core/testdata/enforce_mpp_suite_out.json b/planner/core/testdata/enforce_mpp_suite_out.json
index e62c28ac59cb5..045683345af3b 100644
--- a/planner/core/testdata/enforce_mpp_suite_out.json
+++ b/planner/core/testdata/enforce_mpp_suite_out.json
@@ -173,6 +173,17 @@
           "      └─TableFullScan_25 10000.00 928000.00 batchCop[tiflash] table:t keep order:false, stats:pseudo"
         ],
         "Warn": null
+      },
+      {
+        "SQL": "explain select /*+ READ_FROM_STORAGE(TIFLASH[s]) */ a from s where a = 10 and b is null; -- index path huristic rule will prune tiflash path",
+        "Plan": [
+          "TableReader_12 0.10 root  data:ExchangeSender_11",
+          "└─ExchangeSender_11 0.10 mpp[tiflash]  ExchangeType: PassThrough",
+          "  └─Projection_5 0.10 mpp[tiflash]  test.s.a",
+          "    └─Selection_10 0.10 mpp[tiflash]  eq(test.s.a, 10), isnull(test.s.b)",
+          "      └─TableFullScan_9 10000.00 mpp[tiflash] table:s keep order:false, stats:pseudo"
+        ],
+        "Warn": null
       }
     ]
   },