planner: check virtual column for tiflash (#36771) (#36785)

close pingcap/tiflash#5513, close #36773

expression/integration_test.go

@@ -3688,6 +3688,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()

planner/core/find_best_task.go

@@ -1873,6 +1873,19 @@ func (ds *DataSource) convertToTableScan(prop *property.PhysicalProperty, candid
  if ts.KeepOrder && ts.Desc && ts.StoreType == kv.TiFlash {
  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