Merge pull request pingcap#3 from guo-shaoge/tidb_plan_to_substrait

Author: XuHuaiyu

executor/adapter.go

@@ -19,6 +19,7 @@ import (
 	"context"
 	"fmt"
 	substraitgo "github.com/AilinKid/substrait-go/proto"
+	"github.com/substrait-io/substrait-go/proto/extensions"
 	"runtime/trace"
 	"strconv"
 	"strings"
@@ -483,9 +484,16 @@ func (a *ExecStmt) Exec(ctx context.Context) (_ sqlexec.RecordSet, err error) {
 		sctx.GetSessionVars().StmtCtx.MemTracker.SetBytesLimit(sctx.GetSessionVars().StmtCtx.MemQuotaQuery)
 	}
 	if pp, ok := a.Plan.(plannercore.PhysicalPlan); ok {
-		rel, err := pp.ToSubstraitPB(sctx)
-		if err == nil && sctx.GetSessionVars().ConnectionID != 0 {
+		ssHandler := plannercore.NewSubstraitHandler()
+		rel, err := pp.ToSubstraitPB(sctx, ssHandler)
+		if rel != nil && err == nil && sctx.GetSessionVars().ConnectionID != 0 {
 			plan := &substraitgo.Plan{
+				ExtensionUris: []*extensions.SimpleExtensionURI{
+					{
+						ExtensionUriAnchor: uint32(1),
+						Uri:                "whatever",
+					},
+				},
 				Relations: []*substraitgo.PlanRel{
 					{
 						RelType: &substraitgo.PlanRel_Root{
@@ -494,7 +502,19 @@ func (a *ExecStmt) Exec(ctx context.Context) (_ sqlexec.RecordSet, err error) {
 					},
 				},
 			}
-			// 把这个 rel 传给 velox
+			for funcName, funcAnchor := range ssHandler.SigMap {
+				plan.Extensions = append(plan.Extensions, &extensions.SimpleExtensionDeclaration{
+					MappingType: &extensions.SimpleExtensionDeclaration_ExtensionFunction_{
+						ExtensionFunction: &extensions.SimpleExtensionDeclaration_ExtensionFunction{
+							ExtensionUriReference: 1,
+							FunctionAnchor:        funcAnchor,
+							Name:                  funcName,
+						},
+					},
+				})
+			}
+
+			// 把这个 plan 传给 velox
 			logutil.BgLogger().Error(plan.String())
 		}
 	}

planner/core/plan.go

@@ -341,7 +341,7 @@ type PhysicalPlan interface {
 	// current task. If the child's task is cop task, some operator may close this task and return a new rootTask.
 	attach2Task(...task) task
 
-	ToSubstraitPB(ctx sessionctx.Context) (*substraitgo.Rel, error)
+	ToSubstraitPB(ctx sessionctx.Context, ssHandler *SubstraitHandler) (*substraitgo.Rel, error)
 
 	// ToPB converts physical plan to tipb executor.
 	ToPB(ctx sessionctx.Context, storeType kv.StoreType) (*tipb.Executor, error)

planner/core/plan_to_pb.go

@@ -20,7 +20,9 @@ import (
 	"github.com/pingcap/tidb/expression"
 	"github.com/pingcap/tidb/expression/aggregation"
 	"github.com/pingcap/tidb/kv"
+	"github.com/pingcap/tidb/parser/ast"
 	"github.com/pingcap/tidb/parser/model"
+	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/table"
 	"github.com/pingcap/tidb/table/tables"
@@ -29,12 +31,14 @@ import (
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util"
 	"github.com/pingcap/tidb/util/codec"
+	"github.com/pingcap/tidb/util/logutil"
 	"github.com/pingcap/tidb/util/ranger"
 	"github.com/pingcap/tipb/go-tipb"
+	"go.uber.org/zap"
 )
 
 // ToSubstraitPB implements PhysicalPlan ToSubstraitPB interface.
-func (p *basePhysicalPlan) ToSubstraitPB(ctx sessionctx.Context) (_ *substraitgo.Rel, err error) {
+func (p *basePhysicalPlan) ToSubstraitPB(ctx sessionctx.Context, ssHandler *SubstraitHandler) (_ *substraitgo.Rel, err error) {
 	return nil, errors.Errorf("plan %s fails converts to PB", p.basePlan.ExplainID())
 }
 
@@ -125,11 +129,239 @@ func (p *PhysicalSelection) ToPB(ctx sessionctx.Context, storeType kv.StoreType)
 	}
 	return &tipb.Executor{Tp: tipb.ExecType_TypeSelection, Selection: selExec, ExecutorId: &executorID}, nil
 }
+func getSubStraitType(tp byte) string {
+	if tp == mysql.TypeLong {
+		return "i32"
+	}
+	if tp == mysql.TypeLonglong {
+		return "i64"
+	}
+	if tp == mysql.TypeDouble {
+		return "fp64"
+	}
+	panic("not supproted type")
+}
+
+func getSubstraitPBFunctionArguments(offsets []int32) (funcArgs []*substraitgo.FunctionArgument) {
+	getFunctionArgument := func(offset int32) *substraitgo.FunctionArgument {
+		return &substraitgo.FunctionArgument{
+			ArgType: &substraitgo.FunctionArgument_Value{
+				Value: &substraitgo.Expression{
+					RexType: &substraitgo.Expression_Selection{
+						Selection: &substraitgo.Expression_FieldReference{
+							ReferenceType: &substraitgo.Expression_FieldReference_DirectReference{
+								DirectReference: &substraitgo.Expression_ReferenceSegment{
+									ReferenceType: &substraitgo.Expression_ReferenceSegment_StructField_{
+										StructField: &substraitgo.Expression_ReferenceSegment_StructField{
+											// col offset from child chunk
+											Field: offset,
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		}
+	}
+	for _, off := range offsets {
+		funcArgs = append(funcArgs, getFunctionArgument(off))
+	}
+	return
+}
+
+type SubstraitHandler struct {
+	autoIncId uint32
+	SigMap    map[string]uint32
+}
+
+func NewSubstraitHandler() *SubstraitHandler {
+	return &SubstraitHandler{
+		autoIncId: uint32(1),
+		SigMap:    map[string]uint32{},
+	}
+}
+
+func (h *SubstraitHandler) insertSig(sigName string) uint32 {
+	if _, ok := h.SigMap[sigName]; !ok {
+		h.SigMap[sigName] = h.autoIncId
+		h.autoIncId++
+	}
+	return h.getSigID(sigName)
+}
+
+func (h *SubstraitHandler) getSigID(sigName string) uint32 {
+	return h.SigMap[sigName]
+}
 
-//func (p *PhysicalProjection) ToSubstraitPB() (*substraitgo.Rel, error){
-//	proj := &substraitgo.ProjectRel{}
-//	proj.Input
-//}
+func scalarFuncFieldTypeToSubstraitOutputType(sf *expression.ScalarFunction) (outputType *substraitgo.Type) {
+	switch sf.GetType().EvalType() {
+	case types.ETInt:
+		outputType = &substraitgo.Type{
+			Kind: &substraitgo.Type_I64_{
+				I64: &substraitgo.Type_I64{},
+			},
+		}
+	case types.ETReal:
+		outputType = &substraitgo.Type{
+			Kind: &substraitgo.Type_Fp64{
+				Fp64: &substraitgo.Type_FP64{},
+			},
+		}
+	}
+	switch sf.FuncName.L {
+	case "lt", "lte", "gt", "gte", "and", "or", "eq", "not":
+		outputType = &substraitgo.Type{
+			Kind: &substraitgo.Type_Bool{
+				Bool: &substraitgo.Type_Boolean{
+					Nullability: substraitgo.Type_NULLABILITY_NULLABLE,
+				},
+			},
+		}
+	}
+	return
+}
+
+func (h *SubstraitHandler) scalarFuncToSubstraitgoExpr(sf *expression.ScalarFunction) (*substraitgo.Expression, error) {
+	funcSig := tiDBFuncNameToVeloxFuncName[sf.FuncName.L]
+	var offsets []int32
+	for _, arg := range sf.GetArgs() {
+		col, ok := arg.(*expression.Column)
+		if !ok {
+			return nil, errors.Errorf("fail to convert proj to substrait, only support arg of column type")
+		}
+		offsets = append(offsets, int32(col.Index))
+	}
+	funcSig = funcSig + ":" + getSubStraitType(sf.GetArgs()[0].GetType().GetType()) + "_" + getSubStraitType(sf.GetArgs()[1].GetType().GetType())
+
+	sExpr := &substraitgo.Expression{
+		RexType: &substraitgo.Expression_ScalarFunction_{
+			ScalarFunction: &substraitgo.Expression_ScalarFunction{
+				// function anchor
+				FunctionReference: h.insertSig(funcSig),
+				// function args
+				Arguments:  getSubstraitPBFunctionArguments(offsets),
+				OutputType: scalarFuncFieldTypeToSubstraitOutputType(sf),
+			},
+		},
+	}
+	return sExpr, nil
+}
+
+func (h *SubstraitHandler) columnToSubstraitgoExpr(col *expression.Column) *substraitgo.Expression {
+	sExpr := &substraitgo.Expression{
+		RexType: &substraitgo.Expression_Selection{
+			Selection: &substraitgo.Expression_FieldReference{
+				ReferenceType: &substraitgo.Expression_FieldReference_DirectReference{
+					DirectReference: &substraitgo.Expression_ReferenceSegment{
+						ReferenceType: &substraitgo.Expression_ReferenceSegment_StructField_{
+							StructField: &substraitgo.Expression_ReferenceSegment_StructField{
+								Field: int32(col.Index),
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+	return sExpr
+}
+
+func tidbConstToVeloxExpressionLiteral(cont *expression.Constant) (el *substraitgo.Expression_Literal) {
+	switch cont.GetType().GetType() {
+	case mysql.TypeLonglong:
+		el = &substraitgo.Expression_Literal{
+			LiteralType: &substraitgo.Expression_Literal_I64{
+				I64: cont.Value.GetInt64(),
+			},
+		}
+	case mysql.TypeFloat:
+		el = &substraitgo.Expression_Literal{
+			LiteralType: &substraitgo.Expression_Literal_Fp32{
+				Fp32: cont.Value.GetFloat32(),
+			},
+		}
+	case mysql.TypeDouble:
+		el = &substraitgo.Expression_Literal{
+			LiteralType: &substraitgo.Expression_Literal_Fp64{
+				Fp64: cont.Value.GetFloat64(),
+			},
+		}
+	}
+	return
+}
+
+func (h *SubstraitHandler) constToSubstraitgoExpr(cont *expression.Constant) *substraitgo.Expression {
+	sExpr := &substraitgo.Expression{
+		RexType: &substraitgo.Expression_Literal_{
+			Literal: tidbConstToVeloxExpressionLiteral(cont),
+		},
+	}
+	return sExpr
+}
+
+func (h *SubstraitHandler) buildSubstraitProjExpression(tidbExpr []expression.Expression) (substraitgoExprs []*substraitgo.Expression, err error) {
+	for _, expr := range tidbExpr {
+		var sExpr *substraitgo.Expression
+		var err error
+		switch x := expr.(type) {
+		case *expression.Column:
+			sExpr = h.columnToSubstraitgoExpr(x)
+		case *expression.Constant:
+			sExpr = h.constToSubstraitgoExpr(x)
+		case *expression.ScalarFunction:
+			sExpr, err = h.scalarFuncToSubstraitgoExpr(x)
+			if err != nil {
+				return nil, err
+			}
+		}
+		substraitgoExprs = append(substraitgoExprs, sExpr)
+	}
+	return substraitgoExprs, nil
+}
+
+var tiDBFuncNameToVeloxFuncName = map[string]string{
+	ast.LE:       "lte",
+	ast.LT:       "lt",
+	ast.GE:       "gte",
+	ast.GT:       "gt",
+	ast.Plus:     "plus",
+	ast.Minus:    "minus",
+	ast.Mul:      "multiply",
+	ast.Mod:      "mod",
+	ast.EQ:       "eq",
+	ast.UnaryNot: "not",
+	ast.And:      "and",
+	ast.LogicOr:  "or",
+}
+
+// ToSubstraitPB implements the substraitgo interface.
+func (p *PhysicalProjection) ToSubstraitPB(ctx sessionctx.Context, ssHandler *SubstraitHandler) (*substraitgo.Rel, error) {
+	childRel, err := p.children[0].ToSubstraitPB(ctx, ssHandler)
+	if err != nil {
+		return nil, err
+	}
+	sspb := &substraitgo.ProjectRel{
+		Common: &substraitgo.RelCommon{
+			EmitKind: &substraitgo.RelCommon_Direct_{
+				Direct: &substraitgo.RelCommon_Direct{},
+			},
+		},
+		Input: childRel,
+	}
+	sspbExpression, err := ssHandler.buildSubstraitProjExpression(p.Exprs)
+	if err != nil {
+		logutil.BgLogger().Error("error", zap.Error(err))
+		return nil, err
+	}
+	sspb.Expressions = sspbExpression
+	logutil.BgLogger().Warn("expression", zap.Int("len", len(sspb.Expressions)))
+	// extract function map (assume there are all simple functions)
+	return &substraitgo.Rel{
+		RelType: &substraitgo.Rel_Project{Project: sspb},
+	}, nil
+}
 
 // ToPB implements PhysicalPlan ToPB interface.
 func (p *PhysicalProjection) ToPB(ctx sessionctx.Context, storeType kv.StoreType) (*tipb.Executor, error) {
@@ -195,7 +427,7 @@ func (p *PhysicalLimit) ToPB(ctx sessionctx.Context, storeType kv.StoreType) (*t
 	return &tipb.Executor{Tp: tipb.ExecType_TypeLimit, Limit: limitExec, ExecutorId: &executorID}, nil
 }
 
-func (p *PhysicalTableReader) ToSubstraitPB(ctx sessionctx.Context) (rel *substraitgo.Rel, err error) {
+func (p *PhysicalTableReader) ToSubstraitPB(ctx sessionctx.Context, ssHandler *SubstraitHandler) (rel *substraitgo.Rel, err error) {
 	tableScan, ok := p.TablePlans[0].(*PhysicalTableScan)
 	if !ok {
 		return nil, nil

planner/core/point_get_plan.go

@@ -93,7 +93,7 @@ type PointGetPlan struct {
 	accessCols []*expression.Column
 }
 
-func (p *PointGetPlan) ToSubstraitPB(ctx sessionctx.Context) (*substraitgo.Rel, error) {
+func (p *PointGetPlan) ToSubstraitPB(ctx sessionctx.Context, ssHandler *SubstraitHandler) (*substraitgo.Rel, error) {
 	return nil, nil
 }
 
@@ -324,7 +324,7 @@ type BatchPointGetPlan struct {
 	accessCols []*expression.Column
 }
 
-func (p *BatchPointGetPlan) ToSubstraitPB(ctx sessionctx.Context) (*substraitgo.Rel, error) {
+func (p *BatchPointGetPlan) ToSubstraitPB(ctx sessionctx.Context, ssHandler *SubstraitHandler) (*substraitgo.Rel, error) {
 	return nil, nil
 }