diff --git a/go/vt/vtgate/planbuilder/operators/SQL_builder.go b/go/vt/vtgate/planbuilder/operators/SQL_builder.go index 248064b70fa..53aab5eb105 100644 --- a/go/vt/vtgate/planbuilder/operators/SQL_builder.go +++ b/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...) diff --git a/go/vt/vtgate/planbuilder/operators/apply_join.go b/go/vt/vtgate/planbuilder/operators/apply_join.go index 2297227a22d..d57b9a740f2 100644 --- a/go/vt/vtgate/planbuilder/operators/apply_join.go +++ b/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 } diff --git a/go/vt/vtgate/planbuilder/operators/join_merging.go b/go/vt/vtgate/planbuilder/operators/join_merging.go index cb3569cf79e..6952dffda0f 100644 --- a/go/vt/vtgate/planbuilder/operators/join_merging.go +++ b/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, } diff --git a/go/vt/vtgate/planbuilder/operators/predicates/predicate.go b/go/vt/vtgate/planbuilder/operators/predicates/predicate.go new file mode 100644 index 00000000000..1ec56b459fd --- /dev/null +++ b/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) +} diff --git a/go/vt/vtgate/planbuilder/operators/predicates/tracker.go b/go/vt/vtgate/planbuilder/operators/predicates/tracker.go new file mode 100644 index 00000000000..36bf4c35e12 --- /dev/null +++ b/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 +} diff --git a/go/vt/vtgate/planbuilder/operators/query_planning.go b/go/vt/vtgate/planbuilder/operators/query_planning.go index 935e4d4204d..2ff2527d729 100644 --- a/go/vt/vtgate/planbuilder/operators/query_planning.go +++ b/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. diff --git a/go/vt/vtgate/planbuilder/operators/route.go b/go/vt/vtgate/planbuilder/operators/route.go index 9aeafec2799..d41d168cf0f 100644 --- a/go/vt/vtgate/planbuilder/operators/route.go +++ b/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 { diff --git a/go/vt/vtgate/planbuilder/plancontext/planning_context.go b/go/vt/vtgate/planbuilder/plancontext/planning_context.go index 016f5c877cf..42be02ec5cf 100644 --- a/go/vt/vtgate/planbuilder/plancontext/planning_context.go +++ b/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 diff --git a/go/vt/vtgate/planbuilder/plancontext/planning_context_test.go b/go/vt/vtgate/planbuilder/plancontext/planning_context_test.go index b13c2112e3a..7016c07dc0f 100644 --- a/go/vt/vtgate/planbuilder/plancontext/planning_context_test.go +++ b/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(), } } diff --git a/go/vt/vtgate/planbuilder/simplifier_test.go b/go/vt/vtgate/planbuilder/simplifier_test.go index 012475ba021..f1bc6c337b4 100644 --- a/go/vt/vtgate/planbuilder/simplifier_test.go +++ b/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) diff --git a/go/vt/vtgate/planbuilder/testdata/onecase.json b/go/vt/vtgate/planbuilder/testdata/onecase.json index 9d653b2f6e9..09a83b3b5b7 100644 --- a/go/vt/vtgate/planbuilder/testdata/onecase.json +++ b/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" + ] } } ] \ No newline at end of file