*: a little fix and add 3 functions. (pingcap#3486)

expression/schema.go

@@ -74,6 +74,13 @@ func (s *Schema) Clone() *Schema {
 // FindColumn finds an Column from schema for a ast.ColumnName. It compares the db/table/column names.
 // If there are more than one result, it will raise ambiguous error.
 func (s *Schema) FindColumn(astCol *ast.ColumnName) (*Column, error) {
+	col, _, err := s.FindColumnAndIndex(astCol)
+	return col, errors.Trace(err)
+}
+
+// FindColumnAndIndex finds an Column and its index from schema for a ast.ColumnName.
+// It compares the db/table/column names. If there are more than one result, raise ambiguous error.
+func (s *Schema) FindColumnAndIndex(astCol *ast.ColumnName) (*Column, int, error) {
 	dbName, tblName, colName := astCol.Schema, astCol.Table, astCol.Name
 	idx := -1
 	for i, col := range s.Columns {
@@ -83,14 +90,14 @@ func (s *Schema) FindColumn(astCol *ast.ColumnName) (*Column, error) {
 			if idx == -1 {
 				idx = i
 			} else {
-				return nil, errors.Errorf("Column %s is ambiguous", col.String())
+				return nil, -1, errors.Errorf("Column %s is ambiguous", col.String())
 			}
 		}
 	}
 	if idx == -1 {
-		return nil, nil
+		return nil, idx, nil
 	}
-	return s.Columns[idx], nil
+	return s.Columns[idx], idx, nil
 }
 
 // RetrieveColumn retrieves column in expression from the columns in schema.

model/model.go

@@ -135,6 +135,18 @@ func (t *TableInfo) GetPkName() CIStr {
 	return CIStr{}
 }
 
+// ColumnIsInIndex checks whether c is included in any indices of t.
+func (t *TableInfo) ColumnIsInIndex(c *ColumnInfo) bool {
+	for _, index := range t.Indices {
+		for _, column := range index.Columns {
+			if column.Name.L == c.Name.L {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 // IndexColumn provides index column info.
 type IndexColumn struct {
 	Name   CIStr `json:"name"`   // Index name

plan/expression_rewriter.go

@@ -60,13 +60,22 @@ func evalAstExpr(expr ast.ExprNode, ctx context.Context) (types.Datum, error) {
 // asScalar means whether this expression must be treated as a scalar expression.
 // And this function returns a result expression, a new plan that may have apply or semi-join.
 func (b *planBuilder) rewrite(expr ast.ExprNode, p LogicalPlan, aggMapper map[*ast.AggregateFuncExpr]int, asScalar bool) (
+	expression.Expression, LogicalPlan, error) {
+	return b.rewriteWithPreprocess(expr, p, aggMapper, asScalar, nil)
+}
+
+// rewriteWithPreprocess is for handling the situation that we need to adjust the input ast tree
+// before really using its node in `expressionRewriter.Leave`. In that case, we first call
+// er.preprocess(expr), which returns a new expr. Then we use the new expr in `Leave`.
+func (b *planBuilder) rewriteWithPreprocess(expr ast.ExprNode, p LogicalPlan, aggMapper map[*ast.AggregateFuncExpr]int, asScalar bool, preprocess func(ast.Node) ast.Node) (
 	expression.Expression, LogicalPlan, error) {
 	er := &expressionRewriter{
-		p:        p,
-		aggrMap:  aggMapper,
-		b:        b,
-		asScalar: asScalar,
-		ctx:      b.ctx,
+		p:          p,
+		aggrMap:    aggMapper,
+		b:          b,
+		asScalar:   asScalar,
+		ctx:        b.ctx,
+		preprocess: preprocess,
 	}
 	if p != nil {
 		er.schema = p.Schema()
@@ -98,6 +107,9 @@ type expressionRewriter struct {
 	ctx      context.Context
 	// asScalar means the return value must be a scalar value.
 	asScalar bool
+
+	// preprocess is called for every ast.Node in Leave.
+	preprocess func(ast.Node) ast.Node
 }
 
 func getRowLen(e expression.Expression) int {
@@ -599,11 +611,14 @@ func (er *expressionRewriter) handleScalarSubquery(v *ast.SubqueryExpr) (ast.Nod
 }
 
 // Leave implements Visitor interface.
-func (er *expressionRewriter) Leave(inNode ast.Node) (retNode ast.Node, ok bool) {
+func (er *expressionRewriter) Leave(originInNode ast.Node) (retNode ast.Node, ok bool) {
 	if er.err != nil {
 		return retNode, false
 	}
-
+	var inNode = originInNode
+	if er.preprocess != nil {
+		inNode = er.preprocess(inNode)
+	}
 	switch v := inNode.(type) {
 	case *ast.AggregateFuncExpr, *ast.ColumnNameExpr, *ast.ParenthesesExpr, *ast.WhenClause,
 		*ast.SubqueryExpr, *ast.ExistsSubqueryExpr, *ast.CompareSubqueryExpr, *ast.ValuesExpr:
@@ -658,7 +673,7 @@ func (er *expressionRewriter) Leave(inNode ast.Node) (retNode ast.Node, ok bool)
 	if er.err != nil {
 		return retNode, false
 	}
-	return inNode, true
+	return originInNode, true
 }
 
 func datumToConstant(d types.Datum, tp byte) *expression.Constant {