vtgate: merge empty reference branches through rewritten copies

go/test/endtoend/vtgate/plan_tests/plan_e2e_test.go

@@ -29,9 +29,12 @@ func TestE2ECases(t *testing.T) {
 	err := utils.WaitForAuthoritative(t, "main", "source_of_ref", clusterInstance.VtgateProcess.ReadVSchema)
 	require.NoError(t, err)
 
+	// union_cases must run before dml_cases: the dml cases empty user_extra,
+	// and several union cases need its rows to reach their join primitives.
 	e2eTestCaseFiles := []string{
 		"select_cases.json",
 		"filter_cases.json",
+		"union_cases.json",
 		"dml_cases.json",
 		"reference_cases.json",
 	}

go/test/endtoend/vtgate/queries/union/union_test.go

@@ -67,6 +67,8 @@ func TestUnionDistinct(t *testing.T) {
 			mcmp.AssertMatchesNoOrder("select id1, id2 from t1 union select 827, 452 union select id3,id4 from t2",
 				"[[INT64(4) INT64(4)] [INT64(1) INT64(1)] [INT64(2) INT64(2)] [INT64(3) INT64(3)] [INT64(827) INT64(452)] [INT64(2) INT64(3)] [INT64(3) INT64(4)] [INT64(5) INT64(5)]]")
 			mcmp.AssertMatches("select 1 from dual where 1 IN (select 1 as col union select 2)", "[[INT64(1)]]")
+			mcmp.AssertMatches("select 1 from dual union select id1 from t1 where id1 in (null)", "[[INT64(1)]]")
+			mcmp.AssertMatches("select id1 from t1 where id1 in (null) union select 1 from dual", "[[INT64(1)]]")
 			mcmp.AssertMatches(`SELECT 1 from t1 UNION SELECT 2 from t1`, `[[INT64(1)] [INT64(2)]]`)
 			mcmp.AssertMatches(`SELECT 5 from t1 UNION SELECT 6 from t1`, `[[INT64(5)] [INT64(6)]]`)
 			mcmp.AssertMatchesNoOrder(`SELECT id1 from t1 UNION SELECT id2 from t1`, `[[INT64(1)] [INT64(2)] [INT64(3)] [INT64(4)]]`)

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

@@ -616,7 +616,7 @@ func buildProjection(op *Projection, qb *queryBuilder) {
 func buildApplyJoin(op *ApplyJoin, qb *queryBuilder) {
 	preds := slice.Map(op.JoinPredicates.columns, func(jc applyJoinColumn) sqlparser.Expr {
 		if jc.JoinPredicateID != nil {
-			qb.ctx.PredTracker.Skip(*jc.JoinPredicateID)
+			qb.ctx.PredTracker.SkipWithDescendants(*jc.JoinPredicateID)
 		}
 		return jc.Original
 	})

go/vt/vtgate/planbuilder/operators/cte_merging.go

@@ -22,6 +22,20 @@ import (
 )
 
 func tryMergeRecurse(ctx *plancontext.PlanningContext, in *RecurseCTE) (Operator, *ApplyResult) {
+	// A recursive predicate that was pushed into every source of a UNION inside
+	// the term leaves per-source copies behind. Those copies cannot be restored
+	// to column form (their column belongs to another scope) and cannot be
+	// skipped (they carry the recursion condition), so mergeCTE would emit them
+	// as arguments no primitive produces. Keep the RecurseCTE primitive instead.
+	for _, predicate := range in.Predicates {
+		if predicate.JoinPredicateID == nil {
+			continue
+		}
+		if len(ctx.PredTracker.DescendantIDs(*predicate.JoinPredicateID)) > 0 {
+			return in, NoRewrite
+		}
+	}
+
 	op := tryMergeCTE(ctx, in.Seed(), in.Term(), in)
 	if op == nil {
 		return in, NoRewrite
@@ -84,7 +98,7 @@ func mergeCTE(ctx *plancontext.PlanningContext, seed, term *Route, r Routing, in
 	newTerm, _ := expandHorizon(ctx, hz)
 	for _, predicate := range in.Predicates {
 		if predicate.JoinPredicateID != nil {
-			ctx.PredTracker.Set(*predicate.JoinPredicateID, predicate.Original)
+			ctx.PredTracker.ResetToOriginal(*predicate.JoinPredicateID, predicate.Original)
 		}
 	}
 

go/vt/vtgate/planbuilder/operators/misc_routing.go

@@ -29,6 +29,11 @@ type (
 	// Can be merged with any other route going to the same keyspace
 	NoneRouting struct {
 		keyspace *vindexes.Keyspace
+
+		// inferredKeyspace is set when the routing this replaced had no keyspace
+		// of its own (information_schema, dual): keyspace is then only a
+		// placeholder giving the engine route a target, not a genuine read.
+		inferredKeyspace bool
 	}
 
 	// TargetedRouting is used when the user has used syntax to target the

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

@@ -59,10 +59,14 @@ func (j *JoinPredicate) Current() sqlparser.Expr {
 func (j *JoinPredicate) IsExpr() {}
 
 func (j *JoinPredicate) Format(buf *sqlparser.TrackedBuffer) {
-	j.Current().Format(buf)
+	if expr := j.Current(); expr != nil {
+		expr.Format(buf)
+	}
 }
 
 func (j *JoinPredicate) FormatFast(buf *sqlparser.TrackedBuffer) {
 	fmt.Fprintf(buf, "JP(%d):", j.ID)
-	j.Current().FormatFast(buf)
+	if expr := j.Current(); expr != nil {
+		expr.FormatFast(buf)
+	}
 }

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

@@ -17,9 +17,10 @@ limitations under the License.
 package predicates
 
 import (
-	"vitess.io/vitess/go/vt/vterrors"
+	"maps"
 
 	"vitess.io/vitess/go/vt/sqlparser"
+	"vitess.io/vitess/go/vt/vterrors"
 )
 
 type (
@@ -29,15 +30,27 @@ type (
 	Tracker struct {
 		lastID      ID
 		expressions map[ID]sqlparser.Expr
+
+		// children tracks per-source copies of a predicate, such as the copies
+		// created when a predicate is pushed into every source of a UNION.
+		// When the original predicate is restored or skipped, the copies must
+		// be skipped as well - they cannot be restored to column form inside
+		// their source, and after a merge no one produces their arguments.
+		children map[ID][]ID
 	}
 
 	// ID is a unique key that references the current expression a join predicate represents.
 	ID int
+
+	// Snapshot holds the previous expressions of a predicate and its copies,
+	// captured by ResetToOriginal so an abandoned rewrite can be rolled back.
+	Snapshot map[ID]sqlparser.Expr
 )
 
 func NewTracker() *Tracker {
 	return &Tracker{
 		expressions: make(map[ID]sqlparser.Expr),
+		children:    make(map[ID][]ID),
 	}
 }
 
@@ -50,6 +63,24 @@ func (t *Tracker) NewJoinPredicate(org sqlparser.Expr) *JoinPredicate {
 	}
 }
 
+// NewChildJoinPredicate creates a new JoinPredicate that is tracked as a copy of the
+// given parent predicate, so that Skip/restore operations on the parent cascade to it.
+func (t *Tracker) NewChildJoinPredicate(parent *JoinPredicate, org sqlparser.Expr) *JoinPredicate {
+	jp := t.NewJoinPredicate(org)
+	t.children[parent.ID] = append(t.children[parent.ID], jp.ID)
+	return jp
+}
+
+// DescendantIDs returns the IDs of all transitive copies of the given predicate.
+func (t *Tracker) DescendantIDs(id ID) []ID {
+	var ids []ID
+	for _, child := range t.children[id] {
+		ids = append(ids, child)
+		ids = append(ids, t.DescendantIDs(child)...)
+	}
+	return ids
+}
+
 func (t *Tracker) nextID() ID {
 	id := t.lastID
 	t.lastID++
@@ -71,3 +102,31 @@ func (t *Tracker) Get(id ID) (sqlparser.Expr, error) {
 func (t *Tracker) Skip(id ID) {
 	t.expressions[id] = nil
 }
+
+// SkipWithDescendants skips the given predicate and all its transitive copies.
+func (t *Tracker) SkipWithDescendants(id ID) {
+	t.Skip(id)
+	for _, child := range t.children[id] {
+		t.SkipWithDescendants(child)
+	}
+}
+
+// ResetToOriginal returns the predicate to its pre-push expression and skips
+// all of its transitive copies: they cannot be restored to column form inside
+// their source, and once the parent is restored nothing produces their
+// arguments. It returns a Snapshot of the previous state so callers that may
+// abandon the rewrite can undo it with Rollback.
+func (t *Tracker) ResetToOriginal(id ID, expr sqlparser.Expr) Snapshot {
+	snapshot := Snapshot{id: t.expressions[id]}
+	t.expressions[id] = expr
+	for _, child := range t.DescendantIDs(id) {
+		snapshot[child] = t.expressions[child]
+		t.Skip(child)
+	}
+	return snapshot
+}
+
+// Rollback puts back the expressions captured in a Snapshot.
+func (t *Tracker) Rollback(s Snapshot) {
+	maps.Copy(t.expressions, s)
+}

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

@@ -0,0 +1,104 @@
+/*
+Copyright 2026 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 (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/vt/sqlparser"
+)
+
+func TestSkipWithDescendantsCascadesToChildren(t *testing.T) {
+	tracker := NewTracker()
+	parent := tracker.NewJoinPredicate(sqlparser.NewIntLiteral("1"))
+	childA := tracker.NewChildJoinPredicate(parent, sqlparser.NewIntLiteral("2"))
+	childB := tracker.NewChildJoinPredicate(parent, sqlparser.NewIntLiteral("3"))
+
+	tracker.SkipWithDescendants(parent.ID)
+
+	require.Nil(t, parent.Current())
+	require.Nil(t, childA.Current())
+	require.Nil(t, childB.Current())
+}
+
+func TestSkipWithDescendantsCascadesTransitively(t *testing.T) {
+	tracker := NewTracker()
+	parent := tracker.NewJoinPredicate(sqlparser.NewIntLiteral("1"))
+	child := tracker.NewChildJoinPredicate(parent, sqlparser.NewIntLiteral("2"))
+	grandchild := tracker.NewChildJoinPredicate(child, sqlparser.NewIntLiteral("3"))
+
+	tracker.SkipWithDescendants(parent.ID)
+
+	require.Nil(t, parent.Current())
+	require.Nil(t, child.Current())
+	require.Nil(t, grandchild.Current())
+}
+
+func TestSkipWithDescendantsLeavesUnrelatedPredicates(t *testing.T) {
+	tracker := NewTracker()
+	parent := tracker.NewJoinPredicate(sqlparser.NewIntLiteral("1"))
+	child := tracker.NewChildJoinPredicate(parent, sqlparser.NewIntLiteral("2"))
+	other := tracker.NewJoinPredicate(sqlparser.NewIntLiteral("3"))
+
+	tracker.SkipWithDescendants(parent.ID)
+
+	require.Nil(t, parent.Current())
+	require.Nil(t, child.Current())
+	require.NotNil(t, other.Current())
+}
+
+func TestDescendantIDsReturnsTransitiveClosure(t *testing.T) {
+	tracker := NewTracker()
+	parent := tracker.NewJoinPredicate(sqlparser.NewIntLiteral("1"))
+	childA := tracker.NewChildJoinPredicate(parent, sqlparser.NewIntLiteral("2"))
+	childB := tracker.NewChildJoinPredicate(parent, sqlparser.NewIntLiteral("3"))
+	grandchild := tracker.NewChildJoinPredicate(childA, sqlparser.NewIntLiteral("4"))
+
+	require.ElementsMatch(t, []ID{childA.ID, childB.ID, grandchild.ID}, tracker.DescendantIDs(parent.ID))
+	require.Empty(t, tracker.DescendantIDs(childB.ID))
+}
+
+func TestResetToOriginalRestoresParentAndSkipsDescendants(t *testing.T) {
+	tracker := NewTracker()
+	parent := tracker.NewJoinPredicate(sqlparser.NewIntLiteral("1"))
+	child := tracker.NewChildJoinPredicate(parent, sqlparser.NewIntLiteral("2"))
+	grandchild := tracker.NewChildJoinPredicate(child, sqlparser.NewIntLiteral("3"))
+	original := sqlparser.NewIntLiteral("42")
+
+	snapshot := tracker.ResetToOriginal(parent.ID, original)
+
+	require.Same(t, original, parent.Current())
+	require.Nil(t, child.Current())
+	require.Nil(t, grandchild.Current())
+	require.Len(t, snapshot, 3)
+}
+
+func TestRollbackRestoresSnapshotState(t *testing.T) {
+	tracker := NewTracker()
+	parent := tracker.NewJoinPredicate(sqlparser.NewIntLiteral("1"))
+	child := tracker.NewChildJoinPredicate(parent, sqlparser.NewIntLiteral("2"))
+	pushed := parent.Current()
+	copied := child.Current()
+
+	snapshot := tracker.ResetToOriginal(parent.ID, sqlparser.NewIntLiteral("42"))
+	tracker.Rollback(snapshot)
+
+	require.Same(t, pushed, parent.Current())
+	require.Same(t, copied, child.Current())
+}

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

@@ -19,6 +19,7 @@ package operators
 import (
 	"fmt"
 	"io"
+	"maps"
 	"strconv"
 
 	"vitess.io/vitess/go/slice"
@@ -151,26 +152,18 @@ func tryMergeApplyJoin(in *ApplyJoin, ctx *plancontext.PlanningContext) (_ Opera
 
 	//  - Rewrite join predicates already pushed down &&
 	//  - Save original join predicates if we have to bail out of the rewrite
-	original := map[predicates.ID]sqlparser.Expr{}
+	original := predicates.Snapshot{}
 	for _, col := range aj.JoinPredicates.columns {
 		if col.JoinPredicateID != nil {
 			// if we have pushed down a join predicate, we need to restore it to its original shape, without the argument from the LHS
-			id := *col.JoinPredicateID
-			oldExpr, err := ctx.PredTracker.Get(id)
-			if err != nil {
-				panic(err)
-			}
-			original[id] = oldExpr
-			ctx.PredTracker.Set(id, col.Original)
+			maps.Copy(original, ctx.PredTracker.ResetToOriginal(*col.JoinPredicateID, col.Original))
 		}
 	}
 
 	// Defer restoration of original predicates if no successful rewrite happens.
 	defer func() {
 		if res == NoRewrite {
-			for id, expr := range original {
-				ctx.PredTracker.Set(id, expr)
-			}
+			ctx.PredTracker.Rollback(original)
 		}
 	}()
 

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

@@ -52,6 +52,15 @@ type (
 		// this field will contain the conditions under which this route is valid
 		Conditions []engine.Condition
 
+		// MergeFallback is set when a UNION was merged into this route on a
+		// routing that ignores join predicates pushed down from an ApplyJoin
+		// above. It records the routing an argument-based merge would have
+		// installed instead, so later union merge attempts against sources
+		// routed elsewhere can still merge this route the way they would have
+		// before. It is never simultaneously this route's Routing, so
+		// predicates pushed into this route cannot mutate it.
+		MergeFallback *ShardedRouting
+
 		ResultColumns int
 	}
 
@@ -118,21 +127,34 @@ type (
 
 // UpdateRoutingLogic first checks if we are dealing with a predicate that
 func UpdateRoutingLogic(ctx *plancontext.PlanningContext, in sqlparser.Expr, r Routing) Routing {
+	if nr, ok := r.(*NoneRouting); ok {
+		// a none routing stays none no matter what further predicates say, and
+		// returning it as-is keeps its inferred-keyspace marker intact.
+		return nr
+	}
+
 	ks := r.Keyspace()
+	inferred := false
 	if ks == nil {
 		var err error
 		ks, err = ctx.VSchema.AnyKeyspace()
 		if err != nil {
 			panic(err)
 		}
+		inferred = true
 	}
-	nr := &NoneRouting{keyspace: ks}
+	nr := &NoneRouting{keyspace: ks, inferredKeyspace: inferred}
 
 	expr := in
 	// If we have a JoinPredicate, let's get the inner expression
 	pred, isJP := in.(*predicates.JoinPredicate)
 	if isJP {
 		expr = pred.Current()
+		if expr == nil {
+			// the predicate has been skipped - the join it belonged to has been
+			// merged away, so it no longer applies and must not influence routing
+			return r
+		}
 	}
 
 	if b := ctx.IsConstantBool(expr); b != nil && !*b {
@@ -186,6 +208,9 @@ func (r *Route) Clone(inputs []Operator) Operator {
 	cloneRoute := *r
 	cloneRoute.Source = inputs[0]
 	cloneRoute.Routing = r.Routing.Clone()
+	if r.MergeFallback != nil {
+		cloneRoute.MergeFallback = r.MergeFallback.Clone().(*ShardedRouting)
+	}
 	return &cloneRoute
 }