From 29fbcdf3da0121acf4e773eb802755231c1932e5 Mon Sep 17 00:00:00 2001
From: ti-srebot <66930949+ti-srebot@users.noreply.github.com>
Date: Tue, 2 Aug 2022 11:18:06 +0800
Subject: [PATCH] planner: check virtual column for tiflash (#36771) (#36786)

close pingcap/tiflash#5513, close pingcap/tidb#36773
---
 expression/integration_test.go | 37 ++++++++++++++++++++++++++++++++++
 planner/core/find_best_task.go | 13 ++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/expression/integration_test.go b/expression/integration_test.go
index 8b2aa0139347f..cf1ab5efa08e5 100644
--- a/expression/integration_test.go
+++ b/expression/integration_test.go
@@ -3716,6 +3716,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 1d97e73cf8541..79b3c6a2d5113 100644
--- a/planner/core/find_best_task.go
+++ b/planner/core/find_best_task.go
@@ -1882,6 +1882,19 @@ 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 {
+			// In theory, TiFlash does not support virtual expr, but in non-mpp mode, if the cop request only contain table scan, then
+			// TiDB will fill the virtual column after decoding the cop response(executor.FillVirtualColumnValue), that is to say, the virtual
+			// columns in Cop request is just a placeholder, so TiFlash can support virtual column in cop request mode. However, virtual column
+			// with TiDBShard is special, it can be added using create index statement, TiFlash's ddl does not handle create index statement, so
+			// there is a chance that the TiDBShard's virtual column is not seen by TiFlash, in this case, TiFlash will throw column not found error
+			if ds.containExprPrefixUk && expression.GcColumnExprIsTidbShard(col.VirtualExpr) {
+				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