Golang expression.Schema类代码示例

func (a *havingAndOrderbyExprResolver) resolveFromSchema(v *ast.ColumnNameExpr, schema expression.Schema) (int, error) {
	col, err := schema.FindColumn(v.Name)
	if err != nil {
		return -1, errors.Trace(err)
	}
	if col == nil {
		return -1, nil
	}
	newColName := &ast.ColumnName{
		Schema: col.DBName,
		Table:  col.TblName,
		Name:   col.ColName,
	}
	for i, field := range a.selectFields {
		if c, ok := field.Expr.(*ast.ColumnNameExpr); ok && colMatch(newColName, c.Name) {
			return i, nil
		}
	}
	sf := &ast.SelectField{
		Expr:      &ast.ColumnNameExpr{Name: newColName},
		Auxiliary: true,
	}
	sf.Expr.SetType(col.GetType())
	a.selectFields = append(a.selectFields, sf)
	return len(a.selectFields) - 1, nil
}

func makeUsedList(usedCols []*expression.Column, schema expression.Schema) []bool {
	used := make([]bool, len(schema))
	for _, col := range usedCols {
		idx := schema.GetIndex(col)
		used[idx] = true
	}
	return used
}

func initColumnIndexInExpr(expr expression.Expression, schema expression.Schema) {
	switch assign := expr; assign.(type) {
	case (*expression.Column):
		assign.(*expression.Column).Index = schema.GetIndex(assign.(*expression.Column))
	case (*expression.ScalarFunction):
		for i, args := 0, assign.(*expression.ScalarFunction).Args; i < len(args); i++ {
			initColumnIndexInExpr(args[i], schema)
		}
	}
}

func getUsedList(usedCols []*expression.Column, schema expression.Schema) []bool {
	used := make([]bool, len(schema))
	for _, col := range usedCols {
		idx := schema.GetIndex(col)
		if idx == -1 {
			log.Errorf("Can't find column %s from schema %s.", col, schema)
		}
		used[idx] = true
	}
	return used
}

// columnSubstitute substitutes the columns in filter to expressions in select fields.
// e.g. select * from (select b as a from t) k where a < 10 => select * from (select b as a from t where b < 10) k.
func columnSubstitute(expr expression.Expression, schema expression.Schema, newExprs []expression.Expression) expression.Expression {
	switch v := expr.(type) {
	case *expression.Column:
		id := schema.GetIndex(v)
		return newExprs[id]
	case *expression.ScalarFunction:
		for i, arg := range v.Args {
			v.Args[i] = columnSubstitute(arg, schema, newExprs)
		}
	}
	return expr
}

func (e *havingAndOrderbyResolver) Enter(inNode ast.Node) (ast.Node, bool) {
	switch v := inNode.(type) {
	case *ast.ValueExpr, *ast.ColumnName, *ast.ParenthesesExpr:
	case *ast.ColumnNameExpr:
		var first, second expression.Schema
		var col *expression.Column
		fromSchema := e.proj.GetChildByIndex(0).GetSchema()
		fromSchemaFirst := e.orderBy && e.inExpr
		if fromSchemaFirst {
			first, second = fromSchema, e.proj.GetSchema()
			col, e.err = first.FindColumn(v.Name)
		} else {
			first, second = e.proj.GetSchema(), fromSchema
			col, e.err = first.FindSelectFieldColumn(v.Name, e.proj.Exprs)
		}
		if e.err != nil {
			return inNode, true
		}
		if col == nil {
			if fromSchemaFirst {
				col, e.err = second.FindSelectFieldColumn(v.Name, e.proj.Exprs)
			} else {
				col, e.err = second.FindColumn(v.Name)
			}

			if e.err != nil {
				return inNode, true
			}
			if col == nil {
				e.err = errors.Errorf("Can't find Column %s", v.Name.Name)
				return inNode, true
			}
			if !fromSchemaFirst {
				e.addProjectionExpr(v, col)
			} else {
				e.mapper[v] = col
			}
		} else if fromSchemaFirst {
			e.addProjectionExpr(v, col)
		} else {
			e.mapper[v] = col
		}
	case *ast.SubqueryExpr, *ast.CompareSubqueryExpr, *ast.ExistsSubqueryExpr:
		return inNode, true
	default:
		e.inExpr = true
	}
	return inNode, false
}

