inputs method to return additional information about the input primit…

…ive (#13883) Signed-off-by: Harshit Gangal <[email protected]>
vitessio · Aug 31, 2023 · 8452262 · 8452262
1 parent 2442a0f
commit 8452262
Show file tree

Hide file tree

Showing 37 changed files with 616 additions and 482 deletions.
diff --git a/go/vt/vtgate/engine/concatenate.go b/go/vt/vtgate/engine/concatenate.go
@@ -340,8 +340,8 @@ func (c *Concatenate) NeedsTransaction() bool {
 }
 
 // Inputs returns the input primitives for this
-func (c *Concatenate) Inputs() []Primitive {
-	return c.Sources
+func (c *Concatenate) Inputs() ([]Primitive, []map[string]any) {
+	return c.Sources, nil
 }
 
 func (c *Concatenate) description() PrimitiveDescription {

diff --git a/go/vt/vtgate/engine/distinct.go b/go/vt/vtgate/engine/distinct.go
@@ -246,8 +246,8 @@ func (d *Distinct) NeedsTransaction() bool {
 }
 
 // Inputs implements the Primitive interface
-func (d *Distinct) Inputs() []Primitive {
-	return []Primitive{d.Source}
+func (d *Distinct) Inputs() ([]Primitive, []map[string]any) {
+	return []Primitive{d.Source}, nil
 }
 
 func (d *Distinct) description() PrimitiveDescription {

diff --git a/go/vt/vtgate/engine/exec_prepared_statement.go b/go/vt/vtgate/engine/exec_prepared_statement.go
@@ -61,8 +61,8 @@ func (e *ExecStmt) TryStreamExecute(ctx context.Context, vcursor VCursor, bindVa
 	return vcursor.StreamExecutePrimitive(ctx, e.Input, bindVars, wantfields, callback)
 }
 
-func (e *ExecStmt) Inputs() []Primitive {
-	return []Primitive{e.Input}
+func (e *ExecStmt) Inputs() ([]Primitive, []map[string]any) {
+	return []Primitive{e.Input}, nil
 }
 
 func (e *ExecStmt) description() PrimitiveDescription {

diff --git a/go/vt/vtgate/engine/fake_primitive_test.go b/go/vt/vtgate/engine/fake_primitive_test.go
@@ -43,8 +43,8 @@ type fakePrimitive struct {
 	allResultsInOneCall bool
 }
 
-func (f *fakePrimitive) Inputs() []Primitive {
-	return []Primitive{}
+func (f *fakePrimitive) Inputs() ([]Primitive, []map[string]any) {
+	return []Primitive{}, nil
 }
 
 var _ Primitive = (*fakePrimitive)(nil)

diff --git a/go/vt/vtgate/engine/filter.go b/go/vt/vtgate/engine/filter.go
@@ -108,8 +108,8 @@ func (f *Filter) GetFields(ctx context.Context, vcursor VCursor, bindVars map[st
 }
 
 // Inputs returns the input to limit
-func (f *Filter) Inputs() []Primitive {
-	return []Primitive{f.Input}
+func (f *Filter) Inputs() ([]Primitive, []map[string]any) {
+	return []Primitive{f.Input}, nil
 }
 
 func (f *Filter) description() PrimitiveDescription {

diff --git a/go/vt/vtgate/engine/fk_cascade.go b/go/vt/vtgate/engine/fk_cascade.go
@@ -18,6 +18,7 @@ package engine
 
 import (
 	"context"
+	"fmt"
 
 	"vitess.io/vitess/go/sqltypes"
 	querypb "vitess.io/vitess/go/vt/proto/query"
@@ -167,32 +168,30 @@ func (fkc *FkCascade) TryStreamExecute(ctx context.Context, vcursor VCursor, bin
 }
 
 // Inputs implements the Primitive interface.
-func (fkc *FkCascade) Inputs() []Primitive {
-	return nil
+func (fkc *FkCascade) Inputs() ([]Primitive, []map[string]any) {
+	var inputs []Primitive
+	var inputsMap []map[string]any
+	inputs = append(inputs, fkc.Selection)
+	inputsMap = append(inputsMap, map[string]any{
+		inputName: "Selection",
+	})
+	for idx, child := range fkc.Children {
+		inputsMap = append(inputsMap, map[string]any{
+			inputName: fmt.Sprintf("CascadeChild-%d", idx+1),
+			"BvName":  child.BVName,
+			"Cols":    child.Cols,
+		})
+		inputs = append(inputs, child.Exec)
+	}
+	inputs = append(inputs, fkc.Parent)
+	inputsMap = append(inputsMap, map[string]any{
+		inputName: "Parent",
+	})
+	return inputs, inputsMap
 }
 
 func (fkc *FkCascade) description() PrimitiveDescription {
-	var childrenDesc []PrimitiveDescription
-	for _, child := range fkc.Children {
-		childrenDesc = append(childrenDesc, PrimitiveDescription{
-			OperatorType: "FkCascadeChild",
-			Inputs: []PrimitiveDescription{
-				PrimitiveToPlanDescription(child.Exec),
-			},
-			Other: map[string]any{
-				"BvName": child.BVName,
-				"Cols":   child.Cols,
-			},
-		})
-	}
-	return PrimitiveDescription{
-		OperatorType: fkc.RouteType(),
-		Other: map[string]any{
-			"Selection": PrimitiveToPlanDescription(fkc.Selection),
-			"Parent":    PrimitiveToPlanDescription(fkc.Parent),
-			"Children":  childrenDesc,
-		},
-	}
+	return PrimitiveDescription{OperatorType: fkc.RouteType()}
 }
 
 var _ Primitive = (*FkCascade)(nil)
diff --git a/go/vt/vtgate/engine/fk_verify.go b/go/vt/vtgate/engine/fk_verify.go
@@ -18,6 +18,7 @@ package engine
 
 import (
 	"context"
+	"fmt"
 
 	"vitess.io/vitess/go/sqltypes"
 	querypb "vitess.io/vitess/go/vt/proto/query"
@@ -143,31 +144,27 @@ func (f *FkVerify) TryStreamExecute(ctx context.Context, vcursor VCursor, bindVa
 }
 
 // Inputs implements the Primitive interface
-func (f *FkVerify) Inputs() []Primitive {
-	return nil
+func (f *FkVerify) Inputs() ([]Primitive, []map[string]any) {
+	var inputs []Primitive
+	var inputsMap []map[string]any
+	for idx, parent := range f.Verify {
+		inputsMap = append(inputsMap, map[string]any{
+			inputName: fmt.Sprintf("VerifyParent-%d", idx+1),
+			"BvName":  parent.BvName,
+			"Cols":    parent.Cols,
+		})
+		inputs = append(inputs, parent.Exec)
+	}
+	inputs = append(inputs, f.Exec)
+	inputsMap = append(inputsMap, map[string]any{
+		inputName: "Child",
+	})
+	return inputs, inputsMap
+
 }
 
 func (f *FkVerify) description() PrimitiveDescription {
-	var parentDesc []PrimitiveDescription
-	for _, parent := range f.Verify {
-		parentDesc = append(parentDesc, PrimitiveDescription{
-			OperatorType: "FkVerifyParent",
-			Inputs: []PrimitiveDescription{
-				PrimitiveToPlanDescription(parent.Exec),
-			},
-			Other: map[string]any{
-				"BvName": parent.BvName,
-				"Cols":   parent.Cols,
-			},
-		})
-	}
-	return PrimitiveDescription{
-		OperatorType: f.RouteType(),
-		Other: map[string]any{
-			"Parent": parentDesc,
-			"Child":  PrimitiveToPlanDescription(f.Exec),
-		},
-	}
+	return PrimitiveDescription{OperatorType: f.RouteType()}
 }
 
 var _ Primitive = (*FkVerify)(nil)

diff --git a/go/vt/vtgate/engine/hash_join.go b/go/vt/vtgate/engine/hash_join.go
@@ -234,8 +234,8 @@ func (hj *HashJoin) NeedsTransaction() bool {
 }
 
 // Inputs implements the Primitive interface
-func (hj *HashJoin) Inputs() []Primitive {
-	return []Primitive{hj.Left, hj.Right}
+func (hj *HashJoin) Inputs() ([]Primitive, []map[string]any) {
+	return []Primitive{hj.Left, hj.Right}, nil
 }
 
 // description implements the Primitive interface

diff --git a/go/vt/vtgate/engine/insert.go b/go/vt/vtgate/engine/insert.go
@@ -108,11 +108,11 @@ type (
 	ksID = []byte
 )
 
-func (ins *Insert) Inputs() []Primitive {
+func (ins *Insert) Inputs() ([]Primitive, []map[string]any) {
 	if ins.Input == nil {
-		return nil
+		return nil, nil
 	}
-	return []Primitive{ins.Input}
+	return []Primitive{ins.Input}, nil
 }
 
 // NewQueryInsert creates an Insert with a query string.

diff --git a/go/vt/vtgate/engine/join.go b/go/vt/vtgate/engine/join.go
@@ -189,8 +189,8 @@ func (jn *Join) GetFields(ctx context.Context, vcursor VCursor, bindVars map[str
 }
 
 // Inputs returns the input primitives for this join
-func (jn *Join) Inputs() []Primitive {
-	return []Primitive{jn.Left, jn.Right}
+func (jn *Join) Inputs() ([]Primitive, []map[string]any) {
+	return []Primitive{jn.Left, jn.Right}, nil
 }
 
 func joinFields(lfields, rfields []*querypb.Field, cols []int) []*querypb.Field {

diff --git a/go/vt/vtgate/engine/limit.go b/go/vt/vtgate/engine/limit.go
@@ -154,8 +154,8 @@ func (l *Limit) GetFields(ctx context.Context, vcursor VCursor, bindVars map[str
 }
 
 // Inputs returns the input to limit
-func (l *Limit) Inputs() []Primitive {
-	return []Primitive{l.Input}
+func (l *Limit) Inputs() ([]Primitive, []map[string]any) {
+	return []Primitive{l.Input}, nil
 }
 
 // NeedsTransaction implements the Primitive interface.

diff --git a/go/vt/vtgate/engine/memory_sort.go b/go/vt/vtgate/engine/memory_sort.go
@@ -150,8 +150,8 @@ func (ms *MemorySort) GetFields(ctx context.Context, vcursor VCursor, bindVars m
 }
 
 // Inputs returns the input to memory sort
-func (ms *MemorySort) Inputs() []Primitive {
-	return []Primitive{ms.Input}
+func (ms *MemorySort) Inputs() ([]Primitive, []map[string]any) {
+	return []Primitive{ms.Input}, nil
 }
 
 // NeedsTransaction implements the Primitive interface

diff --git a/go/vt/vtgate/engine/ordered_aggregate.go b/go/vt/vtgate/engine/ordered_aggregate.go
@@ -250,8 +250,8 @@ func (oa *OrderedAggregate) GetFields(ctx context.Context, vcursor VCursor, bind
 }
 
 // Inputs returns the Primitive input for this aggregation
-func (oa *OrderedAggregate) Inputs() []Primitive {
-	return []Primitive{oa.Input}
+func (oa *OrderedAggregate) Inputs() ([]Primitive, []map[string]any) {
+	return []Primitive{oa.Input}, nil
 }
 
 // NeedsTransaction implements the Primitive interface

diff --git a/go/vt/vtgate/engine/plan_description.go b/go/vt/vtgate/engine/plan_description.go
@@ -28,6 +28,8 @@ import (
 	"vitess.io/vitess/go/vt/vtgate/vindexes"
 )
 
+const inputName = "InputName"
+
 // PrimitiveDescription is used to create a serializable representation of the Primitive tree
 // Using this structure, all primitives can share json marshalling code, which gives us an uniform output
 type PrimitiveDescription struct {
@@ -41,7 +43,9 @@ type PrimitiveDescription struct {
 	// this is only used in conjunction with TargetDestination
 	TargetTabletType topodatapb.TabletType
 	Other            map[string]any
-	Inputs           []PrimitiveDescription
+
+	InputName string
+	Inputs    []PrimitiveDescription
 }
 
 // MarshalJSON serializes the PlanDescription into a JSON representation.
@@ -51,29 +55,37 @@ func (pd PrimitiveDescription) MarshalJSON() ([]byte, error) {
 	buf := &bytes.Buffer{}
 	buf.WriteString("{")
 
-	if err := marshalAdd("", buf, "OperatorType", pd.OperatorType); err != nil {
+	prepend := ""
+	if pd.InputName != "" {
+		if err := marshalAdd(prepend, buf, "InputName", pd.InputName); err != nil {
+			return nil, err
+		}
+		prepend = ","
+	}
+	if err := marshalAdd(prepend, buf, "OperatorType", pd.OperatorType); err != nil {
 		return nil, err
 	}
+	prepend = ","
 	if pd.Variant != "" {
-		if err := marshalAdd(",", buf, "Variant", pd.Variant); err != nil {
+		if err := marshalAdd(prepend, buf, "Variant", pd.Variant); err != nil {
 			return nil, err
 		}
 	}
 	if pd.Keyspace != nil {
-		if err := marshalAdd(",", buf, "Keyspace", pd.Keyspace); err != nil {
+		if err := marshalAdd(prepend, buf, "Keyspace", pd.Keyspace); err != nil {
 			return nil, err
 		}
 	}
 	if pd.TargetDestination != nil {
 		s := pd.TargetDestination.String()
 		dest := s[11:] // TODO: All these start with Destination. We should fix that instead if trimming it out here
 
-		if err := marshalAdd(",", buf, "TargetDestination", dest); err != nil {
+		if err := marshalAdd(prepend, buf, "TargetDestination", dest); err != nil {
 			return nil, err
 		}
 	}
 	if pd.TargetTabletType != topodatapb.TabletType_UNKNOWN {
-		if err := marshalAdd(",", buf, "TargetTabletType", pd.TargetTabletType.String()); err != nil {
+		if err := marshalAdd(prepend, buf, "TargetTabletType", pd.TargetTabletType.String()); err != nil {
 			return nil, err
 		}
 	}
@@ -83,7 +95,7 @@ func (pd PrimitiveDescription) MarshalJSON() ([]byte, error) {
 	}
 
 	if len(pd.Inputs) > 0 {
-		if err := marshalAdd(",", buf, "Inputs", pd.Inputs); err != nil {
+		if err := marshalAdd(prepend, buf, "Inputs", pd.Inputs); err != nil {
 			return nil, err
 		}
 	}
@@ -172,11 +184,25 @@ func marshalAdd(prepend string, buf *bytes.Buffer, name string, obj any) error {
 func PrimitiveToPlanDescription(in Primitive) PrimitiveDescription {
 	this := in.description()
 
-	for _, input := range in.Inputs() {
-		this.Inputs = append(this.Inputs, PrimitiveToPlanDescription(input))
+	inputs, infos := in.Inputs()
+	for idx, input := range inputs {
+		pd := PrimitiveToPlanDescription(input)
+		if infos != nil {
+			for k, v := range infos[idx] {
+				if k == inputName {
+					pd.InputName = v.(string)
+					continue
+				}
+				if pd.Other == nil {
+					pd.Other = map[string]any{}
+				}
+				pd.Other[k] = v
+			}
+		}
+		this.Inputs = append(this.Inputs, pd)
 	}
 
-	if len(in.Inputs()) == 0 {
+	if len(inputs) == 0 {
 		this.Inputs = []PrimitiveDescription{}
 	}
 

diff --git a/go/vt/vtgate/engine/primitive.go b/go/vt/vtgate/engine/primitive.go
@@ -223,7 +223,7 @@ type (
 		TryStreamExecute(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error
 
 		// Inputs is a slice containing the inputs to this Primitive
-		Inputs() []Primitive
+		Inputs() ([]Primitive, []map[string]any)
 
 		// description is the description, sans the inputs, of this Primitive.
 		// to get the plan description with all children, use PrimitiveToPlanDescription()
@@ -245,7 +245,8 @@ func Find(isMatch Match, start Primitive) Primitive {
 	if isMatch(start) {
 		return start
 	}
-	for _, input := range start.Inputs() {
+	inputs, _ := start.Inputs()
+	for _, input := range inputs {
 		result := Find(isMatch, input)
 		if result != nil {
 			return result
@@ -260,8 +261,8 @@ func Exists(m Match, p Primitive) bool {
 }
 
 // Inputs implements no inputs
-func (noInputs) Inputs() []Primitive {
-	return nil
+func (noInputs) Inputs() ([]Primitive, []map[string]any) {
+	return nil, nil
 }
 
 func (noTxNeeded) NeedsTransaction() bool {

diff --git a/go/vt/vtgate/engine/projection.go b/go/vt/vtgate/engine/projection.go
@@ -165,8 +165,8 @@ func (p *Projection) evalFields(env *evalengine.ExpressionEnv, infields []*query
 }
 
 // Inputs implements the Primitive interface
-func (p *Projection) Inputs() []Primitive {
-	return []Primitive{p.Input}
+func (p *Projection) Inputs() ([]Primitive, []map[string]any) {
+	return []Primitive{p.Input}, nil
 }
 
 // description implements the Primitive interface