ddl: fix invalid index on multi-layer virtual columns

ddl/db_integration_test.go

@@ -797,6 +797,36 @@ func (s *testIntegrationSuite5) TestModifyingColumnOption(c *C) {
 	assertErrCode("alter table t2 modify column c int references t1(a)", errMsg)
 }
 
+func (s *testIntegrationSuite1) TestIndexOnMultipleGeneratedColumn(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+
+	tk.MustExec("create database if not exists test")
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t (a int, b int as (a + 1), c int as (b + 1))")
+	tk.MustExec("insert into t (a) values (1)")
+	tk.MustExec("create index idx on t (c)")
+	tk.MustQuery("select * from t where c > 1").Check(testkit.Rows("1 2 3"))
+	tk.MustExec("admin check table t")
+
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t (a int, b int as (a + 1), c int as (b + 1), d int as (c + 1))")
+	tk.MustExec("insert into t (a) values (1)")
+	tk.MustExec("create index idx on t (d)")
+	tk.MustQuery("select * from t where d > 2").Check(testkit.Rows("1 2 3 4"))
+	tk.MustExec("admin check table t")
+
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t(a bigint, b bigint as (a+1) virtual, c bigint as (b+1) virtual)")
+	tk.MustExec("alter table t add index idx_b(b)")
+	tk.MustExec("alter table t add index idx_c(c)")
+	tk.MustExec("insert into t(a) values(1)")
+	tk.MustExec("alter table t add column(d bigint as (c+1) virtual)")
+	tk.MustExec("alter table t add index idx_d(d)")
+	tk.MustQuery("select * from t where d > 2").Check(testkit.Rows("1 2 3 4"))
+	tk.MustExec("admin check table t")
+}
+
 func (s *testIntegrationSuite2) TestCaseInsensitiveCharsetAndCollate(c *C) {
 	tk := testkit.NewTestKit(c, s.store)
 

ddl/index.go

@@ -908,27 +908,48 @@ func (w *addIndexWorker) run(d *ddlCtx) {
 
 func makeupDecodeColMap(sessCtx sessionctx.Context, t table.Table, indexInfo *model.IndexInfo) (map[int64]decoder.Column, error) {
 	cols := t.Cols()
-	decodeColMap := make(map[int64]decoder.Column, len(indexInfo.Columns))
-	for _, v := range indexInfo.Columns {
-		col := cols[v.Offset]
-		tpExpr := decoder.Column{
-			Col: col,
-		}
+	indexedCols := make([]*table.Column, len(indexInfo.Columns))
+	for i, v := range indexInfo.Columns {
+		indexedCols[i] = cols[v.Offset]
+	}
+
+	indexedColsCpy := make([]*table.Column, len(indexedCols))
+	copy(indexedColsCpy, indexedCols)
+	decodeColMap, err := buildFullDecodeColMap(indexedColsCpy, sessCtx, t)
+	if err != nil {
+		return nil, err
+	}
+
+	decoder.SubstituteGenColsInDecodeColMap(decodeColMap)
+	decoder.RemoveUnusedVirtualCols(decodeColMap, indexedCols)
+
+	return decodeColMap, nil
+}
+
+func buildFullDecodeColMap(pendingCols []*table.Column, sessCtx sessionctx.Context, t table.Table) (map[int64]decoder.Column, error) {
+	decodeColMap := make(map[int64]decoder.Column, len(pendingCols))
+	for i := 0; i < len(pendingCols); i++ {
+		col := pendingCols[i]
 		if col.IsGenerated() && !col.GeneratedStored {
-			for _, c := range cols {
+			// Find depended columns and put them into pendingCols.
+			for _, c := range t.Cols() {
 				if _, ok := col.Dependences[c.Name.L]; ok {
-					decodeColMap[c.ID] = decoder.Column{
-						Col: c,
-					}
+					pendingCols = append(pendingCols, c)
 				}
 			}
 			e, err := expression.ParseSimpleExprCastWithTableInfo(sessCtx, col.GeneratedExprString, t.Meta(), &col.FieldType)
 			if err != nil {
 				return nil, errors.Trace(err)
 			}
-			tpExpr.GenExpr = e
+			decodeColMap[col.ID] = decoder.Column{
+				Col:     col,
+				GenExpr: e,
+			}
+		} else {
+			decodeColMap[col.ID] = decoder.Column{
+				Col: col,
+			}
 		}
-		decodeColMap[col.ID] = tpExpr
 	}
 	return decodeColMap, nil
 }

util/admin/admin.go

@@ -589,27 +589,37 @@ func CompareTableRecord(sessCtx sessionctx.Context, txn kv.Transaction, t table.
 }
 
 func makeRowDecoder(t table.Table, decodeCol []*table.Column, genExpr map[model.TableColumnID]expression.Expression) *decoder.RowDecoder {
-	cols := t.Cols()
-	tblInfo := t.Meta()
-	decodeColsMap := make(map[int64]decoder.Column, len(decodeCol))
-	for _, v := range decodeCol {
-		col := cols[v.Offset]
-		tpExpr := decoder.Column{
-			Col: col,
-		}
+	decodeColCpy := make([]*table.Column, len(decodeCol))
+	copy(decodeColCpy, decodeCol)
+	decodeColsMap := buildFullDecodeColMap(decodeColCpy, t, genExpr)
+
+	decoder.SubstituteGenColsInDecodeColMap(decodeColsMap)
+	decoder.RemoveUnusedVirtualCols(decodeColsMap, decodeCol)
+	return decoder.NewRowDecoder(t, decodeColsMap)
+}
+
+func buildFullDecodeColMap(pendingCols []*table.Column, t table.Table, genExpr map[model.TableColumnID]expression.Expression) map[int64]decoder.Column {
+	decodeColMap := make(map[int64]decoder.Column, len(pendingCols))
+	for i := 0; i < len(pendingCols); i++ {
+		col := pendingCols[i]
 		if col.IsGenerated() && !col.GeneratedStored {
-			for _, c := range cols {
+			// Find depended columns and put them into pendingCols.
+			for _, c := range t.Cols() {
 				if _, ok := col.Dependences[c.Name.L]; ok {
-					decodeColsMap[c.ID] = decoder.Column{
-						Col: c,
-					}
+					pendingCols = append(pendingCols, c)
 				}
 			}
-			tpExpr.GenExpr = genExpr[model.TableColumnID{TableID: tblInfo.ID, ColumnID: col.ID}]
+			decodeColMap[col.ID] = decoder.Column{
+				Col:     col,
+				GenExpr: genExpr[model.TableColumnID{TableID: t.Meta().ID, ColumnID: col.ID}],
+			}
+		} else {
+			decodeColMap[col.ID] = decoder.Column{
+				Col: col,
+			}
 		}
-		decodeColsMap[col.ID] = tpExpr
 	}
-	return decoder.NewRowDecoder(t, decodeColsMap)
+	return decodeColMap
 }
 
 // genExprs use to calculate generated column value.

util/rowDecoder/decoder.go

@@ -14,8 +14,10 @@
 package decoder
 
 import (
+	"sort"
 	"time"
 
+	"github.com/pingcap/parser/ast"
 	"github.com/pingcap/parser/mysql"
 	"github.com/pingcap/tidb/expression"
 	"github.com/pingcap/tidb/sessionctx"
@@ -134,3 +136,74 @@ func (rd *RowDecoder) DecodeAndEvalRowWithMap(ctx sessionctx.Context, handle int
 	}
 	return row, nil
 }
+
+// SubstituteGenColsInDecodeColMap substitutes generated columns in every expression in decodeColMap
+// with non-generated one by looking up decodeColMap.
+func SubstituteGenColsInDecodeColMap(decodeColMap map[int64]Column) {
+	// Sort columns by ID in ascending order.
+	var orderedCols []int64
+	for colID := range decodeColMap {
+		orderedCols = append(orderedCols, colID)
+	}
+	sort.Slice(orderedCols, func(i, j int) bool { return orderedCols[i] < orderedCols[j] })
+
+	// Iterate over decodeColMap, substitution only happens once because
+	// the column with smaller columnID can not refer to those with larger.
+	for _, colID := range orderedCols {
+		decCol := decodeColMap[colID]
+		if decCol.GenExpr != nil {
+			decodeColMap[colID] = Column{
+				Col:     decCol.Col,
+				GenExpr: substituteGeneratedColumn(decCol.GenExpr, decodeColMap),
+			}
+		} else {
+			decodeColMap[colID] = Column{
+				Col: decCol.Col,
+			}
+		}
+	}
+}
+
+// substituteGeneratedColumn substitutes generated columns in an expression with non-generated one by looking up decodeColMap.
+func substituteGeneratedColumn(expr expression.Expression, decodeColMap map[int64]Column) expression.Expression {
+	switch v := expr.(type) {
+	case *expression.Column:
+		if c, ok := decodeColMap[v.ID]; ok && c.GenExpr != nil {
+			return c.GenExpr
+		}
+		return v
+	case *expression.ScalarFunction:
+		if v.FuncName.L == ast.Cast {
+			newFunc := v.Clone().(*expression.ScalarFunction)
+			newFunc.GetArgs()[0] = substituteGeneratedColumn(newFunc.GetArgs()[0], decodeColMap)
+			return newFunc
+		}
+		newArgs := make([]expression.Expression, 0, len(v.GetArgs()))
+		for _, arg := range v.GetArgs() {
+			newArgs = append(newArgs, substituteGeneratedColumn(arg, decodeColMap))
+		}
+		return expression.NewFunctionInternal(v.GetCtx(), v.FuncName.L, v.RetType, newArgs...)
+	}
+	return expr
+}
+
+// RemoveUnusedVirtualCols removes all virtual columns in decodeColMap that cannot found in indexedCols.
+func RemoveUnusedVirtualCols(decodeColMap map[int64]Column, indexedCols []*table.Column) {
+	for colID, decCol := range decodeColMap {
+		col := decCol.Col
+		if !col.IsGenerated() || col.GeneratedStored {
+			continue
+		}
+
+		found := false
+		for _, v := range indexedCols {
+			if v.Offset == col.Offset {
+				found = true
+			}
+		}
+
+		if !found {
+			delete(decodeColMap, colID)
+		}
+	}
+}