wip

Signed-off-by: Florent Poinsard <florent.poinsard@outlook.fr>

go/vt/vtgate/planbuilder/operators/SQL_builder.go

@@ -610,11 +610,7 @@ func buildProjection(op *Projection, qb *queryBuilder) {
 
 func buildApplyJoin(op *ApplyJoin, qb *queryBuilder) {
 	predicates := slice.Map(op.JoinPredicates.columns, func(jc applyJoinColumn) sqlparser.Expr {
-		// since we are adding these join predicates, we need to mark to broken up version (RHSExpr) of it as done
-		err := qb.ctx.SkipJoinPredicates(jc.Original)
-		if err != nil {
-			panic(err)
-		}
+		qb.ctx.SkipJoinPredicatesTODO(jc.JoinPredicateID)
 		return jc.Original
 	})
 	pred := sqlparser.AndExpressions(predicates...)

go/vt/vtgate/planbuilder/operators/apply_join.go

@@ -22,6 +22,8 @@ import (
 	"slices"
 	"strings"
 
+	"vitess.io/vitess/go/vt/vtgate/planbuilder/operators/predicates"
+
 	"vitess.io/vitess/go/slice"
 	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vterrors"
@@ -69,10 +71,11 @@ type (
 	//     so they can be used for the result of this expression that is using data from both sides.
 	//     All fields will be used for these
 	applyJoinColumn struct {
-		Original sqlparser.Expr // this is the original expression being passed through
-		LHSExprs []BindVarExpr  // These are the expressions we are pushing to the left hand side which we'll receive as bind variables
-		RHSExpr  sqlparser.Expr // This the expression that we'll evaluate on the right hand side. This is nil, if the right hand side has nothing.
-		GroupBy  bool           // if this is true, we need to push this down to our inputs with addToGroupBy set to true
+		JoinPredicateID predicates.ID
+		Original        sqlparser.Expr // this is the original expression being passed through
+		LHSExprs        []BindVarExpr  // These are the expressions we are pushing to the left hand side which we'll receive as bind variables
+		RHSExpr         sqlparser.Expr // This the expression that we'll evaluate on the right hand side. This is nil, if the right hand side has nothing.
+		GroupBy         bool           // if this is true, we need to push this down to our inputs with addToGroupBy set to true
 	}
 
 	// BindVarExpr is an expression needed from one side of a join/subquery, and the argument name for it.
@@ -144,12 +147,15 @@ func (aj *ApplyJoin) AddJoinPredicate(ctx *plancontext.PlanningContext, expr sql
 		return
 	}
 	rhs := aj.RHS
-	predicates := sqlparser.SplitAndExpression(nil, expr)
-	for _, pred := range predicates {
-		col := breakExpressionInLHSandRHS(ctx, pred, TableID(aj.LHS))
+	preds := sqlparser.SplitAndExpression(nil, expr)
+	for _, pred := range preds {
+		lhsID := TableID(aj.LHS)
+		col := breakExpressionInLHSandRHS(ctx, pred, lhsID)
+		newPred := ctx.PredTracker.NewJoinPredicate(col.RHSExpr)
+		col.JoinPredicateID = newPred.ID
 		aj.JoinPredicates.add(col)
-		ctx.AddJoinPredicates(pred, col.RHSExpr)
-		rhs = rhs.AddPredicate(ctx, col.RHSExpr)
+		ctx.AddJoinPredicates(pred, newPred)
+		rhs = rhs.AddPredicate(ctx, newPred)
 	}
 	aj.RHS = rhs
 }

go/vt/vtgate/planbuilder/operators/join_merging.go

@@ -239,8 +239,12 @@ func (jm *joinMerger) getApplyJoin(ctx *plancontext.PlanningContext, op1, op2 *R
 }
 
 func (jm *joinMerger) merge(ctx *plancontext.PlanningContext, op1, op2 *Route, r Routing) *Route {
+	aj := jm.getApplyJoin(ctx, op1, op2)
+	for _, column := range aj.JoinPredicates.columns {
+		ctx.PredTracker.Set(column.JoinPredicateID, column.Original)
+	}
 	return &Route{
-		unaryOperator: newUnaryOp(jm.getApplyJoin(ctx, op1, op2)),
+		unaryOperator: newUnaryOp(aj),
 		MergedWith:    []*Route{op2},
 		Routing:       r,
 	}

go/vt/vtgate/planbuilder/operators/predicates/predicate.go

@@ -0,0 +1,50 @@
+/*
+Copyright 2025 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package predicates
+
+import (
+	"vitess.io/vitess/go/vt/sqlparser"
+)
+
+// JoinPredicate represents a condition relating two parts of the query,
+// typically used when expressions span multiple tables. We track it in
+// different shapes so we can rewrite or push it differently during plan generation.
+type JoinPredicate struct {
+	ID      ID
+	tracker *Tracker
+}
+
+var _ sqlparser.Expr = (*JoinPredicate)(nil)
+var _ sqlparser.Visitable = (*JoinPredicate)(nil)
+
+func (j *JoinPredicate) VisitThis() sqlparser.SQLNode {
+	return j.Current()
+}
+
+func (j *JoinPredicate) Current() sqlparser.Expr {
+	return j.tracker.Expressions[j.ID]
+}
+
+func (j *JoinPredicate) IsExpr() {}
+
+func (j *JoinPredicate) Format(buf *sqlparser.TrackedBuffer) {
+	j.Current().Format(buf)
+}
+
+func (j *JoinPredicate) FormatFast(buf *sqlparser.TrackedBuffer) {
+	j.Current().FormatFast(buf)
+}

go/vt/vtgate/planbuilder/operators/predicates/tracker.go

@@ -0,0 +1,66 @@
+/*
+Copyright 2025 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package predicates
+
+import (
+	"sync"
+
+	"vitess.io/vitess/go/vt/sqlparser"
+)
+
+type (
+	// Tracker manages a global mapping of expression IDs to their "Shape".
+	// This allows the same logical expression to take different forms (shapes)
+	// depending on pushdown or join strategies. We lock around 'lastID' to ensure
+	// unique IDs in concurrent planning contexts.
+	Tracker struct {
+		mu          sync.Mutex
+		lastID      ID
+		Expressions map[ID]sqlparser.Expr
+	}
+
+	// ID is a unique key that references the shape of a single expression.
+	// We use it so multiple references to an expression can share the same shape entry.
+	ID int
+)
+
+func NewTracker() *Tracker {
+	return &Tracker{
+		Expressions: make(map[ID]sqlparser.Expr),
+	}
+}
+
+func (t *Tracker) NewJoinPredicate(org sqlparser.Expr) *JoinPredicate {
+	nextID := t.NextID()
+	t.Expressions[nextID] = org
+	return &JoinPredicate{
+		ID:      nextID,
+		tracker: t,
+	}
+}
+
+func (t *Tracker) NextID() ID {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	id := t.lastID
+	t.lastID++
+	return id
+}
+
+func (t *Tracker) Set(id ID, expr sqlparser.Expr) {
+	t.Expressions[id] = expr
+}

go/vt/vtgate/planbuilder/operators/query_planning.go

@@ -333,6 +333,9 @@ func mergeOffsetExpressions(e1, e2 sqlparser.Expr) (expr sqlparser.Expr, failed
 }
 
 // mergeLimitExpressions merges two LIMIT expressions with an OFFSET expression.
+
+// select tbl1.foo = tbl2.bar from tbl1, tbl2 where tbl1.foo = tbl2.bar
+
 // l1: First LIMIT expression.
 // l2: Second LIMIT expression.
 // off2: Second OFFSET expression.

go/vt/vtgate/planbuilder/operators/route.go

@@ -26,6 +26,7 @@ import (
 	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vtgate/engine"
 	"vitess.io/vitess/go/vt/vtgate/evalengine"
+	"vitess.io/vitess/go/vt/vtgate/planbuilder/operators/predicates"
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
 	"vitess.io/vitess/go/vt/vtgate/semantics"
 	"vitess.io/vitess/go/vt/vtgate/vindexes"
@@ -127,6 +128,11 @@ func UpdateRoutingLogic(ctx *plancontext.PlanningContext, expr sqlparser.Expr, r
 		return r.updateRoutingLogic(ctx, expr)
 	}
 
+	// If we have a JoinPredicate, get it, otherwise do nothing
+	if pred, ok := expr.(*predicates.JoinPredicate); ok {
+		expr = pred.Current()
+	}
+
 	// For some expressions, even if we can't evaluate them, we know that they will always return false or null
 	cmp, ok := expr.(*sqlparser.ComparisonExpr)
 	if !ok {

go/vt/vtgate/planbuilder/plancontext/planning_context.go

@@ -19,6 +19,8 @@ package plancontext
 import (
 	"io"
 
+	"vitess.io/vitess/go/vt/vtgate/planbuilder/operators/predicates"
+
 	"vitess.io/vitess/go/sqltypes"
 	querypb "vitess.io/vitess/go/vt/proto/query"
 	"vitess.io/vitess/go/vt/sqlparser"
@@ -41,7 +43,8 @@ type PlanningContext struct {
 
 	// skipPredicates tracks predicates that should be skipped, typically when
 	// a join predicate is reverted to its original form during planning.
-	skipPredicates map[sqlparser.Expr]any
+	skipPredicates     map[sqlparser.Expr]any
+	skipJoinPredicates map[predicates.ID]any
 
 	PlannerVersion querypb.ExecuteOptions_PlannerVersion
 
@@ -79,6 +82,8 @@ type PlanningContext struct {
 
 	emptyEnv    *evalengine.ExpressionEnv
 	constantCfg *evalengine.Config
+
+	PredTracker *predicates.Tracker
 }
 
 // CreatePlanningContext initializes a new PlanningContext with the given parameters.
@@ -104,14 +109,16 @@ func CreatePlanningContext(stmt sqlparser.Statement,
 	vschema.PlannerWarning(semTable.Warning)
 
 	return &PlanningContext{
-		ReservedVars:      reservedVars,
-		SemTable:          semTable,
-		VSchema:           vschema,
-		joinPredicates:    map[sqlparser.Expr][]sqlparser.Expr{},
-		skipPredicates:    map[sqlparser.Expr]any{},
-		PlannerVersion:    version,
-		ReservedArguments: map[sqlparser.Expr]string{},
-		Statement:         stmt,
+		ReservedVars:       reservedVars,
+		SemTable:           semTable,
+		VSchema:            vschema,
+		joinPredicates:     map[sqlparser.Expr][]sqlparser.Expr{},
+		skipPredicates:     map[sqlparser.Expr]any{},
+		skipJoinPredicates: map[predicates.ID]any{},
+		PlannerVersion:     version,
+		ReservedArguments:  map[sqlparser.Expr]string{},
+		Statement:          stmt,
+		PredTracker:        predicates.NewTracker(),
 	}, nil
 }
 
@@ -146,7 +153,11 @@ func (ctx *PlanningContext) ShouldSkip(expr sqlparser.Expr) bool {
 			return true
 		}
 	}
-	return false
+	var found bool
+	if jp, ok := expr.(*predicates.JoinPredicate); ok {
+		_, found = ctx.skipJoinPredicates[jp.ID]
+	}
+	return found
 }
 
 // AddJoinPredicates associates additional RHS predicates with an existing join predicate.
@@ -176,6 +187,10 @@ func (ctx *PlanningContext) SkipJoinPredicates(joinPred sqlparser.Expr) error {
 	return vterrors.VT13001("predicate does not exist: " + sqlparser.String(joinPred))
 }
 
+func (ctx *PlanningContext) SkipJoinPredicatesTODO(id predicates.ID) {
+	ctx.skipJoinPredicates[id] = struct{}{}
+}
+
 // KeepPredicateInfo transfers join predicate information from another context.
 // This is useful when nesting queries, ensuring consistent predicate handling across contexts.
 func (ctx *PlanningContext) KeepPredicateInfo(other *PlanningContext) {
@@ -456,6 +471,7 @@ func (ctx *PlanningContext) UseMirror() *PlanningContext {
 		OuterTables:       ctx.OuterTables,
 		CurrentCTE:        ctx.CurrentCTE,
 		emptyEnv:          ctx.emptyEnv,
+		PredTracker:       ctx.PredTracker,
 		isMirrored:        true,
 	}
 	return ctx.mirror

go/vt/vtgate/planbuilder/plancontext/planning_context_test.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"fmt"
 	"testing"
+	"vitess.io/vitess/go/vt/vtgate/planbuilder/operators/predicates"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -181,6 +182,7 @@ func createPlanContext(st *semantics.SemTable) *PlanningContext {
 		skipPredicates:    map[sqlparser.Expr]any{},
 		ReservedArguments: map[sqlparser.Expr]string{},
 		VSchema:           &vschema{},
+		PredTracker:       predicates.NewTracker(),
 	}
 }
 

go/vt/vtgate/planbuilder/simplifier_test.go

@@ -61,8 +61,8 @@ func TestSimplifyBuggyQuery(t *testing.T) {
 }
 
 func TestSimplifyPanic(t *testing.T) {
-	t.Skip("not needed to run")
-	query := "(select id from unsharded union select id from unsharded_auto) union (select id from unsharded_auto union select name from unsharded)"
+	// t.Skip("not needed to run")
+	query := "select t.id from (select id, textcol1 as baz from route1) as t join (select id, textcol1+textcol1 as baz from user) as s ON t.id = s.id WHERE t.baz = '3' AND s.baz = '3'"
 
 	env := vtenv.NewTestEnv()
 	vschema := loadSchema(t, "vschemas/schema.json", true)

go/vt/vtgate/planbuilder/testdata/onecase.json

@@ -1,8 +1,25 @@
 [
   {
     "comment": "Add your test case here for debugging and run go test -run=One.",
-    "query": "",
+    "query": "select user.col from user_extra left outer join user on user_extra.user_id = user.id WHERE user.id IS NULL",
     "plan": {
+      "QueryType": "SELECT",
+      "Original": "select user.col from user_extra left outer join user on user_extra.user_id = user.id WHERE user.id IS NULL",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "Scatter",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select `user`.col from user_extra left join `user` on user_extra.user_id = `user`.id where 1 != 1",
+        "Query": "select `user`.col from user_extra left join `user` on user_extra.user_id = `user`.id where `user`.id is null",
+        "Table": "`user`, user_extra"
+      },
+      "TablesUsed": [
+        "user.user",
+        "user.user_extra"
+      ]
     }
   }
 ]