feat: fix pushing of values and filters to a value join

Signed-off-by: Manan Gupta <manan@planetscale.com>
planetscale · Feb 21, 2025 · c558b3f · c558b3f
1 parent df126a2
commit c558b3f
Show file tree

Hide file tree

Showing 6 changed files with 42 additions and 6 deletions.
diff --git a/go/vt/vtgate/planbuilder/operators/phases.go b/go/vt/vtgate/planbuilder/operators/phases.go
@@ -157,6 +157,7 @@ func newValuesJoin(ctx *plancontext.PlanningContext, lhs, rhs Operator, joinType
 	v := &Values{
 		unaryOperator: newUnaryOp(rhs),
 		Name:          bindVariableName,
+		TableID:       TableID(lhs),
 	}
 	return &ValuesJoin{
 		binaryOperator:    newBinaryOp(lhs, v),

diff --git a/go/vt/vtgate/planbuilder/operators/query_planning.go b/go/vt/vtgate/planbuilder/operators/query_planning.go
@@ -130,8 +130,7 @@ func tryPushValues(ctx *plancontext.PlanningContext, in *Values) (Operator, *App
 	switch src := in.Source.(type) {
 	case *ValuesJoin:
 		src.LHS = in.Clone([]Operator{src.LHS})
-		src.RHS = in.Clone([]Operator{src.RHS})
-		return src, Rewrote("pushed values under value join")
+		return src, Rewrote("pushed values to the LHS of values join")
 	case *Filter:
 		return Swap(in, src, "pushed values under filter")
 	case *Route:
@@ -724,7 +723,7 @@ func tryPushFilter(ctx *plancontext.PlanningContext, in *Filter) (Operator, *App
 		return src, Rewrote("push filter to outer query in subquery container")
 	case *ValuesJoin:
 		for _, pred := range in.Predicates {
-			src.AddJoinPredicate(ctx, pred)
+			src.AddPredicate(ctx, pred)
 		}
 		return src, Rewrote("pushed filter predicates through values join")
 	}

diff --git a/go/vt/vtgate/planbuilder/operators/values.go b/go/vt/vtgate/planbuilder/operators/values.go
@@ -21,12 +21,14 @@ import (
 	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
+	"vitess.io/vitess/go/vt/vtgate/semantics"
 )
 
 type Values struct {
 	unaryOperator
 
-	Name string
+	Name    string
+	TableID semantics.TableSet
 }
 
 func (v *Values) Clone(inputs []Operator) Operator {
@@ -97,3 +99,7 @@ func (v *Values) ShortDescription() string {
 func (v *Values) GetOrdering(ctx *plancontext.PlanningContext) []OrderBy {
 	return v.Source.GetOrdering(ctx)
 }
+
+func (v *Values) introducesTableID() semantics.TableSet {
+	return v.TableID
+}
diff --git a/go/vt/vtgate/planbuilder/operators/values_join.go b/go/vt/vtgate/planbuilder/operators/values_join.go
@@ -117,6 +117,10 @@ func (vj *ValuesJoin) MakeInner() {
 	// no-op for values-join
 }
 
+func (vj *ValuesJoin) AddPredicate(ctx *plancontext.PlanningContext, expr sqlparser.Expr) Operator {
+	return AddPredicate(ctx, vj, expr, false, newFilterSinglePredicate)
+}
+
 func (vj *ValuesJoin) IsInner() bool {
 	return true
 }
@@ -160,7 +164,7 @@ func (vj *ValuesJoin) GetOrdering(ctx *plancontext.PlanningContext) []OrderBy {
 }
 
 func (vj *ValuesJoin) planOffsets(ctx *plancontext.PlanningContext) Operator {
-	exprs := ctx.ValuesJoinColumns[vj.ValuesDestination]
+	exprs := ctx.GetColumns(vj.ValuesDestination)
 	for _, jc := range vj.JoinColumns {
 		newExprs := vj.planOffsetsForLHSExprs(ctx, jc.LHS)
 		exprs = append(exprs, newExprs...)
@@ -172,7 +176,7 @@ func (vj *ValuesJoin) planOffsets(ctx *plancontext.PlanningContext) Operator {
 		newExprs := vj.planOffsetsForLHSExprs(ctx, jc.LHS)
 		exprs = append(exprs, newExprs...)
 	}
-	ctx.ValuesJoinColumns[vj.ValuesDestination] = exprs
+	ctx.SetColumns(vj.ValuesDestination, exprs)
 	return vj
 }
 

diff --git a/go/vt/vtgate/planbuilder/plancontext/planning_context.go b/go/vt/vtgate/planbuilder/plancontext/planning_context.go
@@ -474,6 +474,15 @@ func (ctx *PlanningContext) ActiveCTE() *ContextCTE {
 	return ctx.CurrentCTE[len(ctx.CurrentCTE)-1]
 }
 
+func (ctx *PlanningContext) GetColumns(joinName string) []*sqlparser.AliasedExpr {
+	valuesName := ctx.ValueJoins[joinName]
+	return ctx.ValuesJoinColumns[valuesName]
+}
+func (ctx *PlanningContext) SetColumns(joinName string, cols []*sqlparser.AliasedExpr) {
+	valuesName := ctx.ValueJoins[joinName]
+	ctx.ValuesJoinColumns[valuesName] = cols
+}
+
 func (ctx *PlanningContext) UseMirror() *PlanningContext {
 	if ctx.isMirrored {
 		panic(vterrors.VT13001("cannot mirror already mirrored planning context"))

diff --git a/go/vt/vtgate/semantics/table_set.go b/go/vt/vtgate/semantics/table_set.go
@@ -18,6 +18,7 @@ package semantics
 
 import (
 	"fmt"
+	"strings"
 
 	"vitess.io/vitess/go/vt/vtgate/semantics/bitset"
 )
@@ -41,6 +42,22 @@ func (ts TableSet) Format(f fmt.State, verb rune) {
 	fmt.Fprintf(f, "}")
 }
 
+func (ts TableSet) DebugString() string {
+	var f strings.Builder
+	first := true
+	f.WriteString("TableSet{")
+	bitset.Bitset(ts).ForEach(func(tid int) {
+		if first {
+			f.WriteString(fmt.Sprintf("%d", tid))
+			first = false
+		} else {
+			f.WriteString(fmt.Sprintf(",%d", tid))
+		}
+	})
+	f.WriteString("}")
+	return f.String()
+}
+
 // IsOverlapping returns true if at least one table exists in both sets
 func (ts TableSet) IsOverlapping(other TableSet) bool {
 	return bitset.Bitset(ts).Overlaps(bitset.Bitset(other))