Merge cockroachdb#157031

157031: schemachanger/rel: fix notJoin subquery depth calculation for non-entity variables r=rafiss a=rafiss

Previously, the subquery depth calculation in setSubqueryDepth only considered
direct entity dependencies when determining when a notJoin subquery should
execute. It tracked the maximum entity slot among the subquery's inputs but
ignored when non-entity variables (like table-id, index-id) were actually bound.
This caused notJoin subqueries to execute too early, before all their required
variables were available, leading to "unbound variable" errors.

This commit fixes the depth calculation by tracking which entity provides each
variable through the facts. When a subquery needs a non-entity variable, we now
find which entity binds that variable and use that entity's slot to determine
the correct execution depth. This ensures notJoin subqueries only execute after
all their input variables have been bound by the appropriate entities in the
join order.

The fix is needed for rules that pass non-entity variables to notJoin
predicates, particularly in constraint rename scenarios where variables may
be bound at different join depths. (At the time of this writing, we
don't yet have rules like this, but we soon will add one.)

informs cockroachdb#148341
Release note: None

Co-authored-by: Rafi Shamim <[email protected]>

Author: craig[bot]

pkg/sql/schemachanger/rel/query_build.go

@@ -427,19 +427,41 @@ func (p *queryBuilder) setSubQueryDepths(entitySlots []slotIdx) {
 }
 
 func (p *queryBuilder) setSubqueryDepth(s *subQuery, entitySlots []slotIdx) {
-	var max int
+	// First, figure out which entity binds each variable.
+	// slotToEntityProvider maps slot -> entity slot that provides it.
+	slotToEntityProvider := make(map[int]int)
+	for _, f := range p.facts {
+		if p.slotIsEntity[f.variable] {
+			// This fact is about an entity, so it might bind a variable
+			slotToEntityProvider[int(f.value)] = int(f.variable)
+		}
+	}
+
+	var maxEntitySlot int
 	s.inputSlotMappings.ForEach(func(key, _ int) {
-		if p.slotIsEntity[key] && key > max {
-			max = key
+		if p.slotIsEntity[key] {
+			if key > maxEntitySlot {
+				maxEntitySlot = key
+			}
+		} else {
+			// For non-entity variables, find which entity provides them
+			provider, hasProvider := slotToEntityProvider[key]
+			if hasProvider {
+				if provider > maxEntitySlot {
+					maxEntitySlot = provider
+				}
+			}
 		}
 	})
+
 	got := sort.Search(len(entitySlots), func(i int) bool {
-		return int(entitySlots[i]) >= max
+		return int(entitySlots[i]) >= maxEntitySlot
 	})
 	if got == len(entitySlots) {
 		panic(errors.AssertionFailedf("failed to find maximum entity in entitySlots: %v not in %v",
-			max, entitySlots))
+			maxEntitySlot, entitySlots))
 	}
+
 	s.depth = queryDepth(got + 1)
 }
 

pkg/sql/schemachanger/rel/query_eval.go