// columnSubstitute substitutes the columns in filter to expressions in select fields.
// e.g. select * from (select b as a from t) k where a < 10 => select * from (select b as a from t where b < 10) k.
func columnSubstitute(expr expression.Expression, schema expression.Schema, newExprs []expression.Expression) expression.Expression {
	switch v := expr.(type) {
	case *expression.Column:
		id := schema.GetIndex(v)
		if id == -1 {
			log.Errorf("Can't find columns %s in schema %s", v.ToString(), schema.ToString())
		}
		return newExprs[id]
	case *expression.ScalarFunction:
		for i, arg := range v.Args {
			v.Args[i] = columnSubstitute(arg, schema, newExprs)
		}
	}
	return expr
}

// replaceColsInPropBySchema replaces the columns in original prop with the columns in schema.
func replaceColsInPropBySchema(prop *requiredProperty, schema expression.Schema) *requiredProperty {
	newProps := make([]*columnProp, 0, len(prop.props))
	for _, p := range prop.props {
		idx := schema.GetIndex(p.col)
		if idx == -1 {
			log.Errorf("Can't find column %s in schema", p.col)
		}
		newProps = append(newProps, &columnProp{col: schema[idx], desc: p.desc})
	}
	return &requiredProperty{
		props:      newProps,
		sortKeyLen: prop.sortKeyLen,
		limit:      prop.limit,
	}
}

// getAggFuncChildIdx gets which children it belongs to, 0 stands for left, 1 stands for right, -1 stands for both.
func (a *aggPushDownSolver) getAggFuncChildIdx(aggFunc expression.AggregationFunction, schema expression.Schema) int {
	fromLeft, fromRight := false, false
	var cols []*expression.Column
	for _, arg := range aggFunc.GetArgs() {
		cols = append(cols, expression.ExtractColumns(arg)...)
	}
	for _, col := range cols {
		if schema.GetIndex(col) != -1 {
			fromLeft = true
		} else {
			fromRight = true
		}
	}
	if fromLeft && fromRight {
		return -1
	} else if fromLeft {
		return 0
	}
	return 1
}

func retrieveColumnsInExpression(expr expression.Expression, schema expression.Schema) (
	expression.Expression, error) {
	switch v := expr.(type) {
	case *expression.ScalarFunction:
		for i, arg := range v.Args {
			newExpr, err := retrieveColumnsInExpression(arg, schema)
			if err != nil {
				return nil, errors.Trace(err)
			}
			v.Args[i] = newExpr
		}
	case *expression.Column:
		if !v.Correlated {
			newColumn := schema.RetrieveColumn(v)
			if newColumn == nil {
				return nil, errors.Errorf("Can't Find column %s.", expr.ToString())
			}
			return newColumn, nil
		}
	}
	return expr, nil
}

// calculateResultOfExpression set inner table columns in a expression as null and calculate the finally result of the scalar function.
func calculateResultOfExpression(schema expression.Schema, expr expression.Expression) (expression.Expression, error) {
	switch x := expr.(type) {
	case *expression.ScalarFunction:
		var err error
		args := make([]expression.Expression, len(x.Args))
		for i, arg := range x.Args {
			args[i], err = calculateResultOfExpression(schema, arg)
		}
		if err != nil {
			return nil, errors.Trace(err)
		}
		return expression.NewFunction(x.FuncName.L, types.NewFieldType(mysql.TypeTiny), args...)
	case *expression.Column:
		if schema.GetIndex(x) == -1 {
			return x, nil
		}
		constant := &expression.Constant{Value: types.Datum{}}
		constant.Value.SetNull()
		return constant, nil
	default:
		return x.DeepCopy(), nil
	}
}

