add more information to vexplain keys

Signed-off-by: Andres Taylor <[email protected]>

go/vt/sqlparser/ast_funcs.go

@@ -1498,6 +1498,25 @@ func (op ComparisonExprOperator) ToString() string {
 	}
 }
 
+// JSONString returns a string representation for this operator that does not need escaping in JSON
+func (op ComparisonExprOperator) JSONString() string {
+	switch op {
+	case EqualOp, NotEqualOp, NullSafeEqualOp, InOp, NotInOp, LikeOp, NotLikeOp, RegexpOp, NotRegexpOp:
+		// These operators are safe for JSON output, so we delegate to ToString
+		return op.ToString()
+	case LessThanOp:
+		return "lt"
+	case GreaterThanOp:
+		return "gt"
+	case LessEqualOp:
+		return "le"
+	case GreaterEqualOp:
+		return "ge"
+	default:
+		panic("unreachable")
+	}
+}
+
 // ToString returns the operator as a string
 func (op IsExprOperator) ToString() string {
 	switch op {

go/vt/sqlparser/constants.go

@@ -724,6 +724,29 @@ func (op ComparisonExprOperator) Inverse() ComparisonExprOperator {
 	panic("unreachable")
 }
 
+// SwitchSides returns the reversed comparison operator if applicable, along with a boolean indicating success.
+// For symmetric operators like '=', '!=', and '<=>', it returns the same operator and true.
+// For directional comparison operators ('<', '>', '<=', '>='), it returns the opposite operator and true.
+// For operators that imply directionality or cannot be logically reversed (such as 'IN', 'LIKE', 'REGEXP'),
+// it returns the original operator and false, indicating that switching sides is not valid.
+func (op ComparisonExprOperator) SwitchSides() (ComparisonExprOperator, bool) {
+	switch op {
+	case EqualOp, NotEqualOp, NullSafeEqualOp, LikeOp, NotLikeOp, RegexpOp, NotRegexpOp, InOp, NotInOp:
+		// These operators are symmetric, so switching sides has no effect
+		return op, true
+	case LessThanOp:
+		return GreaterThanOp, true
+	case GreaterThanOp:
+		return LessThanOp, true
+	case LessEqualOp:
+		return GreaterEqualOp, true
+	case GreaterEqualOp:
+		return LessEqualOp, true
+	default:
+		return op, false
+	}
+}
+
 func (op ComparisonExprOperator) IsCommutative() bool {
 	switch op {
 	case EqualOp, NotEqualOp, NullSafeEqualOp:

go/vt/vtgate/executor_vexplain_test.go

@@ -122,54 +122,61 @@ func TestVExplainKeys(t *testing.T) {
 		{
 			query: "select count(*), col2 from music group by col2",
 			expectedRowString: `{
+	"statementType": "SELECT",
 	"groupingColumns": [
 		"music.col2"
 	],
-	"statementType": "SELECT"
+	"selectColumns": [
+		"music.col2"
+	]
 }`,
 		}, {
 			query: "select * from user u join user_extra ue on u.id = ue.user_id where u.col1 > 100 and ue.noLimit = 'foo'",
 			expectedRowString: `{
+	"statementType": "SELECT",
 	"joinColumns": [
-		"user.id",
-		"user_extra.user_id"
+		"user.id =",
+		"user_extra.user_id ="
 	],
 	"filterColumns": [
-		"user.col1",
-		"user_extra.noLimit"
-	],
-	"statementType": "SELECT"
+		"user.col1 gt",
+		"user_extra.noLimit ="
+	]
 }`,
 		}, {
 			// same as above, but written differently
 			query: "select * from user_extra ue, user u where ue.noLimit = 'foo' and u.col1 > 100 and ue.user_id = u.id",
 			expectedRowString: `{
+	"statementType": "SELECT",
 	"joinColumns": [
-		"user.id",
-		"user_extra.user_id"
+		"user.id =",
+		"user_extra.user_id ="
 	],
 	"filterColumns": [
-		"user.col1",
-		"user_extra.noLimit"
-	],
-	"statementType": "SELECT"
+		"user.col1 gt",
+		"user_extra.noLimit ="
+	]
 }`,
 		},
 		{
 			query: "select u.foo, ue.bar, count(*) from user u join user_extra ue on u.id = ue.user_id where u.name = 'John Doe' group by 1, 2",
 			expectedRowString: `{
+	"statementType": "SELECT",
 	"groupingColumns": [
 		"user.foo",
 		"user_extra.bar"
 	],
 	"joinColumns": [
-		"user.id",
-		"user_extra.user_id"
+		"user.id =",
+		"user_extra.user_id ="
 	],
 	"filterColumns": [
-		"user.name"
+		"user.name ="
 	],
-	"statementType": "SELECT"
+	"selectColumns": [
+		"user.foo",
+		"user_extra.bar"
+	]
 }`,
 		},
 		{
@@ -181,47 +188,54 @@ func TestVExplainKeys(t *testing.T) {
 		{
 			query: "select name, sum(amount) from user group by name",
 			expectedRowString: `{
+	"statementType": "SELECT",
 	"groupingColumns": [
 		"user.name"
 	],
-	"statementType": "SELECT"
+	"selectColumns": [
+		"user.amount",
+		"user.name"
+	]
 }`,
 		},
 		{
 			query: "select name from user where age > 30",
 			expectedRowString: `{
+	"statementType": "SELECT",
 	"filterColumns": [
-		"user.age"
+		"user.age gt"
 	],
-	"statementType": "SELECT"
+	"selectColumns": [
+		"user.name"
+	]
 }`,
 		},
 		{
 			query: "select * from user where name = 'apa' union select * from user_extra where name = 'monkey'",
 			expectedRowString: `{
+	"statementType": "SELECT",
 	"filterColumns": [
-		"user.name",
-		"user_extra.name"
-	],
-	"statementType": "SELECT"
+		"user.name =",
+		"user_extra.name ="
+	]
 }`,
 		},
 		{
 			query: "update user set name = 'Jane Doe' where id = 1",
 			expectedRowString: `{
+	"statementType": "UPDATE",
 	"filterColumns": [
-		"user.id"
-	],
-	"statementType": "UPDATE"
+		"user.id ="
+	]
 }`,
 		},
 		{
 			query: "delete from user where order_date < '2023-01-01'",
 			expectedRowString: `{
+	"statementType": "DELETE",
 	"filterColumns": [
-		"user.order_date"
-	],
-	"statementType": "DELETE"
+		"user.order_date lt"
+	]
 }`,
 		},
 	}

go/vt/vtgate/planbuilder/operators/keys.go

@@ -17,23 +17,54 @@ limitations under the License.
 package operators
 
 import (
+	"encoding/json"
 	"fmt"
 	"slices"
+	"sort"
+
+	"vitess.io/vitess/go/slice"
 
 	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
 )
 
-type VExplainKeys struct {
-	GroupingColumns []string `json:"groupingColumns,omitempty"`
-	TableName       []string `json:"tableName,omitempty"`
-	JoinColumns     []string `json:"joinColumns,omitempty"`
-	FilterColumns   []string `json:"filterColumns,omitempty"`
-	StatementType   string   `json:"statementType"`
+type (
+	Column struct {
+		Table string
+		Name  string
+	}
+
+	ColumnUse struct {
+		Column Column
+		Uses   sqlparser.ComparisonExprOperator
+	}
+
+	VExplainKeys struct {
+		StatementType   string
+		TableName       []string
+		GroupingColumns []Column
+		JoinColumns     []ColumnUse
+		FilterColumns   []ColumnUse
+		SelectColumns   []Column
+	}
+)
+
+func (c Column) String() string {
+	return fmt.Sprintf("%s.%s", c.Table, c.Name)
+}
+
+func (c ColumnUse) String() string {
+	return fmt.Sprintf("%s %s", c.Column, c.Uses.JSONString())
+}
+
+type columnUse struct {
+	col *sqlparser.ColName
+	use sqlparser.ComparisonExprOperator
 }
 
 func GetVExplainKeys(ctx *plancontext.PlanningContext, stmt sqlparser.Statement) (result VExplainKeys) {
-	var filterColumns, joinColumns, groupingColumns []*sqlparser.ColName
+	var groupingColumns, selectColumns []*sqlparser.ColName
+	var filterColumns, joinColumns []columnUse
 
 	addPredicate := func(predicate sqlparser.Expr) {
 		predicates := sqlparser.SplitAndExpression(nil, predicate)
@@ -44,15 +75,19 @@ func GetVExplainKeys(ctx *plancontext.PlanningContext, stmt sqlparser.Statement)
 			}
 			lhs, lhsOK := cmp.Left.(*sqlparser.ColName)
 			rhs, rhsOK := cmp.Right.(*sqlparser.ColName)
+
+			var output = &filterColumns
 			if lhsOK && rhsOK && ctx.SemTable.RecursiveDeps(lhs) != ctx.SemTable.RecursiveDeps(rhs) {
-				joinColumns = append(joinColumns, lhs, rhs)
-				continue
+				// If the columns are from different tables, they are considered join columns
+				output = &joinColumns
 			}
+
 			if lhsOK {
-				filterColumns = append(filterColumns, lhs)
+				*output = append(*output, columnUse{lhs, cmp.Operator})
 			}
-			if rhsOK {
-				filterColumns = append(filterColumns, rhs)
+
+			if switchedOp, ok := cmp.Operator.SwitchSides(); rhsOK && ok {
+				*output = append(*output, columnUse{rhs, switchedOp})
 			}
 		}
 	}
@@ -65,30 +100,34 @@ func GetVExplainKeys(ctx *plancontext.PlanningContext, stmt sqlparser.Statement)
 			addPredicate(node.On)
 		case *sqlparser.GroupBy:
 			for _, expr := range node.Exprs {
-				predicates := sqlparser.SplitAndExpression(nil, expr)
-				for _, expr := range predicates {
-					col, ok := expr.(*sqlparser.ColName)
-					if ok {
-						groupingColumns = append(groupingColumns, col)
-					}
+				col, ok := expr.(*sqlparser.ColName)
+				if ok {
+					groupingColumns = append(groupingColumns, col)
 				}
 			}
+		case *sqlparser.AliasedExpr:
+			_ = sqlparser.VisitSQLNode(node, func(e sqlparser.SQLNode) (kontinue bool, err error) {
+				if col, ok := e.(*sqlparser.ColName); ok {
+					selectColumns = append(selectColumns, col)
+				}
+				return true, nil
+			})
 		}
 
 		return true, nil
 	})
 
 	return VExplainKeys{
+		SelectColumns:   getUniqueColNames(ctx, selectColumns),
 		GroupingColumns: getUniqueColNames(ctx, groupingColumns),
-		JoinColumns:     getUniqueColNames(ctx, joinColumns),
-		FilterColumns:   getUniqueColNames(ctx, filterColumns),
+		JoinColumns:     getUniqueColUsages(ctx, joinColumns),
+		FilterColumns:   getUniqueColUsages(ctx, filterColumns),
 		StatementType:   sqlparser.ASTToStatementType(stmt).String(),
 	}
 }
 
-func getUniqueColNames(ctx *plancontext.PlanningContext, columns []*sqlparser.ColName) []string {
-	var colNames []string
-	for _, col := range columns {
+func getUniqueColNames(ctx *plancontext.PlanningContext, inCols []*sqlparser.ColName) (columns []Column) {
+	for _, col := range inCols {
 		tableInfo, err := ctx.SemTable.TableInfoForExpr(col)
 		if err != nil {
 			continue
@@ -97,9 +136,56 @@ func getUniqueColNames(ctx *plancontext.PlanningContext, columns []*sqlparser.Co
 		if table == nil {
 			continue
 		}
-		colNames = append(colNames, fmt.Sprintf("%s.%s", table.Name.String(), col.Name.String()))
+		columns = append(columns, Column{Table: table.Name.String(), Name: col.Name.String()})
+	}
+	sort.Slice(columns, func(i, j int) bool {
+		return columns[i].String() < columns[j].String()
+	})
+
+	return slices.Compact(columns)
+}
+
+func getUniqueColUsages(ctx *plancontext.PlanningContext, inCols []columnUse) (columns []ColumnUse) {
+	for _, col := range inCols {
+		tableInfo, err := ctx.SemTable.TableInfoForExpr(col.col)
+		if err != nil {
+			continue
+		}
+		table := tableInfo.GetVindexTable()
+		if table == nil {
+			continue
+		}
+
+		columns = append(columns, ColumnUse{
+			Column: Column{Table: table.Name.String(), Name: col.col.Name.String()},
+			Uses:   col.use,
+		})
+	}
+
+	sort.Slice(columns, func(i, j int) bool {
+		return columns[i].Column.String() < columns[j].Column.String()
+	})
+	return slices.Compact(columns)
+}
+
+func (v VExplainKeys) MarshalJSON() ([]byte, error) {
+	// Create a custom struct to marshal with conditional fields
+	aux := struct {
+		StatementType   string   `json:"statementType"`
+		TableName       []string `json:"tableName,omitempty"`
+		GroupingColumns []string `json:"groupingColumns,omitempty"`
+		JoinColumns     []string `json:"joinColumns,omitempty"`
+		FilterColumns   []string `json:"filterColumns,omitempty"`
+		SelectColumns   []string `json:"selectColumns,omitempty"`
+	}{
+		StatementType:   v.StatementType,
+		TableName:       v.TableName,
+		SelectColumns:   slice.Map(v.SelectColumns, func(c Column) string { return c.String() }),
+		GroupingColumns: slice.Map(v.GroupingColumns, func(c Column) string { return c.String() }),
+		JoinColumns:     slice.Map(v.JoinColumns, func(c ColumnUse) string { return c.String() }),
+		FilterColumns:   slice.Map(v.FilterColumns, func(c ColumnUse) string { return c.String() }),
 	}
 
-	slices.Sort(colNames)
-	return slices.Compact(colNames)
+	// Marshal the aux struct into JSON
+	return json.Marshal(aux)
 }