@@ -450,6 +450,7 @@ func (ec *evalContext) maybeVisitSubqueries() (nextSubQuery int, done bool, erro
 func (ec *evalContext) visitSubquery(query int) (done bool, _ error) {
 	sub := ec.q.notJoins[query]
 	sec := sub.query.getEvalContext()
+
 	defer sub.query.putEvalContext(sec)
 	defer func() { // reset the slots populated to run the subquery
 		sub.inputSlotMappings.ForEach(func(_, subSlot int) {

pkg/sql/schemachanger/rel/rel_test.go

@@ -577,3 +577,188 @@ func TestConcurrentQueryInDifferentDatabases(t *testing.T) {
 	}
 	require.NoError(t, g.Wait())
 }
+
+type notJoinTestAttr string
+
+func (a notJoinTestAttr) String() string { return string(a) }
+
+// TestNotJoinSubqueryDepthWithNonEntityVariables tests that notJoin subqueries
+// execute at the correct depth when they depend on non-entity variables that
+// are bound by entities at different join depths.
+//
+// This test catches a bug where notJoin subqueries would execute too early,
+// before their required non-entity variables were bound. The fix ensures that
+// we track which entity provides each variable through facts, so notJoin
+// subqueries wait until all their input variables are available.
+func TestNotJoinSubqueryDepthWithNonEntityVariables(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	type FirstEntity struct {
+		ID       int
+		SharedID int
+	}
+	type SecondEntity struct {
+		ID       int
+		SharedID int
+		Value    string
+	}
+	type ThirdEntity struct {
+		ID       int
+		SharedID int
+		Flag     int
+	}
+
+	schema := rel.MustSchema("test_notjoin_depth",
+		rel.EntityMapping(reflect.TypeOf((*FirstEntity)(nil)),
+			rel.EntityAttr(notJoinTestAttr("id"), "ID"),
+			rel.EntityAttr(notJoinTestAttr("shared_id"), "SharedID"),
+		),
+		rel.EntityMapping(reflect.TypeOf((*SecondEntity)(nil)),
+			rel.EntityAttr(notJoinTestAttr("id"), "ID"),
+			rel.EntityAttr(notJoinTestAttr("shared_id"), "SharedID"),
+			rel.EntityAttr(notJoinTestAttr("value"), "Value"),
+		),
+		rel.EntityMapping(reflect.TypeOf((*ThirdEntity)(nil)),
+			rel.EntityAttr(notJoinTestAttr("id"), "ID"),
+			rel.EntityAttr(notJoinTestAttr("shared_id"), "SharedID"),
+			rel.EntityAttr(notJoinTestAttr("flag"), "Flag"),
+		),
+	)
+
+	// Define a notJoin rule that depends on a non-entity variable (shared_id).
+	// This rule checks if there's no ThirdEntity with the given shared_id and flag=1.
+	noThirdWithFlag := schema.DefNotJoin1("no_third_with_flag", "shared_id_var", func(
+		sharedIDVar rel.Var,
+	) rel.Clauses {
+		return rel.Clauses{
+			rel.Var("third").Type((*ThirdEntity)(nil)),
+			rel.Var("third").AttrEqVar(notJoinTestAttr("shared_id"), sharedIDVar),
+			rel.Var("third").AttrEq(notJoinTestAttr("flag"), 1),
+		}
+	})
+
+	first1 := &FirstEntity{ID: 1, SharedID: 100}
+	second1 := &SecondEntity{ID: 2, SharedID: 100, Value: "test"}
+	third1 := &ThirdEntity{ID: 3, SharedID: 100, Flag: 0}
+	third2 := &ThirdEntity{ID: 4, SharedID: 200, Flag: 1}
+
+	db, err := rel.NewDatabase(schema,
+		rel.Index{Attrs: []rel.Attr{rel.Type}},
+		rel.Index{Attrs: []rel.Attr{rel.Self}},
+		rel.Index{Attrs: []rel.Attr{notJoinTestAttr("id")}},
+		rel.Index{Attrs: []rel.Attr{notJoinTestAttr("shared_id")}},
+		rel.Index{Attrs: []rel.Attr{notJoinTestAttr("flag")}},
+		rel.Index{Attrs: []rel.Attr{notJoinTestAttr("value")}},
+	)
+	require.NoError(t, err)
+	require.NoError(t, db.Insert(first1))
+	require.NoError(t, db.Insert(second1))
+	require.NoError(t, db.Insert(third1))
+	require.NoError(t, db.Insert(third2))
+
+	// Test case 1: Query where shared_id is bound by SecondEntity (at depth 2).
+	// The notJoin should execute after SecondEntity is joined.
+	t.Run("notjoin_executes_after_second_entity", func(t *testing.T) {
+		q, err := rel.NewQuery(schema,
+			// FirstEntity is joined first (depth 1).
+			rel.Var("first").Type((*FirstEntity)(nil)),
+			rel.Var("first").AttrEq(notJoinTestAttr("id"), 1),
+			// SecondEntity is joined second (depth 2) and binds shared_id_var.
+			rel.Var("second").Type((*SecondEntity)(nil)),
+			rel.Var("second").AttrEqVar(notJoinTestAttr("shared_id"), "shared_id_var"),
+			// Join FirstEntity and SecondEntity on shared_id.
+			rel.Var("first").AttrEqVar(notJoinTestAttr("shared_id"), "shared_id_var"),
+			// This notJoin depends on shared_id_var, which is bound by SecondEntity
+			// It should execute at depth 2 or later, not before
+			noThirdWithFlag("shared_id_var"),
+		)
+		require.NoError(t, err)
+
+		var results [][]interface{}
+		err = q.Iterate(db, nil, func(r rel.Result) error {
+			results = append(results, []interface{}{
+				r.Var(rel.Var("first")),
+				r.Var(rel.Var("second")),
+				r.Var(rel.Var("shared_id_var")),
+			})
+			return nil
+		})
+		require.NoError(t, err)
+		// Should find the combination where shared_id=100.
+		require.Len(t, results, 1)
+		require.Equal(t, first1, results[0][0])
+		require.Equal(t, second1, results[0][1])
+		require.Equal(t, 100, results[0][2])
+	})
+
+	// Test case 2: Query where shared_id would cause the notJoin to fail.
+	t.Run("notjoin_filters_results_correctly", func(t *testing.T) {
+		// Add a FirstEntity with shared_id=200.
+		first2 := &FirstEntity{ID: 5, SharedID: 200}
+		second2 := &SecondEntity{ID: 6, SharedID: 200, Value: "test2"}
+		require.NoError(t, db.Insert(first2))
+		require.NoError(t, db.Insert(second2))
+
+		q, err := rel.NewQuery(schema,
+			rel.Var("first").Type((*FirstEntity)(nil)),
+			rel.Var("first").AttrEq(notJoinTestAttr("id"), 5),
+			rel.Var("second").Type((*SecondEntity)(nil)),
+			rel.Var("second").AttrEqVar(notJoinTestAttr("shared_id"), "shared_id_var"),
+			rel.Var("first").AttrEqVar(notJoinTestAttr("shared_id"), "shared_id_var"),
+			noThirdWithFlag("shared_id_var"),
+		)
+		require.NoError(t, err)
+
+		var results [][]interface{}
+		err = q.Iterate(db, nil, func(r rel.Result) error {
+			results = append(results, []interface{}{
+				r.Var(rel.Var("first")),
+				r.Var(rel.Var("second")),
+			})
+			return nil
+		})
+		require.NoError(t, err)
+		// Should find no results because third2 has shared_id=200 and flag=1.
+		require.Empty(t, results)
+	})
+
+	// Test case 3: Complex case with non-entity variable bound at depth 3.
+	t.Run("notjoin_with_variable_bound_at_depth_3", func(t *testing.T) {
+		// Define a more complex notJoin that uses two variables.
+		complexNotJoin := schema.DefNotJoin2("complex_not_join", "sid", "val", func(
+			sidVar, valVar rel.Var,
+		) rel.Clauses {
+			return rel.Clauses{
+				rel.Var("e").Type((*SecondEntity)(nil)),
+				rel.Var("e").AttrEqVar(notJoinTestAttr("shared_id"), sidVar),
+				rel.Var("e").AttrEqVar(notJoinTestAttr("value"), valVar),
+			}
+		})
+
+		q, err := rel.NewQuery(schema,
+			// Depth 1
+			rel.Var("f1").Type((*FirstEntity)(nil)),
+			// Depth 2
+			rel.Var("f2").Type((*FirstEntity)(nil)),
+			rel.Var("f2").AttrNeq(notJoinTestAttr("id"), 1),
+			// Depth 3 - this is where shared_id and value are bound.
+			rel.Var("s").Type((*SecondEntity)(nil)),
+			rel.Var("s").AttrEqVar(notJoinTestAttr("shared_id"), "sid"),
+			rel.Var("s").AttrEqVar(notJoinTestAttr("value"), "val"),
+			// The notJoin should only execute after depth 3.
+			complexNotJoin("sid", "val"),
+			// Add a filter to limit results.
+			rel.Filter("limit", "f1")(func(e *FirstEntity) bool {
+				return e.ID == 1
+			}),
+		)
+		require.NoError(t, err)
+
+		// This should execute without "unbound variable" errors
+		// even though the notJoin depends on variables bound at depth 3
+		err = q.Iterate(db, nil, func(r rel.Result) error {
+			return nil
+		})
+		require.NoError(t, err)
+	})
+}