planner: forbid [batch] point get with update read choose tiflash as …

…datasource (#39553) (#39642)

close #39543

planner/core/find_best_task.go

@@ -806,7 +806,30 @@ func (ds *DataSource) findBestTask(prop *property.PhysicalProperty, planCounter
 		planCounter.Dec(1)
 		return nil, 1, nil
 	}
-
+	if ds.isForUpdateRead && ds.ctx.GetSessionVars().TxnCtx.IsExplicit {
+		hasPointGetPath := false
+		for _, path := range ds.possibleAccessPaths {
+			if ds.isPointGetPath(path) {
+				hasPointGetPath = true
+				break
+			}
+		}
+		tblName := ds.tableInfo.Name
+		ds.possibleAccessPaths, err = filterPathByIsolationRead(ds.ctx, ds.possibleAccessPaths, tblName, ds.DBName)
+		if err != nil {
+			return nil, 1, err
+		}
+		if hasPointGetPath {
+			newPaths := make([]*util.AccessPath, 0)
+			for _, path := range ds.possibleAccessPaths {
+				// if the path is the point get range path with for update lock, we should forbid tiflash as it's store path (#39543)
+				if path.StoreType != kv.TiFlash {
+					newPaths = append(newPaths, path)
+				}
+			}
+			ds.possibleAccessPaths = newPaths
+		}
+	}
 	t = ds.getTask(prop)
 	if t != nil {
 		cntPlan = 1
@@ -914,13 +937,6 @@ func (ds *DataSource) findBestTask(prop *property.PhysicalProperty, planCounter
 			}, cntPlan, nil
 		}
 
-		// if the path is the point get range path with for update lock, we should forbid tiflash as it's store path.
-		if path.StoreType == kv.TiFlash && ds.isForUpdateRead && ds.ctx.GetSessionVars().TxnCtx.IsExplicit {
-			if ds.isPointGetConditions() {
-				continue
-			}
-		}
-
 		canConvertPointGet := len(path.Ranges) > 0 && path.StoreType == kv.TiKV && ds.isPointGetConvertableSchema()
 
 		if canConvertPointGet && expression.MaybeOverOptimized4PlanCache(ds.ctx, path.AccessConds) {
@@ -1912,64 +1928,34 @@ func (s *LogicalIndexScan) GetPhysicalIndexScan(_ *expression.Schema, stats *pro
 	return is
 }
 
-// isPointGetConditions indicates whether the conditions are point-get-able.
+// isPointGetPath indicates whether the conditions are point-get-able.
 // eg: create table t(a int, b int,c int unique, primary (a,b))
 // select * from t where a = 1 and b = 1 and c =1;
-// the datasource can access by primary key(a,b) or unique key (c) which are both point-get-able
-func (ds *DataSource) isPointGetConditions() bool {
-	t, _ := ds.is.TableByID(ds.physicalTableID)
-	columns := map[string]struct{}{}
-	for _, cond := range ds.allConds {
-		s, ok := cond.(*expression.ScalarFunction)
-		if !ok {
+// the datasource can access by primary key(a,b) or unique key c which are both point-get-able
+func (ds *DataSource) isPointGetPath(path *util.AccessPath) bool {
+	if len(path.Ranges) < 1 {
+		return false
+	}
+	if !path.IsIntHandlePath {
+		if path.Index == nil {
 			return false
 		}
-		if s.FuncName.L != ast.EQ || (s.FuncName.L == ast.In && len(s.GetArgs()) != 2) {
+		if !path.Index.Unique || path.Index.HasPrefixIndex() {
 			return false
 		}
-		arg0 := s.GetArgs()[0]
-		arg1 := s.GetArgs()[1]
-		_, ok1 := arg0.(*expression.Constant)
-		col, ok2 := arg1.(*expression.Column)
-		if ok1 && ok2 {
-			columns[t.Meta().FindColumnNameByID(col.ID)] = struct{}{}
-			continue
-		}
-		col, ok1 = arg0.(*expression.Column)
-		_, ok2 = arg1.(*expression.Constant)
-		if ok1 && ok2 {
-			columns[t.Meta().FindColumnNameByID(col.ID)] = struct{}{}
-			continue
-		}
-	}
-	return ds.findPKOrUniqueIndexMatchColumns(columns)
-}
-
-func (ds *DataSource) findPKOrUniqueIndexMatchColumns(columns map[string]struct{}) bool {
-	for _, idx := range ds.tableInfo.Indices {
-		if !idx.Unique && !idx.Primary {
-			continue
-		}
-		if idx.HasPrefixIndex() {
-			continue
-		}
-		if len(idx.Columns) > len(columns) {
-			continue
-		}
-		flag := true
-		for _, idxCol := range idx.Columns {
-			_, ok := columns[idxCol.Name.String()]
-			if !ok {
-				flag = false
-				break
+		idxColsLen := len(path.Index.Columns)
+		for _, ran := range path.Ranges {
+			if len(ran.LowVal) != idxColsLen {
+				return false
 			}
 		}
-		if !flag {
-			continue
+	}
+	for _, ran := range path.Ranges {
+		if !ran.IsPointNonNullable(ds.ctx) {
+			return false
 		}
-		return true
 	}
-	return false
+	return true
 }
 
 // convertToTableScan converts the DataSource to table scan.

planner/core/integration_test.go

@@ -7638,23 +7638,33 @@ func TestPointGetWithSelectLock(t *testing.T) {
 	// Set the hacked TiFlash replica for explain tests.
 	tbl1.Meta().TiFlashReplica = &model.TiFlashReplicaInfo{Count: 1, Available: true}
 
+	sqls := []string{
+		"explain select a, b from t where (a = 1 and b = 2) or (a =2 and b = 1) for update;",
+		"explain select a, b from t where a = 1 and b = 2 for update;",
+		"explain select c, d from t1 where c = 1 for update;",
+		"explain select c, d from t1 where c = 1 and d = 1 for update;",
+		"explain select c, d from t1 where (c = 1 or c = 2 )and d = 1 for update;",
+		"explain select c, d from t1 where c in (1,2,3,4) for update;",
+	}
 	tk.MustExec("set @@tidb_enable_tiflash_read_for_write_stmt = on;")
 	tk.MustExec("set @@tidb_isolation_read_engines='tidb,tiflash';")
 	tk.MustExec("begin;")
-	// assert point get condition with txn commit and tiflash store
-	tk.MustGetErrMsg("explain select a, b from t where a = 1 and b = 2 for update;", "[planner:1815]Internal : Can't find a proper physical plan for this query")
-	tk.MustGetErrMsg("explain select c, d from t1 where c = 1 for update;", "[planner:1815]Internal : Can't find a proper physical plan for this query")
-	tk.MustGetErrMsg("explain select c, d from t1 where c = 1 and d = 1 for update;", "[planner:1815]Internal : Can't find a proper physical plan for this query")
-	tk.MustQuery("explain select a, b from t where a = 1 for update;")
-	tk.MustQuery("explain select c, d from t1 where c > 1 for update;")
-	tk.MustExec("set tidb_isolation_read_engines='tidb,tikv,tiflash';")
-	tk.MustQuery("explain select a, b from t where a = 1 and b = 2 for update;")
-	tk.MustQuery("explain select c, d from t1 where c = 1 for update;")
+	// assert point get / batch point get can't work with tiflash in interaction txn
+	for _, sql := range sqls {
+		err = tk.ExecToErr(sql)
+		require.Error(t, err)
+	}
+	// assert point get / batch point get can work with tikv in interaction txn
+	tk.MustExec("set @@tidb_isolation_read_engines='tidb,tikv,tiflash';")
+	for _, sql := range sqls {
+		tk.MustQuery(sql)
+	}
 	tk.MustExec("commit")
-	tk.MustExec("set tidb_isolation_read_engines='tidb,tiflash';")
-	// assert point get condition with auto commit and tiflash store
-	tk.MustQuery("explain select a, b from t where  a = 1 and b = 2 for update;")
-	tk.MustQuery("explain select c, d from t1 where c = 1 for update;")
+	// assert point get / batch point get can work with tiflash in auto commit
+	tk.MustExec("set @@tidb_isolation_read_engines='tidb,tiflash';")
+	for _, sql := range sqls {
+		tk.MustQuery(sql)
+	}
 }
 
 func TestTableRangeFallback(t *testing.T) {

planner/core/logical_plan_builder.go

@@ -4494,11 +4494,15 @@ func (b *PlanBuilder) buildDataSource(ctx context.Context, tn *ast.TableName, as
 		return nil, ErrPartitionClauseOnNonpartitioned
 	}
 
-	// Skip storage engine check for CreateView.
-	if b.capFlag&canExpandAST == 0 {
-		possiblePaths, err = filterPathByIsolationRead(b.ctx, possiblePaths, tblName, dbName)
-		if err != nil {
-			return nil, err
+	// remain tikv access path to generate point get acceess path if existed
+	// see detail in issue: https://github.com/pingcap/tidb/issues/39543
+	if !(b.isForUpdateRead && b.ctx.GetSessionVars().TxnCtx.IsExplicit) {
+		// Skip storage engine check for CreateView.
+		if b.capFlag&canExpandAST == 0 {
+			possiblePaths, err = filterPathByIsolationRead(b.ctx, possiblePaths, tblName, dbName)
+			if err != nil {
+				return nil, err
+			}
 		}
 	}