From 4436b16307e252c0aa3e6e60cac523d5bbd7e6d3 Mon Sep 17 00:00:00 2001
From: guo-shaoge <shaoge1994@163.com>
Date: Mon, 13 Feb 2023 17:28:03 +0800
Subject: [PATCH] planner: fix tiflash cannot find generated column (#41261)

close pingcap/tidb#40663
---
 expression/integration_test.go |  9 +++++++--
 parser/mysql/type.go           |  1 +
 planner/core/find_best_task.go | 12 +++---------
 3 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/expression/integration_test.go b/expression/integration_test.go
index 55c8f389a5df3..68507621f1e26 100644
--- a/expression/integration_test.go
+++ b/expression/integration_test.go
@@ -3743,18 +3743,23 @@ func TestShardIndexOnTiFlash(t *testing.T) {
 			}
 		}
 	}
+	tk.MustExec("set @@session.tidb_isolation_read_engines = 'tiflash'")
 	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")
+		if strings.Contains(line, "TableFullScan") {
+			require.Contains(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")
+		if strings.Contains(line, "TableFullScan") {
+			require.NotContains(t, line, "mpp[tiflash]")
+		}
 	}
 }
 
diff --git a/parser/mysql/type.go b/parser/mysql/type.go
index f79be8ab30d96..8a2531d870d3e 100644
--- a/parser/mysql/type.go
+++ b/parser/mysql/type.go
@@ -74,6 +74,7 @@ const (
 	PreventNullInsertFlag uint = 1 << 20 /* Prevent this Field from inserting NULL values */
 	EnumSetAsIntFlag      uint = 1 << 21 /* Internal: Used for inferring enum eval type. */
 	DropColumnIndexFlag   uint = 1 << 22 /* Internal: Used for indicate the column is being dropped with index */
+	GeneratedColumnFlag   uint = 1 << 23 /* Internal: TiFlash will check this flag and add a placeholder for this column */
 )
 
 // TypeInt24 bounds.
diff --git a/planner/core/find_best_task.go b/planner/core/find_best_task.go
index b75942f3a2639..bf2b0596a3e0d 100644
--- a/planner/core/find_best_task.go
+++ b/planner/core/find_best_task.go
@@ -1973,15 +1973,9 @@ func (ds *DataSource) convertToTableScan(prop *property.PhysicalProperty, candid
 		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
+		for _, col := range ts.Columns {
+			if col.IsGenerated() && !col.GeneratedStored {
+				col.AddFlag(mysql.GeneratedColumnFlag)
 			}
 		}
 	}