From 5ab6a784dd042c9b1aecf8de3f72d16fb249b4f8 Mon Sep 17 00:00:00 2001
From: xufei <xufeixw@mail.ustc.edu.cn>
Date: Mon, 1 Aug 2022 17:26:12 +0800
Subject: [PATCH] check virtual column for tiflash

Signed-off-by: xufei <xufeixw@mail.ustc.edu.cn>
---
 expression/integration_test.go | 37 ++++++++++++++++++++++++++++++++++
 planner/core/find_best_task.go | 14 +++++++------
 2 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/expression/integration_test.go b/expression/integration_test.go
index 2061aab4be26b..1be779cdaca14 100644
--- a/expression/integration_test.go
+++ b/expression/integration_test.go
@@ -3715,6 +3715,43 @@ func TestIssue16973(t *testing.T) {
 		"AND t1.status IN (2,6,10) AND timestampdiff(month, t2.begin_time, date'2020-05-06') = 0;").Check(testkit.Rows("1"))
 }
 
+func TestShardIndexOnTiFlash(t *testing.T) {
+	store, clean := testkit.CreateMockStore(t)
+	defer clean()
+
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t(id int primary key clustered, a int, b int, unique key uk_expr((tidb_shard(a)),a))")
+
+	// Create virtual tiflash replica info.
+	dom := domain.GetDomain(tk.Session())
+	is := dom.InfoSchema()
+	db, exists := is.SchemaByName(model.NewCIStr("test"))
+	require.True(t, exists)
+	for _, tblInfo := range db.Tables {
+		if tblInfo.Name.L == "t" {
+			tblInfo.TiFlashReplica = &model.TiFlashReplicaInfo{
+				Count:     1,
+				Available: true,
+			}
+		}
+	}
+	tk.MustExec("set @@session.tidb_enforce_mpp = 1")
+	rows := tk.MustQuery("explain select max(b) from t").Rows()
+	for _, row := range rows {
+		line := fmt.Sprintf("%v", row)
+		require.NotContains(t, line, "tiflash")
+	}
+	tk.MustExec("set @@session.tidb_enforce_mpp = 0")
+	tk.MustExec("set @@session.tidb_allow_mpp = 0")
+	rows = tk.MustQuery("explain select max(b) from t").Rows()
+	for _, row := range rows {
+		line := fmt.Sprintf("%v", row)
+		require.NotContains(t, line, "tiflash")
+	}
+}
+
 func TestExprPushdownBlacklist(t *testing.T) {
 	store, clean := testkit.CreateMockStore(t)
 	defer clean()
diff --git a/planner/core/find_best_task.go b/planner/core/find_best_task.go
index ac7790fc8ba99..35b4d7c5dea54 100644
--- a/planner/core/find_best_task.go
+++ b/planner/core/find_best_task.go
@@ -1882,6 +1882,14 @@ func (ds *DataSource) convertToTableScan(prop *property.PhysicalProperty, candid
 		// TiFlash fast mode(https://github.com/pingcap/tidb/pull/35851) does not keep order in TableScan
 		return invalidTask, nil
 	}
+	if ts.StoreType == kv.TiFlash {
+		for _, col := range ts.schema.Columns {
+			if col.VirtualExpr != nil {
+				ds.SCtx().GetSessionVars().RaiseWarningWhenMPPEnforced("MPP mode may be blocked because column `" + col.OrigName + "` is a virtual column which is not supported now.")
+				return invalidTask, nil
+			}
+		}
+	}
 	if prop.TaskTp == property.MppTaskType {
 		if ts.KeepOrder {
 			return invalidTask, nil
@@ -1891,12 +1899,6 @@ func (ds *DataSource) convertToTableScan(prop *property.PhysicalProperty, candid
 			ds.SCtx().GetSessionVars().RaiseWarningWhenMPPEnforced("MPP mode may be blocked because table `" + ds.tableInfo.Name.O + "`is a partition table which is not supported when `@@tidb_partition_prune_mode=static`.")
 			return invalidTask, nil
 		}
-		for _, col := range ts.schema.Columns {
-			if col.VirtualExpr != nil {
-				ds.SCtx().GetSessionVars().RaiseWarningWhenMPPEnforced("MPP mode may be blocked because column `" + col.OrigName + "` is a virtual column which is not supported now.")
-				return invalidTask, nil
-			}
-		}
 		mppTask := &mppTask{
 			p:      ts,
 			cst:    cost,