func (a *AggregateFuncExtractor) resolveFromSchema(v *ast.ColumnNameExpr, schema expression.Schema) (int, error) {
	col, err := schema.FindColumn(v.Name)
	if err != nil {
		return -1, errors.Trace(err)
	}
	if col == nil {
		return -1, nil
	}
	newColName := &ast.ColumnName{
		Schema: col.DBName,
		Table:  col.TblName,
		Name:   col.ColName,
	}
	for i, field := range a.selectFields {
		if c, ok := field.Expr.(*ast.ColumnNameExpr); ok && colMatch(newColName, c.Name) {
			return i, nil
		}
	}
	a.selectFields = append(a.selectFields, &ast.SelectField{
		Expr:      &ast.ColumnNameExpr{Name: newColName},
		Auxiliary: true,
	})
	return len(a.selectFields) - 1, nil
}

func (b *planBuilder) buildNewSelect(sel *ast.SelectStmt) Plan {
	oldLen := len(sel.Fields.Fields)
	hasAgg := b.detectSelectAgg(sel)
	var aggFuncs []*ast.AggregateFuncExpr
	var havingMap, orderMap, totalMap map[*ast.AggregateFuncExpr]int
	if hasAgg {
		aggFuncs, havingMap, orderMap, totalMap = b.extractAggFunc(sel)
	}
	// Build subquery
	// Convert subquery to expr with plan
	// TODO: add subquery support.
	//b.buildSubquery(sel)
	var p Plan
	if sel.From != nil {
		p = b.buildResultSetNode(sel.From.TableRefs)
		if b.err != nil {
			return nil
		}
		if sel.Where != nil {
			p = b.buildSelection(p, sel.Where, nil)
		}
		if b.err != nil {
			return nil
		}
		if sel.LockTp != ast.SelectLockNone {
			p = b.buildSelectLock(p, sel.LockTp)
			if b.err != nil {
				return nil
			}
		}
		if hasAgg {
			p = b.buildAggregation(p, aggFuncs, sel.GroupBy)
			if b.err != nil {
				return nil
			}
		}
	} else {
		p = b.buildNewTableDual()
		if b.err != nil {
			return nil
		}
		if sel.Where != nil {
			b.buildSelection(p, sel.Where, nil)
		}
		if hasAgg {
			p = b.buildAggregation(p, aggFuncs, nil)
			if b.err != nil {
				return nil
			}
		}
	}
	p = b.buildProjection(p, sel.Fields.Fields, totalMap)
	if b.err != nil {
		return nil
	}
	if sel.Having != nil {
		p = b.buildSelection(p, sel.Having.Expr, havingMap)
		if b.err != nil {
			return nil
		}
	}
	if sel.Distinct {
		p = b.buildDistinct(p)
		if b.err != nil {
			return nil
		}
	}
	// TODO: implement push order during cbo
	if sel.OrderBy != nil {
		p = b.buildNewSort(p, sel.OrderBy.Items, orderMap)
		if b.err != nil {
			return nil
		}
	}
	if sel.Limit != nil {
		p = b.buildLimit(p, sel.Limit)
		if b.err != nil {
			return nil
		}
	}
	if oldLen != len(sel.Fields.Fields) {
		proj := &Projection{}
		proj.id = b.allocID(proj)
		var newSchema expression.Schema
		oldSchema := p.GetSchema()
		proj.Exprs = make([]expression.Expression, 0, oldLen)
		for _, col := range oldSchema[:oldLen] {
			proj.Exprs = append(proj.Exprs, col)
		}
		newSchema = oldSchema[:oldLen]
		newSchema = newSchema.DeepCopy()
		for _, s := range newSchema {
			s.FromID = proj.id
		}
		proj.SetSchema(newSchema)
		addChild(proj, p)
		proj.correlated = p.IsCorrelated()
		return proj
	}
	return p
//.........这里部分代码省略.........